Compare commits
85
Commits
@@ -0,0 +1,83 @@
|
|||||||
|
name: Build Nanobot OAuth
|
||||||
|
|
||||||
|
on:
|
||||||
|
push:
|
||||||
|
branches: ['main']
|
||||||
|
pull_request:
|
||||||
|
branches: ['main']
|
||||||
|
schedule:
|
||||||
|
- cron: '0 3 * * *'
|
||||||
|
workflow_dispatch:
|
||||||
|
|
||||||
|
env:
|
||||||
|
REGISTRY: git.wylab.me
|
||||||
|
IMAGE_NAME: wylab/nanobot
|
||||||
|
BUILDKIT_PROGRESS: plain
|
||||||
|
|
||||||
|
jobs:
|
||||||
|
build:
|
||||||
|
runs-on: [self-hosted, linux-amd64]
|
||||||
|
timeout-minutes: 15
|
||||||
|
permissions:
|
||||||
|
contents: read
|
||||||
|
packages: write
|
||||||
|
steps:
|
||||||
|
- name: Checkout repository
|
||||||
|
uses: actions/checkout@v4
|
||||||
|
|
||||||
|
- name: Set up Docker Buildx
|
||||||
|
uses: docker/setup-buildx-action@v3
|
||||||
|
|
||||||
|
- name: Log in to the container registry
|
||||||
|
uses: docker/login-action@v3
|
||||||
|
with:
|
||||||
|
registry: ${{ env.REGISTRY }}
|
||||||
|
username: ${{ secrets.REGISTRY_USERNAME || github.actor }}
|
||||||
|
password: ${{ secrets.REGISTRY_PASSWORD || secrets.GITHUB_TOKEN }}
|
||||||
|
|
||||||
|
- name: Build and push Docker image
|
||||||
|
uses: docker/build-push-action@v5
|
||||||
|
with:
|
||||||
|
context: .
|
||||||
|
file: Dockerfile.oauth
|
||||||
|
provenance: false
|
||||||
|
platforms: linux/amd64
|
||||||
|
cache-from: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:buildcache
|
||||||
|
cache-to: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:buildcache,mode=max
|
||||||
|
push: ${{ github.event_name != 'pull_request' }}
|
||||||
|
tags: |
|
||||||
|
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest
|
||||||
|
${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ github.sha }}
|
||||||
|
|
||||||
|
cleanup:
|
||||||
|
if: github.event_name == 'push' || github.event_name == 'schedule'
|
||||||
|
runs-on: [self-hosted, linux-amd64]
|
||||||
|
needs: build
|
||||||
|
steps:
|
||||||
|
- name: Delete images older than 24h
|
||||||
|
env:
|
||||||
|
TOKEN: ${{ secrets.REGISTRY_PASSWORD || secrets.GITHUB_TOKEN }}
|
||||||
|
run: |
|
||||||
|
cutoff=$(date -u -d '24 hours ago' +%s)
|
||||||
|
page=1
|
||||||
|
while true; do
|
||||||
|
versions=$(curl -sf -H "Authorization: token $TOKEN" \
|
||||||
|
"https://${{ env.REGISTRY }}/api/v1/packages/wylab?type=container&q=nanobot&limit=50&page=$page")
|
||||||
|
count=$(echo "$versions" | jq length)
|
||||||
|
[ "$count" = "0" ] && break
|
||||||
|
echo "$versions" | jq -c '.[]' | while read -r pkg; do
|
||||||
|
ver=$(echo "$pkg" | jq -r '.version')
|
||||||
|
# Keep latest and buildcache, only delete SHA tags
|
||||||
|
case "$ver" in latest|buildcache) continue ;; esac
|
||||||
|
created=$(echo "$pkg" | jq -r '.created_at')
|
||||||
|
ts=$(date -u -d "$created" +%s 2>/dev/null || echo 0)
|
||||||
|
if [ "$ts" -lt "$cutoff" ]; then
|
||||||
|
id=$(echo "$pkg" | jq -r '.id')
|
||||||
|
echo "Deleting nanobot:$ver (id=$id, created=$created)"
|
||||||
|
curl -sf -X DELETE -H "Authorization: token $TOKEN" \
|
||||||
|
"https://${{ env.REGISTRY }}/api/v1/packages/wylab/container/nanobot/$ver" || true
|
||||||
|
fi
|
||||||
|
done
|
||||||
|
[ "$count" -lt 50 ] && break
|
||||||
|
page=$((page + 1))
|
||||||
|
done
|
||||||
@@ -0,0 +1,62 @@
|
|||||||
|
FROM birdxs/nanobot:latest
|
||||||
|
|
||||||
|
# ── Skill dependencies ──────────────────────────────────────────────
|
||||||
|
|
||||||
|
# APT: ffmpeg (video-frames, whisper), jq, tmux, build-essential (for go)
|
||||||
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||||
|
ffmpeg jq tmux build-essential procps && rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
# gh CLI via GitHub official apt repo
|
||||||
|
RUN curl -fsSL https://cli.github.com/packages/githubcli-archive-keyring.gpg \
|
||||||
|
| dd of=/usr/share/keyrings/githubcli-archive-keyring.gpg && \
|
||||||
|
echo "deb [arch=amd64 signed-by=/usr/share/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" \
|
||||||
|
> /etc/apt/sources.list.d/github-cli.list && \
|
||||||
|
apt-get update && apt-get install -y gh && rm -rf /var/lib/apt/lists/*
|
||||||
|
|
||||||
|
# Go toolchain
|
||||||
|
RUN curl -fsSL https://go.dev/dl/go1.23.6.linux-amd64.tar.gz | tar -C /usr/local -xzf -
|
||||||
|
ENV PATH="/usr/local/go/bin:/root/go/bin:${PATH}"
|
||||||
|
|
||||||
|
# Go tools: blogwatcher, blu (blucli), gifgrep, sonos (sonoscli), wacli, songsee
|
||||||
|
RUN go install github.com/Hyaxia/blogwatcher/cmd/blogwatcher@latest && \
|
||||||
|
go install github.com/steipete/blucli/cmd/blu@latest && \
|
||||||
|
go install github.com/steipete/gifgrep/cmd/gifgrep@latest && \
|
||||||
|
go install github.com/steipete/sonoscli/cmd/sonos@latest && \
|
||||||
|
go install github.com/steipete/wacli/cmd/wacli@latest && \
|
||||||
|
go install github.com/steipete/songsee/cmd/songsee@latest
|
||||||
|
|
||||||
|
# Pre-built binaries from GitHub releases
|
||||||
|
# gogcli (gog)
|
||||||
|
RUN curl -fsSL https://github.com/steipete/gogcli/releases/download/v0.9.0/gogcli_0.9.0_linux_amd64.tar.gz \
|
||||||
|
| tar -xzf - -C /usr/local/bin gog
|
||||||
|
|
||||||
|
# goplaces
|
||||||
|
RUN curl -fsSL https://github.com/steipete/goplaces/releases/download/v0.2.1/goplaces_0.2.1_linux_amd64.tar.gz \
|
||||||
|
| tar -xzf - -C /usr/local/bin goplaces
|
||||||
|
|
||||||
|
# himalaya (email CLI)
|
||||||
|
RUN curl -fsSL https://github.com/pimalaya/himalaya/releases/download/v1.1.0/himalaya.x86_64-linux.tgz \
|
||||||
|
| tar -xzf - -C /usr/local/bin himalaya
|
||||||
|
|
||||||
|
# obsidian-cli (release binary is named notesmd-cli, skill expects obsidian-cli)
|
||||||
|
RUN curl -fsSL -o /tmp/obsidian.tar.gz https://github.com/yakitrak/obsidian-cli/releases/download/v0.3.0/notesmd-cli_0.3.0_linux_amd64.tar.gz && \
|
||||||
|
tar -xzf /tmp/obsidian.tar.gz -C /tmp notesmd-cli && \
|
||||||
|
mv /tmp/notesmd-cli /usr/local/bin/obsidian-cli && \
|
||||||
|
rm /tmp/obsidian.tar.gz
|
||||||
|
|
||||||
|
# Node tools: oracle, gemini-cli, summarize
|
||||||
|
RUN npm install -g @steipete/oracle @google/gemini-cli @steipete/summarize
|
||||||
|
|
||||||
|
# Python tools: nano-pdf, openai-whisper
|
||||||
|
RUN uv tool install nano-pdf && \
|
||||||
|
uv tool install openai-whisper
|
||||||
|
ENV PATH="/root/.local/bin:${PATH}"
|
||||||
|
|
||||||
|
# ── Nanobot source ──────────────────────────────────────────────────
|
||||||
|
|
||||||
|
COPY pyproject.toml README.md LICENSE /app/
|
||||||
|
COPY nanobot/ /app/nanobot/
|
||||||
|
RUN uv pip install --system --no-cache --reinstall /app psycopg2-binary
|
||||||
|
|
||||||
|
ENTRYPOINT ["nanobot"]
|
||||||
|
CMD ["gateway"]
|
||||||
+18
-16
@@ -45,10 +45,14 @@ class ContextBuilder:
|
|||||||
if bootstrap:
|
if bootstrap:
|
||||||
parts.append(bootstrap)
|
parts.append(bootstrap)
|
||||||
|
|
||||||
# Memory context
|
# Static knowledge context (KNOWLEDGE.md — manually curated, stable for caching)
|
||||||
memory = self.memory.get_memory_context()
|
# MEMORY.md is excluded from system prompt as it changes frequently (consolidator),
|
||||||
if memory:
|
# but the agent can still read/grep it via tools.
|
||||||
parts.append(f"# Memory\n\n{memory}")
|
knowledge_file = self.memory.memory_dir / "KNOWLEDGE.md"
|
||||||
|
if knowledge_file.exists():
|
||||||
|
knowledge = knowledge_file.read_text(encoding="utf-8").strip()
|
||||||
|
if knowledge:
|
||||||
|
parts.append(f"# Knowledge\n\n{knowledge}")
|
||||||
|
|
||||||
# Skills - progressive loading
|
# Skills - progressive loading
|
||||||
# 1. Always-loaded skills: include full content
|
# 1. Always-loaded skills: include full content
|
||||||
@@ -71,27 +75,18 @@ Skills with available="false" need dependencies installed first - you can try in
|
|||||||
return "\n\n---\n\n".join(parts)
|
return "\n\n---\n\n".join(parts)
|
||||||
|
|
||||||
def _get_identity(self) -> str:
|
def _get_identity(self) -> str:
|
||||||
"""Get the core identity section."""
|
"""Get the core identity section with runtime context."""
|
||||||
from datetime import datetime
|
|
||||||
import time as _time
|
|
||||||
now = datetime.now().strftime("%Y-%m-%d %H:%M (%A)")
|
|
||||||
tz = _time.strftime("%Z") or "UTC"
|
|
||||||
workspace_path = str(self.workspace.expanduser().resolve())
|
workspace_path = str(self.workspace.expanduser().resolve())
|
||||||
system = platform.system()
|
system = platform.system()
|
||||||
runtime = f"{'macOS' if system == 'Darwin' else system} {platform.machine()}, Python {platform.python_version()}"
|
runtime = f"{'macOS' if system == 'Darwin' else system} {platform.machine()}, Python {platform.python_version()}"
|
||||||
|
|
||||||
return f"""# nanobot 🐈
|
return f"""You have access to tools that allow you to:
|
||||||
|
|
||||||
You are nanobot, a helpful AI assistant. You have access to tools that allow you to:
|
|
||||||
- Read, write, and edit files
|
- Read, write, and edit files
|
||||||
- Execute shell commands
|
- Execute shell commands
|
||||||
- Search the web and fetch web pages
|
- Search the web and fetch web pages
|
||||||
- Send messages to users on chat channels
|
- Send messages to users on chat channels
|
||||||
- Spawn subagents for complex background tasks
|
- Spawn subagents for complex background tasks
|
||||||
|
|
||||||
## Current Time
|
|
||||||
{now} ({tz})
|
|
||||||
|
|
||||||
## Runtime
|
## Runtime
|
||||||
{runtime}
|
{runtime}
|
||||||
|
|
||||||
@@ -107,7 +102,14 @@ For normal conversation, just respond with text - do not call the message tool.
|
|||||||
|
|
||||||
Always be helpful, accurate, and concise. When using tools, think step by step: what you know, what you need, and why you chose this tool.
|
Always be helpful, accurate, and concise. When using tools, think step by step: what you know, what you need, and why you chose this tool.
|
||||||
When remembering something important, write to {workspace_path}/memory/MEMORY.md
|
When remembering something important, write to {workspace_path}/memory/MEMORY.md
|
||||||
To recall past events, grep {workspace_path}/memory/HISTORY.md"""
|
To recall past events, grep {workspace_path}/memory/HISTORY.md
|
||||||
|
|
||||||
|
## Visibility Markers
|
||||||
|
|
||||||
|
Messages marked with [HIDDEN:{{signature}}] were not sent to the user. These markers
|
||||||
|
are cryptographically signed by the system to track internal reasoning and background
|
||||||
|
tasks. Do NOT generate [HIDDEN:*] patterns yourself - outputs containing forged
|
||||||
|
visibility markers will be rejected."""
|
||||||
|
|
||||||
def _load_bootstrap_files(self) -> str:
|
def _load_bootstrap_files(self) -> str:
|
||||||
"""Load all bootstrap files from workspace."""
|
"""Load all bootstrap files from workspace."""
|
||||||
|
|||||||
+351
-39
@@ -2,6 +2,8 @@
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import json
|
import json
|
||||||
|
import time
|
||||||
|
from datetime import datetime
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -17,9 +19,11 @@ from nanobot.agent.tools.shell import ExecTool
|
|||||||
from nanobot.agent.tools.web import WebSearchTool, WebFetchTool
|
from nanobot.agent.tools.web import WebSearchTool, WebFetchTool
|
||||||
from nanobot.agent.tools.message import MessageTool
|
from nanobot.agent.tools.message import MessageTool
|
||||||
from nanobot.agent.tools.spawn import SpawnTool
|
from nanobot.agent.tools.spawn import SpawnTool
|
||||||
|
from nanobot.agent.tools.wait import WaitForSubagentsTool
|
||||||
from nanobot.agent.tools.cron import CronTool
|
from nanobot.agent.tools.cron import CronTool
|
||||||
from nanobot.agent.memory import MemoryStore
|
from nanobot.agent.memory import MemoryStore
|
||||||
from nanobot.agent.subagent import SubagentManager
|
from nanobot.agent.subagent import SubagentManager
|
||||||
|
from nanobot.agent.visibility import sign_content, has_forged_marker, strip_all_hidden_markers
|
||||||
from nanobot.session.manager import SessionManager
|
from nanobot.session.manager import SessionManager
|
||||||
|
|
||||||
|
|
||||||
@@ -35,6 +39,22 @@ class AgentLoop:
|
|||||||
5. Sends responses back
|
5. Sends responses back
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
# Server-side context management: Anthropic trims old tool results and preserves all
|
||||||
|
# thinking blocks (keep="all" maximises cache hits). Client keeps full history.
|
||||||
|
CONTEXT_MANAGEMENT = {
|
||||||
|
"edits": [
|
||||||
|
{
|
||||||
|
"type": "clear_thinking_20251015",
|
||||||
|
"keep": "all", # Preserve all thinking blocks for cache reuse
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"type": "clear_tool_uses_20250919",
|
||||||
|
"trigger": {"type": "input_tokens", "value": 80000},
|
||||||
|
"keep": {"type": "tool_uses", "value": 5},
|
||||||
|
},
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
bus: MessageBus,
|
bus: MessageBus,
|
||||||
@@ -69,13 +89,14 @@ class AgentLoop:
|
|||||||
provider=provider,
|
provider=provider,
|
||||||
workspace=workspace,
|
workspace=workspace,
|
||||||
bus=bus,
|
bus=bus,
|
||||||
model=self.model,
|
|
||||||
brave_api_key=brave_api_key,
|
brave_api_key=brave_api_key,
|
||||||
exec_config=self.exec_config,
|
exec_config=self.exec_config,
|
||||||
restrict_to_workspace=restrict_to_workspace,
|
restrict_to_workspace=restrict_to_workspace,
|
||||||
)
|
)
|
||||||
|
|
||||||
self._running = False
|
self._running = False
|
||||||
|
self._quota_cache: dict[str, Any] = {} # {model: str, cached_at: float}
|
||||||
|
self._quota_cache_ttl: float = 300.0 # 5 minutes
|
||||||
self._register_default_tools()
|
self._register_default_tools()
|
||||||
|
|
||||||
def _register_default_tools(self) -> None:
|
def _register_default_tools(self) -> None:
|
||||||
@@ -99,12 +120,13 @@ class AgentLoop:
|
|||||||
self.tools.register(WebFetchTool())
|
self.tools.register(WebFetchTool())
|
||||||
|
|
||||||
# Message tool
|
# Message tool
|
||||||
message_tool = MessageTool(send_callback=self.bus.publish_outbound)
|
message_tool = MessageTool(send_callback=self.bus.publish_outbound, sessions=self.sessions)
|
||||||
self.tools.register(message_tool)
|
self.tools.register(message_tool)
|
||||||
|
|
||||||
# Spawn tool (for subagents)
|
# Spawn tool (for subagents)
|
||||||
spawn_tool = SpawnTool(manager=self.subagents)
|
spawn_tool = SpawnTool(manager=self.subagents)
|
||||||
self.tools.register(spawn_tool)
|
self.tools.register(spawn_tool)
|
||||||
|
self.tools.register(WaitForSubagentsTool(manager=self.subagents))
|
||||||
|
|
||||||
# Cron tool (for scheduling)
|
# Cron tool (for scheduling)
|
||||||
if self.cron_service:
|
if self.cron_service:
|
||||||
@@ -126,7 +148,7 @@ class AgentLoop:
|
|||||||
# Process it
|
# Process it
|
||||||
try:
|
try:
|
||||||
response = await self._process_message(msg)
|
response = await self._process_message(msg)
|
||||||
if response:
|
if response and not response.metadata.get("suppressed", False):
|
||||||
await self.bus.publish_outbound(response)
|
await self.bus.publish_outbound(response)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error processing message: {e}")
|
logger.error(f"Error processing message: {e}")
|
||||||
@@ -134,7 +156,8 @@ class AgentLoop:
|
|||||||
await self.bus.publish_outbound(OutboundMessage(
|
await self.bus.publish_outbound(OutboundMessage(
|
||||||
channel=msg.channel,
|
channel=msg.channel,
|
||||||
chat_id=msg.chat_id,
|
chat_id=msg.chat_id,
|
||||||
content=f"Sorry, I encountered an error: {str(e)}"
|
content=f"Sorry, I encountered an error: {str(e)}",
|
||||||
|
metadata=msg.metadata or {},
|
||||||
))
|
))
|
||||||
except asyncio.TimeoutError:
|
except asyncio.TimeoutError:
|
||||||
continue
|
continue
|
||||||
@@ -144,6 +167,94 @@ class AgentLoop:
|
|||||||
self._running = False
|
self._running = False
|
||||||
logger.info("Agent loop stopping")
|
logger.info("Agent loop stopping")
|
||||||
|
|
||||||
|
def _select_model_based_on_quota(self) -> str:
|
||||||
|
"""Select Opus or Sonnet based on rolling weekly quota burn rate."""
|
||||||
|
# Check cache
|
||||||
|
now = time.time()
|
||||||
|
if self._quota_cache and (now - self._quota_cache.get("cached_at", 0)) < self._quota_cache_ttl:
|
||||||
|
return self._quota_cache["model"]
|
||||||
|
|
||||||
|
# Default models
|
||||||
|
OPUS = "claude-opus-4-6"
|
||||||
|
SONNET = "claude-sonnet-4-6"
|
||||||
|
TOLERANCE = 1.17 # 17% overage triggers downgrade
|
||||||
|
|
||||||
|
# Read rate limits
|
||||||
|
rate_limits_path = self.workspace / "memory" / "rate_limits.json"
|
||||||
|
if not rate_limits_path.exists():
|
||||||
|
logger.warning("rate_limits.json not found, defaulting to Sonnet")
|
||||||
|
return SONNET
|
||||||
|
|
||||||
|
try:
|
||||||
|
with open(rate_limits_path) as f:
|
||||||
|
limits = json.load(f)
|
||||||
|
|
||||||
|
actual_usage = limits.get("weekly_all_models")
|
||||||
|
weekly_reset = limits.get("weekly_reset")
|
||||||
|
|
||||||
|
if actual_usage is None or weekly_reset is None:
|
||||||
|
logger.warning("Rate limit data incomplete, defaulting to Sonnet")
|
||||||
|
return SONNET
|
||||||
|
|
||||||
|
# Calculate expected usage
|
||||||
|
actual_pct = actual_usage * 100
|
||||||
|
week_start = weekly_reset - (168 * 3600)
|
||||||
|
hours_elapsed = max(0, min((now - week_start) / 3600, 168))
|
||||||
|
expected_pct = (hours_elapsed / 168) * 100
|
||||||
|
threshold = expected_pct * TOLERANCE
|
||||||
|
|
||||||
|
# Decision logic
|
||||||
|
if actual_pct > threshold:
|
||||||
|
model = SONNET
|
||||||
|
logger.info(
|
||||||
|
f"Quota: {actual_pct:.1f}% used, expected {expected_pct:.1f}%, "
|
||||||
|
f"threshold {threshold:.1f}% → Sonnet"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
model = OPUS
|
||||||
|
logger.info(
|
||||||
|
f"Quota: {actual_pct:.1f}% used, expected {expected_pct:.1f}%, "
|
||||||
|
f"threshold {threshold:.1f}% → Opus"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Cache decision
|
||||||
|
self._quota_cache = {"model": model, "cached_at": now}
|
||||||
|
return model
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error checking quota: {e}, defaulting to Sonnet")
|
||||||
|
return SONNET
|
||||||
|
|
||||||
|
def _get_quota_status(self) -> str:
|
||||||
|
"""Return human-readable quota status."""
|
||||||
|
rate_limits_path = self.workspace / "memory" / "rate_limits.json"
|
||||||
|
if not rate_limits_path.exists():
|
||||||
|
return "⚠️ No quota data available yet."
|
||||||
|
|
||||||
|
try:
|
||||||
|
with open(rate_limits_path) as f:
|
||||||
|
limits = json.load(f)
|
||||||
|
|
||||||
|
actual_pct = limits.get("weekly_all_models", 0) * 100
|
||||||
|
reset_ts = limits.get("weekly_reset", 0)
|
||||||
|
now = time.time()
|
||||||
|
|
||||||
|
hours_until_reset = (reset_ts - now) / 3600
|
||||||
|
week_start = reset_ts - (168 * 3600)
|
||||||
|
hours_elapsed = max(0, (now - week_start) / 3600)
|
||||||
|
expected_pct = (hours_elapsed / 168) * 100
|
||||||
|
|
||||||
|
model = self._select_model_based_on_quota()
|
||||||
|
|
||||||
|
return f"""📊 Quota Status:
|
||||||
|
• Used: {actual_pct:.1f}% (expected {expected_pct:.1f}%)
|
||||||
|
• Resets in: {hours_until_reset:.1f}h
|
||||||
|
• Current model: {model}
|
||||||
|
• Burn rate: {actual_pct / max(expected_pct, 0.01):.2f}x target"""
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
return f"⚠️ Error reading quota: {e}"
|
||||||
|
|
||||||
async def _process_message(self, msg: InboundMessage, session_key: str | None = None) -> OutboundMessage | None:
|
async def _process_message(self, msg: InboundMessage, session_key: str | None = None) -> OutboundMessage | None:
|
||||||
"""
|
"""
|
||||||
Process a single inbound message.
|
Process a single inbound message.
|
||||||
@@ -174,14 +285,16 @@ class AgentLoop:
|
|||||||
session.clear()
|
session.clear()
|
||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id,
|
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id,
|
||||||
content="🐈 New session started. Memory consolidated.")
|
content="🐈 New session started. Memory consolidated.",
|
||||||
|
metadata=msg.metadata or {})
|
||||||
if cmd == "/help":
|
if cmd == "/help":
|
||||||
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id,
|
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id,
|
||||||
content="🐈 nanobot commands:\n/new — Start a new conversation\n/help — Show available commands")
|
content="🐈 nanobot commands:\n/new — Start a new conversation\n/help — Show available commands\n/quota — Show quota status",
|
||||||
|
metadata=msg.metadata or {})
|
||||||
# Consolidate memory before processing if session is too large
|
if cmd == "/quota":
|
||||||
if len(session.messages) > self.memory_window:
|
status = self._get_quota_status()
|
||||||
await self._consolidate_memory(session)
|
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id, content=status,
|
||||||
|
metadata=msg.metadata or {})
|
||||||
|
|
||||||
# Update tool contexts
|
# Update tool contexts
|
||||||
message_tool = self.tools.get("message")
|
message_tool = self.tools.get("message")
|
||||||
@@ -190,34 +303,75 @@ class AgentLoop:
|
|||||||
|
|
||||||
spawn_tool = self.tools.get("spawn")
|
spawn_tool = self.tools.get("spawn")
|
||||||
if isinstance(spawn_tool, SpawnTool):
|
if isinstance(spawn_tool, SpawnTool):
|
||||||
spawn_tool.set_context(msg.channel, msg.chat_id)
|
spawn_tool.set_context(msg.channel, msg.chat_id, msg.metadata)
|
||||||
|
|
||||||
cron_tool = self.tools.get("cron")
|
cron_tool = self.tools.get("cron")
|
||||||
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)
|
||||||
|
|
||||||
|
# Prepend current time + optional time-gap notice to every user message
|
||||||
|
now_dt = datetime.now()
|
||||||
|
tz = time.strftime("%Z") or "UTC"
|
||||||
|
time_str = now_dt.strftime("%Y-%m-%d %H:%M (%A)")
|
||||||
|
current_message = f"[Current time: {time_str} {tz}]\n{msg.content}"
|
||||||
|
|
||||||
|
# Prefix hook messages so the agent can identify them
|
||||||
|
hook_source = msg.metadata.get("hook_source") if msg.metadata else None
|
||||||
|
if hook_source:
|
||||||
|
current_message = f'[HOOK MESSAGE from "{hook_source}"]\n{current_message}'
|
||||||
|
|
||||||
|
last_user_ts = None
|
||||||
|
for m in reversed(session.messages):
|
||||||
|
if m.get("role") == "user":
|
||||||
|
last_user_ts = m.get("timestamp")
|
||||||
|
break
|
||||||
|
if last_user_ts:
|
||||||
|
try:
|
||||||
|
last_dt = datetime.fromisoformat(last_user_ts)
|
||||||
|
elapsed_seconds = (now_dt - last_dt).total_seconds()
|
||||||
|
if elapsed_seconds > 300: # 5 minutes
|
||||||
|
if elapsed_seconds < 3600:
|
||||||
|
gap_str = f"{int(elapsed_seconds // 60)} minutes"
|
||||||
|
elif elapsed_seconds < 86400:
|
||||||
|
gap_str = f"{int(elapsed_seconds // 3600)} hours"
|
||||||
|
else:
|
||||||
|
gap_str = f"{int(elapsed_seconds // 86400)} days"
|
||||||
|
current_message = f"[SYSTEM ANNOUNCEMENT: {gap_str} have elapsed since last user message; take this into account when replying to user. Ask about it if appropriate]\n\n{current_message}"
|
||||||
|
except (ValueError, TypeError):
|
||||||
|
pass # Malformed timestamp — skip silently
|
||||||
|
|
||||||
# Build initial messages (use get_history for LLM-formatted messages)
|
# Build initial messages (use get_history for LLM-formatted messages)
|
||||||
messages = self.context.build_messages(
|
messages = self.context.build_messages(
|
||||||
history=session.get_history(),
|
history=session.get_history(),
|
||||||
current_message=msg.content,
|
current_message=current_message,
|
||||||
media=msg.media if msg.media else None,
|
media=msg.media if msg.media else None,
|
||||||
channel=msg.channel,
|
channel=msg.channel,
|
||||||
chat_id=msg.chat_id,
|
chat_id=msg.chat_id,
|
||||||
)
|
)
|
||||||
|
# Mark where the current turn starts so we can slice the tool chain for storage
|
||||||
|
turn_start = len(messages)
|
||||||
|
|
||||||
|
# Select model based on quota
|
||||||
|
selected_model = self._select_model_based_on_quota()
|
||||||
|
|
||||||
|
# Check for suppress mode BEFORE the loop so it's available for forgery detection
|
||||||
|
suppress_output = msg.metadata.get("suppress_output", False) if msg.metadata else False
|
||||||
|
|
||||||
# Agent loop
|
# Agent loop
|
||||||
iteration = 0
|
iteration = 0
|
||||||
final_content = None
|
final_content = None
|
||||||
tools_used: list[str] = []
|
final_reasoning = None
|
||||||
|
|
||||||
while iteration < self.max_iterations:
|
while iteration < self.max_iterations:
|
||||||
iteration += 1
|
iteration += 1
|
||||||
|
|
||||||
# Call LLM
|
# Call LLM
|
||||||
|
logger.debug(f"Calling LLM with model={selected_model}, provider.thinking_budget={self.provider.thinking_budget}")
|
||||||
response = await self.provider.chat(
|
response = await self.provider.chat(
|
||||||
messages=messages,
|
messages=messages,
|
||||||
tools=self.tools.get_definitions(),
|
tools=self.tools.get_definitions(),
|
||||||
model=self.model
|
model=selected_model,
|
||||||
|
context_management=self.CONTEXT_MANAGEMENT,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Handle tool calls
|
# Handle tool calls
|
||||||
@@ -241,18 +395,44 @@ class AgentLoop:
|
|||||||
|
|
||||||
# Execute tools
|
# Execute tools
|
||||||
for tool_call in response.tool_calls:
|
for tool_call in response.tool_calls:
|
||||||
tools_used.append(tool_call.name)
|
|
||||||
args_str = json.dumps(tool_call.arguments, ensure_ascii=False)
|
args_str = json.dumps(tool_call.arguments, ensure_ascii=False)
|
||||||
logger.info(f"Tool call: {tool_call.name}({args_str[:200]})")
|
logger.info(f"Tool call: {tool_call.name}({args_str[:200]})")
|
||||||
result = await self.tools.execute(tool_call.name, tool_call.arguments)
|
result = await self.tools.execute(tool_call.name, tool_call.arguments)
|
||||||
messages = self.context.add_tool_result(
|
messages = self.context.add_tool_result(
|
||||||
messages, tool_call.id, tool_call.name, result
|
messages, tool_call.id, tool_call.name, result
|
||||||
)
|
)
|
||||||
# Interleaved CoT: reflect before next action
|
# Interleaved CoT: reflect before next action (skip when thinking is active)
|
||||||
messages.append({"role": "user", "content": "Reflect on the results and decide next steps."})
|
if not getattr(self.provider, 'thinking_budget', 0):
|
||||||
|
messages.append({"role": "user", "content": "Reflect on the results and decide next steps."})
|
||||||
else:
|
else:
|
||||||
# No tool calls, we're done
|
# No tool calls
|
||||||
final_content = response.content
|
final_content = response.content
|
||||||
|
final_reasoning = response.reasoning_content
|
||||||
|
|
||||||
|
# Check for forged signatures if in suppress mode
|
||||||
|
if suppress_output and has_forged_marker(final_content):
|
||||||
|
# Initialize retry counter if needed
|
||||||
|
if not hasattr(self, '_forge_retry_count'):
|
||||||
|
self._forge_retry_count = 0
|
||||||
|
|
||||||
|
if self._forge_retry_count < 1:
|
||||||
|
# First offense: reject and retry with correction
|
||||||
|
self._forge_retry_count += 1
|
||||||
|
logger.warning("Model attempted to forge visibility marker, rejecting output")
|
||||||
|
messages.append({
|
||||||
|
"role": "user",
|
||||||
|
"content": "[System: Previous response rejected. Do not generate [HIDDEN:*] markers.]"
|
||||||
|
})
|
||||||
|
continue # Back to while loop, will retry LLM call
|
||||||
|
else:
|
||||||
|
# Second offense: strip and log error (fallback)
|
||||||
|
logger.error("Model persisted in forging markers despite correction, stripping")
|
||||||
|
final_content = strip_all_hidden_markers(final_content)
|
||||||
|
|
||||||
|
# Reset retry counter on successful completion
|
||||||
|
if hasattr(self, '_forge_retry_count'):
|
||||||
|
self._forge_retry_count = 0
|
||||||
|
|
||||||
break
|
break
|
||||||
|
|
||||||
if final_content is None:
|
if final_content is None:
|
||||||
@@ -265,17 +445,37 @@ class AgentLoop:
|
|||||||
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}")
|
||||||
|
|
||||||
# Save to session (include tool names so consolidation sees what happened)
|
# suppress_output already defined before the loop
|
||||||
session.add_message("user", msg.content)
|
|
||||||
session.add_message("assistant", final_content,
|
if suppress_output:
|
||||||
tools_used=tools_used if tools_used else None)
|
# Sign content with our secret key (forgery check already done in loop)
|
||||||
|
final_content_for_session = sign_content(final_content)
|
||||||
|
# Mark as suppressed for channel handler
|
||||||
|
outbound_metadata = {**(msg.metadata or {}), "suppressed": True}
|
||||||
|
else:
|
||||||
|
final_content_for_session = final_content
|
||||||
|
outbound_metadata = msg.metadata or {}
|
||||||
|
|
||||||
|
# Append final assistant response to messages so it's captured in the tool chain slice
|
||||||
|
# Use the prefixed version for session storage
|
||||||
|
messages = self.context.add_assistant_message(
|
||||||
|
messages, final_content_for_session, None,
|
||||||
|
reasoning_content=final_reasoning,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Save to session: user message + full tool chain (tool_use, tool_results, thinking, final reply)
|
||||||
|
# Store current_message (not msg.content) so the time prefix is preserved
|
||||||
|
# and cache keys match on subsequent turns
|
||||||
|
session.add_message("user", current_message)
|
||||||
|
for chain_msg in messages[turn_start:]:
|
||||||
|
session.add_raw_message(chain_msg)
|
||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
|
|
||||||
return OutboundMessage(
|
return OutboundMessage(
|
||||||
channel=msg.channel,
|
channel=msg.channel,
|
||||||
chat_id=msg.chat_id,
|
chat_id=msg.chat_id,
|
||||||
content=final_content,
|
content=final_content,
|
||||||
metadata=msg.metadata or {}, # Pass through for channel-specific needs (e.g. Slack thread_ts)
|
metadata=outbound_metadata,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _process_system_message(self, msg: InboundMessage) -> OutboundMessage | None:
|
async def _process_system_message(self, msg: InboundMessage) -> OutboundMessage | None:
|
||||||
@@ -308,7 +508,7 @@ class AgentLoop:
|
|||||||
|
|
||||||
spawn_tool = self.tools.get("spawn")
|
spawn_tool = self.tools.get("spawn")
|
||||||
if isinstance(spawn_tool, SpawnTool):
|
if isinstance(spawn_tool, SpawnTool):
|
||||||
spawn_tool.set_context(origin_channel, origin_chat_id)
|
spawn_tool.set_context(origin_channel, origin_chat_id, msg.metadata)
|
||||||
|
|
||||||
cron_tool = self.tools.get("cron")
|
cron_tool = self.tools.get("cron")
|
||||||
if isinstance(cron_tool, CronTool):
|
if isinstance(cron_tool, CronTool):
|
||||||
@@ -321,10 +521,18 @@ class AgentLoop:
|
|||||||
channel=origin_channel,
|
channel=origin_channel,
|
||||||
chat_id=origin_chat_id,
|
chat_id=origin_chat_id,
|
||||||
)
|
)
|
||||||
|
turn_start = len(messages)
|
||||||
|
|
||||||
# Agent loop (limited for announce handling)
|
# Agent loop (limited for announce handling)
|
||||||
iteration = 0
|
iteration = 0
|
||||||
final_content = None
|
final_content = None
|
||||||
|
final_reasoning = None
|
||||||
|
|
||||||
|
# Check for suppress mode BEFORE the loop so it's available for forgery detection
|
||||||
|
suppress_output = msg.metadata.get("suppress_output", False) if msg.metadata else False
|
||||||
|
|
||||||
|
# Select model based on quota
|
||||||
|
selected_model = self._select_model_based_on_quota()
|
||||||
|
|
||||||
while iteration < self.max_iterations:
|
while iteration < self.max_iterations:
|
||||||
iteration += 1
|
iteration += 1
|
||||||
@@ -332,7 +540,8 @@ class AgentLoop:
|
|||||||
response = await self.provider.chat(
|
response = await self.provider.chat(
|
||||||
messages=messages,
|
messages=messages,
|
||||||
tools=self.tools.get_definitions(),
|
tools=self.tools.get_definitions(),
|
||||||
model=self.model
|
model=selected_model,
|
||||||
|
context_management=self.CONTEXT_MANAGEMENT,
|
||||||
)
|
)
|
||||||
|
|
||||||
if response.has_tool_calls:
|
if response.has_tool_calls:
|
||||||
@@ -359,28 +568,80 @@ class AgentLoop:
|
|||||||
messages = self.context.add_tool_result(
|
messages = self.context.add_tool_result(
|
||||||
messages, tool_call.id, tool_call.name, result
|
messages, tool_call.id, tool_call.name, result
|
||||||
)
|
)
|
||||||
# Interleaved CoT: reflect before next action
|
# Interleaved CoT: reflect before next action (skip when thinking is active)
|
||||||
messages.append({"role": "user", "content": "Reflect on the results and decide next steps."})
|
if not getattr(self.provider, 'thinking_budget', 0):
|
||||||
|
messages.append({"role": "user", "content": "Reflect on the results and decide next steps."})
|
||||||
else:
|
else:
|
||||||
|
# No tool calls
|
||||||
final_content = response.content
|
final_content = response.content
|
||||||
|
final_reasoning = response.reasoning_content
|
||||||
|
|
||||||
|
# Check for forged signatures if in suppress mode
|
||||||
|
if suppress_output and has_forged_marker(final_content):
|
||||||
|
# Initialize retry counter if needed
|
||||||
|
if not hasattr(self, '_forge_retry_count_system'):
|
||||||
|
self._forge_retry_count_system = 0
|
||||||
|
|
||||||
|
if self._forge_retry_count_system < 1:
|
||||||
|
# First offense: reject and retry with correction
|
||||||
|
self._forge_retry_count_system += 1
|
||||||
|
logger.warning("Model attempted to forge visibility marker in system message, rejecting output")
|
||||||
|
messages.append({
|
||||||
|
"role": "user",
|
||||||
|
"content": "[System: Previous response rejected. Do not generate [HIDDEN:*] markers.]"
|
||||||
|
})
|
||||||
|
continue # Back to while loop, will retry LLM call
|
||||||
|
else:
|
||||||
|
# Second offense: strip and log error (fallback)
|
||||||
|
logger.error("Model persisted in forging markers in system message despite correction, stripping")
|
||||||
|
final_content = strip_all_hidden_markers(final_content)
|
||||||
|
|
||||||
|
# Reset retry counter on successful completion
|
||||||
|
if hasattr(self, '_forge_retry_count_system'):
|
||||||
|
self._forge_retry_count_system = 0
|
||||||
|
|
||||||
break
|
break
|
||||||
|
|
||||||
if final_content is None:
|
if final_content is None:
|
||||||
final_content = "Background task completed."
|
final_content = "Background task completed."
|
||||||
|
|
||||||
# Save to session (mark as system message in history)
|
# suppress_output already defined before the loop
|
||||||
|
|
||||||
|
if suppress_output:
|
||||||
|
# Sign content with our secret key (forgery check already done in loop)
|
||||||
|
final_content_for_session = sign_content(final_content)
|
||||||
|
# Mark as suppressed for channel handler
|
||||||
|
outbound_metadata = {**(msg.metadata or {}), "suppressed": True}
|
||||||
|
else:
|
||||||
|
final_content_for_session = final_content
|
||||||
|
outbound_metadata = msg.metadata or {}
|
||||||
|
|
||||||
|
# Append final assistant response to messages (use prefixed version for session)
|
||||||
|
messages = self.context.add_assistant_message(
|
||||||
|
messages, final_content_for_session, None,
|
||||||
|
reasoning_content=final_reasoning,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Save to session: user message + full tool chain
|
||||||
session.add_message("user", f"[System: {msg.sender_id}] {msg.content}")
|
session.add_message("user", f"[System: {msg.sender_id}] {msg.content}")
|
||||||
session.add_message("assistant", final_content)
|
for chain_msg in messages[turn_start:]:
|
||||||
|
session.add_raw_message(chain_msg)
|
||||||
self.sessions.save(session)
|
self.sessions.save(session)
|
||||||
|
|
||||||
|
# Return original content (not prefixed) for outbound, but with suppressed metadata
|
||||||
return OutboundMessage(
|
return OutboundMessage(
|
||||||
channel=origin_channel,
|
channel=origin_channel,
|
||||||
chat_id=origin_chat_id,
|
chat_id=origin_chat_id,
|
||||||
content=final_content
|
content=final_content,
|
||||||
|
metadata=outbound_metadata,
|
||||||
)
|
)
|
||||||
|
|
||||||
async def _consolidate_memory(self, session, archive_all: bool = False) -> None:
|
async def _consolidate_memory(self, session, archive_all: bool = False) -> None:
|
||||||
"""Consolidate old messages into MEMORY.md + HISTORY.md, then trim session."""
|
"""Consolidate session into MEMORY.md + HISTORY.md.
|
||||||
|
|
||||||
|
Context window management is now handled server-side via context_management.
|
||||||
|
This only runs on /new to write long-term facts and searchable history.
|
||||||
|
"""
|
||||||
if not session.messages:
|
if not session.messages:
|
||||||
return
|
return
|
||||||
memory = MemoryStore(self.workspace)
|
memory = MemoryStore(self.workspace)
|
||||||
@@ -388,19 +649,50 @@ class AgentLoop:
|
|||||||
old_messages = session.messages
|
old_messages = session.messages
|
||||||
keep_count = 0
|
keep_count = 0
|
||||||
else:
|
else:
|
||||||
|
# Only write truly old messages; keep the recent ones
|
||||||
keep_count = min(10, max(2, self.memory_window // 2))
|
keep_count = min(10, max(2, self.memory_window // 2))
|
||||||
old_messages = session.messages[:-keep_count]
|
old_messages = session.messages[:-keep_count]
|
||||||
if not old_messages:
|
if not old_messages:
|
||||||
return
|
return
|
||||||
logger.info(f"Memory consolidation started: {len(session.messages)} messages, archiving {len(old_messages)}, keeping {keep_count}")
|
logger.info(f"Memory consolidation: archiving {len(old_messages)} messages, keeping {keep_count}")
|
||||||
|
|
||||||
# Format messages for LLM (include tool names when available)
|
# Format messages for LLM — handle full tool chain format
|
||||||
lines = []
|
lines = []
|
||||||
for m in old_messages:
|
for m in old_messages:
|
||||||
if not m.get("content"):
|
role = m.get("role", "?")
|
||||||
|
content = m.get("content")
|
||||||
|
timestamp = m.get("timestamp", "?")[:16]
|
||||||
|
|
||||||
|
if role == "tool":
|
||||||
|
result = str(content or "")[:200]
|
||||||
|
lines.append(f"[{timestamp}] TOOL_RESULT({m.get('name', '?')}): {result}")
|
||||||
continue
|
continue
|
||||||
tools = f" [tools: {', '.join(m['tools_used'])}]" if m.get("tools_used") else ""
|
|
||||||
lines.append(f"[{m.get('timestamp', '?')[:16]}] {m['role'].upper()}{tools}: {m['content']}")
|
# Skip internal reflect prompts
|
||||||
|
if role == "user" and content == "Reflect on the results and decide next steps.":
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Extract text from content (may be list of blocks)
|
||||||
|
if isinstance(content, list):
|
||||||
|
text_parts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"]
|
||||||
|
content_str = " ".join(text_parts)
|
||||||
|
elif isinstance(content, str):
|
||||||
|
content_str = content
|
||||||
|
else:
|
||||||
|
content_str = ""
|
||||||
|
|
||||||
|
# Get tool names from tool_calls or legacy tools_used
|
||||||
|
tool_names = []
|
||||||
|
if m.get("tool_calls"):
|
||||||
|
tool_names = [tc.get("function", {}).get("name", "?") for tc in m["tool_calls"]]
|
||||||
|
elif m.get("tools_used"):
|
||||||
|
tool_names = m["tools_used"]
|
||||||
|
|
||||||
|
if not content_str and not tool_names:
|
||||||
|
continue
|
||||||
|
|
||||||
|
tools_str = f" [tools: {', '.join(tool_names)}]" if tool_names else ""
|
||||||
|
lines.append(f"[{timestamp}] {role.upper()}{tools_str}: {content_str}")
|
||||||
conversation = "\n".join(lines)
|
conversation = "\n".join(lines)
|
||||||
current_memory = memory.read_long_term()
|
current_memory = memory.read_long_term()
|
||||||
|
|
||||||
@@ -424,7 +716,9 @@ Respond with ONLY valid JSON, no markdown fences."""
|
|||||||
{"role": "system", "content": "You are a memory consolidation agent. Respond only with valid JSON."},
|
{"role": "system", "content": "You are a memory consolidation agent. Respond only with valid JSON."},
|
||||||
{"role": "user", "content": prompt},
|
{"role": "user", "content": prompt},
|
||||||
],
|
],
|
||||||
model=self.model,
|
model="claude-haiku-4-5",
|
||||||
|
thinking_budget=0,
|
||||||
|
max_tokens=16384,
|
||||||
)
|
)
|
||||||
text = (response.content or "").strip()
|
text = (response.content or "").strip()
|
||||||
if text.startswith("```"):
|
if text.startswith("```"):
|
||||||
@@ -449,6 +743,7 @@ Respond with ONLY valid JSON, no markdown fences."""
|
|||||||
session_key: str = "cli:direct",
|
session_key: str = "cli:direct",
|
||||||
channel: str = "cli",
|
channel: str = "cli",
|
||||||
chat_id: str = "direct",
|
chat_id: str = "direct",
|
||||||
|
metadata: dict[str, Any] | None = None,
|
||||||
) -> str:
|
) -> str:
|
||||||
"""
|
"""
|
||||||
Process a message directly (for CLI or cron usage).
|
Process a message directly (for CLI or cron usage).
|
||||||
@@ -458,6 +753,7 @@ Respond with ONLY valid JSON, no markdown fences."""
|
|||||||
session_key: Session identifier (overrides channel:chat_id for session lookup).
|
session_key: Session identifier (overrides channel:chat_id for session lookup).
|
||||||
channel: Source channel (for tool context routing).
|
channel: Source channel (for tool context routing).
|
||||||
chat_id: Source chat ID (for tool context routing).
|
chat_id: Source chat ID (for tool context routing).
|
||||||
|
metadata: Optional metadata to pass through (for suppress mode, etc.).
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
The agent's response.
|
The agent's response.
|
||||||
@@ -466,8 +762,24 @@ Respond with ONLY valid JSON, no markdown fences."""
|
|||||||
channel=channel,
|
channel=channel,
|
||||||
sender_id="user",
|
sender_id="user",
|
||||||
chat_id=chat_id,
|
chat_id=chat_id,
|
||||||
content=content
|
content=content,
|
||||||
|
metadata=metadata or {},
|
||||||
)
|
)
|
||||||
|
|
||||||
response = await self._process_message(msg, session_key=session_key)
|
response = await self._process_message(msg, session_key=session_key)
|
||||||
return response.content if response else ""
|
if not response:
|
||||||
|
return ""
|
||||||
|
|
||||||
|
# If suppressed, return signed content from session instead of outbound content
|
||||||
|
if response.metadata.get("suppressed", False):
|
||||||
|
session = self.sessions.get_or_create(session_key)
|
||||||
|
# Get the last assistant message from session (should have signed content)
|
||||||
|
for msg_item in reversed(session.messages):
|
||||||
|
if msg_item.get("role") == "assistant":
|
||||||
|
content = msg_item.get("content", "")
|
||||||
|
if content.startswith("[HIDDEN:"):
|
||||||
|
return content
|
||||||
|
# Fallback to outbound content if signature not found
|
||||||
|
return response.content
|
||||||
|
|
||||||
|
return response.content
|
||||||
|
|||||||
@@ -170,7 +170,7 @@ class SkillsLoader:
|
|||||||
"""Parse nanobot metadata JSON from frontmatter."""
|
"""Parse nanobot metadata JSON from frontmatter."""
|
||||||
try:
|
try:
|
||||||
data = json.loads(raw)
|
data = json.loads(raw)
|
||||||
return data.get("nanobot", {}) if isinstance(data, dict) else {}
|
return (data.get("nanobot") or data.get("clawdbot") or {}) if isinstance(data, dict) else {}
|
||||||
except (json.JSONDecodeError, TypeError):
|
except (json.JSONDecodeError, TypeError):
|
||||||
return {}
|
return {}
|
||||||
|
|
||||||
|
|||||||
+49
-19
@@ -15,6 +15,8 @@ from nanobot.agent.tools.registry import ToolRegistry
|
|||||||
from nanobot.agent.tools.filesystem import ReadFileTool, WriteFileTool, EditFileTool, ListDirTool
|
from nanobot.agent.tools.filesystem import ReadFileTool, WriteFileTool, EditFileTool, ListDirTool
|
||||||
from nanobot.agent.tools.shell import ExecTool
|
from nanobot.agent.tools.shell import ExecTool
|
||||||
from nanobot.agent.tools.web import WebSearchTool, WebFetchTool
|
from nanobot.agent.tools.web import WebSearchTool, WebFetchTool
|
||||||
|
from nanobot.agent.tools.spawn import SpawnTool
|
||||||
|
from nanobot.agent.tools.wait import WaitForSubagentsTool
|
||||||
|
|
||||||
|
|
||||||
class SubagentManager:
|
class SubagentManager:
|
||||||
@@ -40,16 +42,21 @@ class SubagentManager:
|
|||||||
self.provider = provider
|
self.provider = provider
|
||||||
self.workspace = workspace
|
self.workspace = workspace
|
||||||
self.bus = bus
|
self.bus = bus
|
||||||
self.model = model or provider.get_default_model()
|
# Default to Sonnet, not the provider default (Opus).
|
||||||
|
# Quota switching only affects the main agent's own requests, not SubagentManager.
|
||||||
|
# Explicit model overrides (e.g. Haiku workers) still take precedence.
|
||||||
|
self.model = model or "claude-sonnet-4-6"
|
||||||
self.brave_api_key = brave_api_key
|
self.brave_api_key = brave_api_key
|
||||||
self.exec_config = exec_config or ExecToolConfig()
|
self.exec_config = exec_config or ExecToolConfig()
|
||||||
self.restrict_to_workspace = restrict_to_workspace
|
self.restrict_to_workspace = restrict_to_workspace
|
||||||
self._running_tasks: dict[str, asyncio.Task[None]] = {}
|
self._running_tasks: dict[str, asyncio.Task[None]] = {}
|
||||||
|
self._task_results: dict[str, str] = {}
|
||||||
|
|
||||||
async def spawn(
|
async def spawn(
|
||||||
self,
|
self,
|
||||||
task: str,
|
task: str,
|
||||||
label: str | None = None,
|
label: str | None = None,
|
||||||
|
model: str | None = None,
|
||||||
origin_channel: str = "cli",
|
origin_channel: str = "cli",
|
||||||
origin_chat_id: str = "direct",
|
origin_chat_id: str = "direct",
|
||||||
) -> str:
|
) -> str:
|
||||||
@@ -75,7 +82,7 @@ class SubagentManager:
|
|||||||
|
|
||||||
# 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)
|
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
|
||||||
|
|
||||||
@@ -83,7 +90,7 @@ class SubagentManager:
|
|||||||
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 (id: {task_id}). I'll notify you when it completes."
|
return f"Subagent [{display_label}] started. Task ID: {task_id}"
|
||||||
|
|
||||||
async def _run_subagent(
|
async def _run_subagent(
|
||||||
self,
|
self,
|
||||||
@@ -91,12 +98,13 @@ class SubagentManager:
|
|||||||
task: str,
|
task: str,
|
||||||
label: str,
|
label: str,
|
||||||
origin: dict[str, str],
|
origin: dict[str, str],
|
||||||
|
model: str | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Execute the subagent task and announce the result."""
|
"""Execute the subagent task and announce the result."""
|
||||||
logger.info(f"Subagent [{task_id}] starting task: {label}")
|
logger.info(f"Subagent [{task_id}] starting task: {label}")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Build subagent tools (no message tool, no spawn tool)
|
# Build subagent tools (no message tool)
|
||||||
tools = ToolRegistry()
|
tools = ToolRegistry()
|
||||||
allowed_dir = self.workspace if self.restrict_to_workspace else None
|
allowed_dir = self.workspace if self.restrict_to_workspace else None
|
||||||
tools.register(ReadFileTool(allowed_dir=allowed_dir))
|
tools.register(ReadFileTool(allowed_dir=allowed_dir))
|
||||||
@@ -110,6 +118,10 @@ class SubagentManager:
|
|||||||
))
|
))
|
||||||
tools.register(WebSearchTool(api_key=self.brave_api_key))
|
tools.register(WebSearchTool(api_key=self.brave_api_key))
|
||||||
tools.register(WebFetchTool())
|
tools.register(WebFetchTool())
|
||||||
|
spawn_tool = SpawnTool(manager=self)
|
||||||
|
spawn_tool.set_context("subagent", origin["chat_id"])
|
||||||
|
tools.register(spawn_tool)
|
||||||
|
tools.register(WaitForSubagentsTool(manager=self))
|
||||||
|
|
||||||
# Build messages with subagent-specific prompt
|
# Build messages with subagent-specific prompt
|
||||||
system_prompt = self._build_subagent_prompt(task)
|
system_prompt = self._build_subagent_prompt(task)
|
||||||
@@ -119,7 +131,7 @@ class SubagentManager:
|
|||||||
]
|
]
|
||||||
|
|
||||||
# Run agent loop (limited iterations)
|
# Run agent loop (limited iterations)
|
||||||
max_iterations = 15
|
max_iterations = 50
|
||||||
iteration = 0
|
iteration = 0
|
||||||
final_result: str | None = None
|
final_result: str | None = None
|
||||||
|
|
||||||
@@ -129,7 +141,7 @@ class SubagentManager:
|
|||||||
response = await self.provider.chat(
|
response = await self.provider.chat(
|
||||||
messages=messages,
|
messages=messages,
|
||||||
tools=tools.get_definitions(),
|
tools=tools.get_definitions(),
|
||||||
model=self.model,
|
model=model or self.model,
|
||||||
)
|
)
|
||||||
|
|
||||||
if response.has_tool_calls:
|
if response.has_tool_calls:
|
||||||
@@ -189,6 +201,13 @@ class SubagentManager:
|
|||||||
"""Announce the subagent result to the main agent via the message bus."""
|
"""Announce the subagent result to the main agent via the message bus."""
|
||||||
status_text = "completed successfully" if status == "ok" else "failed"
|
status_text = "completed successfully" if status == "ok" else "failed"
|
||||||
|
|
||||||
|
# Child subagents (spawned by other subagents) store results silently.
|
||||||
|
# The parent orchestrator collects them via wait_for_subagents.
|
||||||
|
if origin["channel"] == "subagent":
|
||||||
|
self._task_results[task_id] = result
|
||||||
|
logger.debug(f"Subagent [{task_id}] stored result silently (child subagent)")
|
||||||
|
return
|
||||||
|
|
||||||
announce_content = f"""[Subagent '{label}' {status_text}]
|
announce_content = f"""[Subagent '{label}' {status_text}]
|
||||||
|
|
||||||
Task: {task}
|
Task: {task}
|
||||||
@@ -211,23 +230,16 @@ Summarize this naturally for the user. Keep it brief (1-2 sentences). Do not men
|
|||||||
|
|
||||||
def _build_subagent_prompt(self, task: str) -> str:
|
def _build_subagent_prompt(self, task: str) -> str:
|
||||||
"""Build a focused system prompt for the subagent."""
|
"""Build a focused system prompt for the subagent."""
|
||||||
from datetime import datetime
|
|
||||||
import time as _time
|
|
||||||
now = datetime.now().strftime("%Y-%m-%d %H:%M (%A)")
|
|
||||||
tz = _time.strftime("%Z") or "UTC"
|
|
||||||
|
|
||||||
return f"""# Subagent
|
return f"""# Subagent
|
||||||
|
|
||||||
## Current Time
|
|
||||||
{now} ({tz})
|
|
||||||
|
|
||||||
You are a subagent spawned by the main agent to complete a specific task.
|
You are a subagent spawned by the main agent to complete a specific task.
|
||||||
|
|
||||||
## Rules
|
## Rules
|
||||||
1. Stay focused - complete only the assigned task, nothing else
|
1. Run `exec date` as your very first action to get the current date and time
|
||||||
2. Your final response will be reported back to the main agent
|
2. Stay focused - complete only the assigned task, nothing else
|
||||||
3. Do not initiate conversations or take on side tasks
|
3. Your final response will be reported back to the main agent
|
||||||
4. Be concise but informative in your findings
|
4. Do not initiate conversations or take on side tasks
|
||||||
|
5. Be concise but informative in your findings
|
||||||
|
|
||||||
## What You Can Do
|
## What You Can Do
|
||||||
- Read and write files in the workspace
|
- Read and write files in the workspace
|
||||||
@@ -237,7 +249,6 @@ You are a subagent spawned by the main agent to complete a specific task.
|
|||||||
|
|
||||||
## What You Cannot Do
|
## What You Cannot Do
|
||||||
- Send messages directly to users (no message tool available)
|
- Send messages directly to users (no message tool available)
|
||||||
- Spawn other subagents
|
|
||||||
- Access the main agent's conversation history
|
- Access the main agent's conversation history
|
||||||
|
|
||||||
## Workspace
|
## Workspace
|
||||||
@@ -246,6 +257,25 @@ Skills are available at: {self.workspace}/skills/ (read SKILL.md files as needed
|
|||||||
|
|
||||||
When you have completed the task, provide a clear summary of your findings or actions."""
|
When you have completed the task, provide a clear summary of your findings or actions."""
|
||||||
|
|
||||||
|
async def wait_for(self, task_ids: list[str]) -> str:
|
||||||
|
"""Wait for specified child subagents to complete and return their results."""
|
||||||
|
tasks_to_wait = [
|
||||||
|
self._running_tasks[tid]
|
||||||
|
for tid in task_ids
|
||||||
|
if tid in self._running_tasks
|
||||||
|
]
|
||||||
|
if tasks_to_wait:
|
||||||
|
await asyncio.gather(*tasks_to_wait, return_exceptions=True)
|
||||||
|
|
||||||
|
results = []
|
||||||
|
for tid in task_ids:
|
||||||
|
result = self._task_results.get(tid)
|
||||||
|
if result is not None:
|
||||||
|
results.append(f"[{tid}]:\n{result}")
|
||||||
|
else:
|
||||||
|
results.append(f"[{tid}]: No result found (invalid ID or task failed before storing)")
|
||||||
|
return "\n\n---\n\n".join(results)
|
||||||
|
|
||||||
def get_running_count(self) -> int:
|
def get_running_count(self) -> int:
|
||||||
"""Return the number of currently running subagents."""
|
"""Return the number of currently running subagents."""
|
||||||
return len(self._running_tasks)
|
return len(self._running_tasks)
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ from typing import Any, Callable, Awaitable
|
|||||||
|
|
||||||
from nanobot.agent.tools.base import Tool
|
from nanobot.agent.tools.base import Tool
|
||||||
from nanobot.bus.events import OutboundMessage
|
from nanobot.bus.events import OutboundMessage
|
||||||
|
from nanobot.session import SessionManager
|
||||||
|
|
||||||
|
|
||||||
class MessageTool(Tool):
|
class MessageTool(Tool):
|
||||||
@@ -12,10 +13,12 @@ class MessageTool(Tool):
|
|||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None,
|
send_callback: Callable[[OutboundMessage], Awaitable[None]] | None = None,
|
||||||
|
sessions: SessionManager | None = None,
|
||||||
default_channel: str = "",
|
default_channel: str = "",
|
||||||
default_chat_id: str = ""
|
default_chat_id: str = ""
|
||||||
):
|
):
|
||||||
self._send_callback = send_callback
|
self._send_callback = send_callback
|
||||||
|
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
|
||||||
|
|
||||||
@@ -81,6 +84,13 @@ class MessageTool(Tool):
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
await self._send_callback(msg)
|
await self._send_callback(msg)
|
||||||
|
|
||||||
|
if self._sessions:
|
||||||
|
session_key = f"{channel}:{chat_id}"
|
||||||
|
session = self._sessions.get_or_create(session_key)
|
||||||
|
session.add_message("assistant", content)
|
||||||
|
self._sessions.save(session)
|
||||||
|
|
||||||
return f"Message sent to {channel}:{chat_id}"
|
return f"Message sent to {channel}:{chat_id}"
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return f"Error sending message: {str(e)}"
|
return f"Error sending message: {str(e)}"
|
||||||
|
|||||||
@@ -51,15 +51,20 @@ class SpawnTool(Tool):
|
|||||||
"type": "string",
|
"type": "string",
|
||||||
"description": "Optional short label for the task (for display)",
|
"description": "Optional short label for the task (for display)",
|
||||||
},
|
},
|
||||||
|
"model": {
|
||||||
|
"type": "string",
|
||||||
|
"description": "Optional model override for the subagent (e.g. 'claude-haiku-4-5'). Defaults to the main agent's model.",
|
||||||
|
},
|
||||||
},
|
},
|
||||||
"required": ["task"],
|
"required": ["task"],
|
||||||
}
|
}
|
||||||
|
|
||||||
async def execute(self, task: str, label: str | None = None, **kwargs: Any) -> str:
|
async def execute(self, task: str, label: str | None = None, model: str | None = None, **kwargs: Any) -> str:
|
||||||
"""Spawn a subagent to execute the given task."""
|
"""Spawn a subagent to execute the given task."""
|
||||||
return await self._manager.spawn(
|
return await self._manager.spawn(
|
||||||
task=task,
|
task=task,
|
||||||
label=label,
|
label=label,
|
||||||
|
model=model,
|
||||||
origin_channel=self._origin_channel,
|
origin_channel=self._origin_channel,
|
||||||
origin_chat_id=self._origin_chat_id,
|
origin_chat_id=self._origin_chat_id,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -0,0 +1,50 @@
|
|||||||
|
"""Wait-for-subagents tool for orchestrator subagents."""
|
||||||
|
|
||||||
|
from typing import Any, TYPE_CHECKING
|
||||||
|
|
||||||
|
from nanobot.agent.tools.base import Tool
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from nanobot.agent.subagent import SubagentManager
|
||||||
|
|
||||||
|
|
||||||
|
class WaitForSubagentsTool(Tool):
|
||||||
|
"""
|
||||||
|
Tool to wait for child subagents to complete and collect their results.
|
||||||
|
|
||||||
|
Use this after spawning multiple subagents to wait for all of them
|
||||||
|
and get their results for synthesis.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, manager: "SubagentManager"):
|
||||||
|
self._manager = manager
|
||||||
|
|
||||||
|
@property
|
||||||
|
def name(self) -> str:
|
||||||
|
return "wait_for_subagents"
|
||||||
|
|
||||||
|
@property
|
||||||
|
def description(self) -> str:
|
||||||
|
return (
|
||||||
|
"Wait for one or more child subagents to complete and return their results. "
|
||||||
|
"Use this after spawning subagents to collect all results before synthesizing. "
|
||||||
|
"Blocks until all specified subagents finish."
|
||||||
|
)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def parameters(self) -> dict[str, Any]:
|
||||||
|
return {
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"task_ids": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {"type": "string"},
|
||||||
|
"description": "List of task IDs to wait for (from spawn tool responses)",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"required": ["task_ids"],
|
||||||
|
}
|
||||||
|
|
||||||
|
async def execute(self, task_ids: list[str], **kwargs: Any) -> str:
|
||||||
|
"""Wait for the specified subagents and return their results."""
|
||||||
|
return await self._manager.wait_for(task_ids)
|
||||||
@@ -0,0 +1,83 @@
|
|||||||
|
# nanobot/agent/visibility.py
|
||||||
|
"""Cryptographic signing for visibility markers to prevent model forgery."""
|
||||||
|
|
||||||
|
import hmac
|
||||||
|
import hashlib
|
||||||
|
import re
|
||||||
|
from typing import Tuple
|
||||||
|
|
||||||
|
SECRET_KEY = "nanobot_visibility_secret_key_v1"
|
||||||
|
|
||||||
|
|
||||||
|
def sign_content(content: str) -> str:
|
||||||
|
"""
|
||||||
|
Sign content with HMAC and prepend marker.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content: The message content to sign
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Content with signed visibility marker: "[HIDDEN:{sig}] {content}"
|
||||||
|
"""
|
||||||
|
sig = hmac.new(
|
||||||
|
SECRET_KEY.encode(),
|
||||||
|
content.encode(),
|
||||||
|
hashlib.sha256
|
||||||
|
).hexdigest()[:8]
|
||||||
|
return f"[HIDDEN:{sig}] {content}"
|
||||||
|
|
||||||
|
|
||||||
|
def verify_signature(marked_content: str) -> Tuple[bool, str]:
|
||||||
|
"""
|
||||||
|
Verify HMAC signature and extract clean content.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
marked_content: Content potentially with [HIDDEN:{sig}] marker
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Tuple of (is_valid, clean_content)
|
||||||
|
- is_valid: True if signature is valid, False otherwise
|
||||||
|
- clean_content: Content without marker
|
||||||
|
"""
|
||||||
|
match = re.match(r'\[HIDDEN:([a-f0-9]{8})\] (.*)', marked_content, re.DOTALL)
|
||||||
|
if not match:
|
||||||
|
return False, marked_content
|
||||||
|
|
||||||
|
claimed_sig, content = match.groups()
|
||||||
|
expected_sig = hmac.new(
|
||||||
|
SECRET_KEY.encode(),
|
||||||
|
content.encode(),
|
||||||
|
hashlib.sha256
|
||||||
|
).hexdigest()[:8]
|
||||||
|
|
||||||
|
is_valid = hmac.compare_digest(claimed_sig, expected_sig)
|
||||||
|
return is_valid, content
|
||||||
|
|
||||||
|
|
||||||
|
def has_forged_marker(content: str) -> bool:
|
||||||
|
"""
|
||||||
|
Check if content has an invalid [HIDDEN:*] marker at the start.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content: Content to check
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
True if content starts with forged marker, False otherwise
|
||||||
|
"""
|
||||||
|
if not content.startswith("[HIDDEN:"):
|
||||||
|
return False
|
||||||
|
is_valid, _ = verify_signature(content)
|
||||||
|
return not is_valid
|
||||||
|
|
||||||
|
|
||||||
|
def strip_all_hidden_markers(content: str) -> str:
|
||||||
|
"""
|
||||||
|
Remove all [HIDDEN:*] patterns from content (valid or invalid).
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content: Content potentially with markers
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Content with all markers stripped
|
||||||
|
"""
|
||||||
|
return re.sub(r'\[HIDDEN:[a-f0-9]{8}\]\s*', '', content)
|
||||||
@@ -20,6 +20,7 @@ class MessageBus:
|
|||||||
self.inbound: asyncio.Queue[InboundMessage] = asyncio.Queue()
|
self.inbound: asyncio.Queue[InboundMessage] = asyncio.Queue()
|
||||||
self.outbound: asyncio.Queue[OutboundMessage] = asyncio.Queue()
|
self.outbound: asyncio.Queue[OutboundMessage] = asyncio.Queue()
|
||||||
self._outbound_subscribers: dict[str, list[Callable[[OutboundMessage], Awaitable[None]]]] = {}
|
self._outbound_subscribers: dict[str, list[Callable[[OutboundMessage], Awaitable[None]]]] = {}
|
||||||
|
self._correlation_store: dict[str, asyncio.Future] = {}
|
||||||
self._running = False
|
self._running = False
|
||||||
|
|
||||||
async def publish_inbound(self, msg: InboundMessage) -> None:
|
async def publish_inbound(self, msg: InboundMessage) -> None:
|
||||||
@@ -38,6 +39,27 @@ class MessageBus:
|
|||||||
"""Consume the next outbound message (blocks until available)."""
|
"""Consume the next outbound message (blocks until available)."""
|
||||||
return await self.outbound.get()
|
return await self.outbound.get()
|
||||||
|
|
||||||
|
def register_correlation(self, correlation_id: str) -> asyncio.Future:
|
||||||
|
"""Register a Future to be resolved when a matching outbound message appears."""
|
||||||
|
loop = asyncio.get_running_loop()
|
||||||
|
future = loop.create_future()
|
||||||
|
self._correlation_store[correlation_id] = future
|
||||||
|
return future
|
||||||
|
|
||||||
|
def resolve_correlation(self, msg: OutboundMessage) -> None:
|
||||||
|
"""Check if an outbound message has a correlation_id and resolve the matching Future."""
|
||||||
|
cid = msg.metadata.get("correlation_id") if msg.metadata else None
|
||||||
|
if cid and cid in self._correlation_store:
|
||||||
|
future = self._correlation_store.pop(cid)
|
||||||
|
if not future.done():
|
||||||
|
future.set_result(msg.content)
|
||||||
|
|
||||||
|
def cancel_correlation(self, correlation_id: str) -> None:
|
||||||
|
"""Cancel and remove a pending correlation."""
|
||||||
|
future = self._correlation_store.pop(correlation_id, None)
|
||||||
|
if future and not future.done():
|
||||||
|
future.cancel()
|
||||||
|
|
||||||
def subscribe_outbound(
|
def subscribe_outbound(
|
||||||
self,
|
self,
|
||||||
channel: str,
|
channel: str,
|
||||||
|
|||||||
@@ -74,6 +74,10 @@ class BaseChannel(ABC):
|
|||||||
if not allow_list:
|
if not allow_list:
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
# Wildcard allows everyone
|
||||||
|
if "*" in allow_list:
|
||||||
|
return True
|
||||||
|
|
||||||
sender_str = str(sender_id)
|
sender_str = str(sender_id)
|
||||||
if sender_str in allow_list:
|
if sender_str in allow_list:
|
||||||
return True
|
return True
|
||||||
|
|||||||
@@ -0,0 +1,38 @@
|
|||||||
|
"""Hook channel — receives outbound messages from hook-initiated conversations."""
|
||||||
|
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from nanobot.bus.events import OutboundMessage
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
|
||||||
|
|
||||||
|
class HookChannel:
|
||||||
|
"""
|
||||||
|
Minimal channel for hook-initiated conversations.
|
||||||
|
|
||||||
|
The hook HTTP server publishes InboundMessages to the bus.
|
||||||
|
Responses come back as OutboundMessages routed here.
|
||||||
|
send() is a no-op because the HTTP caller gets the response
|
||||||
|
via bus correlation, not channel delivery.
|
||||||
|
"""
|
||||||
|
|
||||||
|
name = "hook"
|
||||||
|
|
||||||
|
def __init__(self, bus: MessageBus):
|
||||||
|
self.bus = bus
|
||||||
|
self._running = False
|
||||||
|
|
||||||
|
async def start(self) -> None:
|
||||||
|
self._running = True
|
||||||
|
logger.info("Hook channel started")
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
self._running = False
|
||||||
|
|
||||||
|
async def send(self, msg: OutboundMessage) -> None:
|
||||||
|
"""No-op — response is returned via bus correlation to the HTTP caller."""
|
||||||
|
logger.debug(f"Hook channel received outbound for {msg.chat_id} (no-op)")
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_running(self) -> bool:
|
||||||
|
return self._running
|
||||||
@@ -137,6 +137,11 @@ class ChannelManager:
|
|||||||
except ImportError as e:
|
except ImportError as e:
|
||||||
logger.warning(f"QQ channel not available: {e}")
|
logger.warning(f"QQ channel not available: {e}")
|
||||||
|
|
||||||
|
def register_channel(self, name: str, channel: BaseChannel) -> None:
|
||||||
|
"""Register an external channel."""
|
||||||
|
self.channels[name] = channel
|
||||||
|
logger.info(f"{name} channel registered")
|
||||||
|
|
||||||
async def _start_channel(self, name: str, channel: BaseChannel) -> None:
|
async def _start_channel(self, name: str, channel: BaseChannel) -> None:
|
||||||
"""Start a channel and log any exceptions."""
|
"""Start a channel and log any exceptions."""
|
||||||
try:
|
try:
|
||||||
@@ -193,6 +198,9 @@ class ChannelManager:
|
|||||||
timeout=1.0
|
timeout=1.0
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Resolve any pending correlation (hook request-response)
|
||||||
|
self.bus.resolve_correlation(msg)
|
||||||
|
|
||||||
channel = self.channels.get(msg.channel)
|
channel = self.channels.get(msg.channel)
|
||||||
if channel:
|
if channel:
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -92,6 +92,7 @@ class TelegramChannel(BaseChannel):
|
|||||||
BotCommand("start", "Start the bot"),
|
BotCommand("start", "Start the bot"),
|
||||||
BotCommand("new", "Start a new conversation"),
|
BotCommand("new", "Start a new conversation"),
|
||||||
BotCommand("help", "Show available commands"),
|
BotCommand("help", "Show available commands"),
|
||||||
|
BotCommand("quota", "Show current quota status"),
|
||||||
]
|
]
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -127,6 +128,7 @@ class TelegramChannel(BaseChannel):
|
|||||||
self._app.add_handler(CommandHandler("start", self._on_start))
|
self._app.add_handler(CommandHandler("start", self._on_start))
|
||||||
self._app.add_handler(CommandHandler("new", self._forward_command))
|
self._app.add_handler(CommandHandler("new", self._forward_command))
|
||||||
self._app.add_handler(CommandHandler("help", self._forward_command))
|
self._app.add_handler(CommandHandler("help", self._forward_command))
|
||||||
|
self._app.add_handler(CommandHandler("quota", self._forward_command))
|
||||||
|
|
||||||
# Add message handler for text, photos, voice, documents
|
# Add message handler for text, photos, voice, documents
|
||||||
self._app.add_handler(
|
self._app.add_handler(
|
||||||
@@ -187,6 +189,11 @@ class TelegramChannel(BaseChannel):
|
|||||||
# Stop typing indicator for this chat
|
# Stop typing indicator for this chat
|
||||||
self._stop_typing(msg.chat_id)
|
self._stop_typing(msg.chat_id)
|
||||||
|
|
||||||
|
# Check for suppression
|
||||||
|
if msg.metadata.get("suppressed", False):
|
||||||
|
logger.debug(f"Suppressed output (not sent to Telegram): {msg.content[:100]}...")
|
||||||
|
return # Don't send to Telegram API
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# chat_id should be the Telegram chat ID (integer)
|
# chat_id should be the Telegram chat ID (integer)
|
||||||
chat_id = int(msg.chat_id)
|
chat_id = int(msg.chat_id)
|
||||||
|
|||||||
+64
-24
@@ -2,23 +2,24 @@
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import os
|
||||||
import signal
|
|
||||||
from pathlib import Path
|
|
||||||
import select
|
import select
|
||||||
|
import signal
|
||||||
import sys
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
import typer
|
import typer
|
||||||
|
from prompt_toolkit import PromptSession
|
||||||
|
from prompt_toolkit.formatted_text import HTML
|
||||||
|
from prompt_toolkit.history import FileHistory
|
||||||
|
from prompt_toolkit.patch_stdout import patch_stdout
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
from rich.markdown import Markdown
|
from rich.markdown import Markdown
|
||||||
from rich.table import Table
|
from rich.table import Table
|
||||||
from rich.text import Text
|
from rich.text import Text
|
||||||
|
|
||||||
from prompt_toolkit import PromptSession
|
from nanobot import __logo__, __version__
|
||||||
from prompt_toolkit.formatted_text import HTML
|
from nanobot.cli.oauth import oauth_app
|
||||||
from prompt_toolkit.history import FileHistory
|
|
||||||
from prompt_toolkit.patch_stdout import patch_stdout
|
|
||||||
|
|
||||||
from nanobot import __version__, __logo__
|
|
||||||
|
|
||||||
app = typer.Typer(
|
app = typer.Typer(
|
||||||
name="nanobot",
|
name="nanobot",
|
||||||
@@ -270,20 +271,21 @@ This file stores important information that should persist across sessions.
|
|||||||
|
|
||||||
|
|
||||||
def _make_provider(config):
|
def _make_provider(config):
|
||||||
"""Create LiteLLMProvider from config. Exits if no API key found."""
|
"""Create LLM provider from config. Uses OAuth for subscription tokens."""
|
||||||
from nanobot.providers.litellm_provider import LiteLLMProvider
|
from nanobot.providers import create_provider
|
||||||
p = config.get_provider()
|
p = config.get_provider()
|
||||||
model = config.agents.defaults.model
|
model = config.agents.defaults.model
|
||||||
if not (p and p.api_key) and not model.startswith("bedrock/"):
|
if not (p and p.api_key) and not model.startswith("bedrock/"):
|
||||||
console.print("[red]Error: No API key configured.[/red]")
|
console.print("[red]Error: No API key configured.[/red]")
|
||||||
console.print("Set one in ~/.nanobot/config.json under providers section")
|
console.print("Set one in ~/.nanobot/config.json under providers section")
|
||||||
raise typer.Exit(1)
|
raise typer.Exit(1)
|
||||||
return LiteLLMProvider(
|
return create_provider(
|
||||||
api_key=p.api_key if p else None,
|
api_key=p.api_key if p else None,
|
||||||
|
model=model,
|
||||||
api_base=config.get_api_base(),
|
api_base=config.get_api_base(),
|
||||||
default_model=model,
|
|
||||||
extra_headers=p.extra_headers if p else None,
|
extra_headers=p.extra_headers if p else None,
|
||||||
provider_name=config.get_provider_name(),
|
provider_name=config.get_provider_name(),
|
||||||
|
thinking_budget=config.agents.defaults.thinking_budget,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -298,14 +300,14 @@ def gateway(
|
|||||||
verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose output"),
|
verbose: bool = typer.Option(False, "--verbose", "-v", help="Verbose output"),
|
||||||
):
|
):
|
||||||
"""Start the nanobot gateway."""
|
"""Start the nanobot gateway."""
|
||||||
from nanobot.config.loader import load_config, get_data_dir
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
from nanobot.agent.loop import AgentLoop
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
from nanobot.channels.manager import ChannelManager
|
from nanobot.channels.manager import ChannelManager
|
||||||
from nanobot.session.manager import SessionManager
|
from nanobot.config.loader import get_data_dir, load_config
|
||||||
from nanobot.cron.service import CronService
|
from nanobot.cron.service import CronService
|
||||||
from nanobot.cron.types import CronJob
|
from nanobot.cron.types import CronJob
|
||||||
from nanobot.heartbeat.service import HeartbeatService
|
from nanobot.heartbeat.service import HeartbeatService
|
||||||
|
from nanobot.session.manager import SessionManager
|
||||||
|
|
||||||
if verbose:
|
if verbose:
|
||||||
import logging
|
import logging
|
||||||
@@ -357,20 +359,50 @@ def gateway(
|
|||||||
cron.on_job = on_cron_job
|
cron.on_job = on_cron_job
|
||||||
|
|
||||||
# Create heartbeat service
|
# Create heartbeat service
|
||||||
async def on_heartbeat(prompt: str) -> str:
|
async def on_heartbeat(prompt: str, metadata: dict[str, Any] | None = None) -> str:
|
||||||
"""Execute heartbeat through the agent."""
|
"""Execute heartbeat through the agent."""
|
||||||
return await agent.process_direct(prompt, session_key="heartbeat")
|
return await agent.process_direct(
|
||||||
|
prompt,
|
||||||
|
session_key="telegram:239824268", # Run in main telegram session
|
||||||
|
channel="telegram",
|
||||||
|
chat_id="239824268",
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
|
||||||
heartbeat = HeartbeatService(
|
heartbeat = HeartbeatService(
|
||||||
workspace=config.workspace_path,
|
workspace=config.workspace_path,
|
||||||
on_heartbeat=on_heartbeat,
|
on_heartbeat=on_heartbeat,
|
||||||
interval_s=30 * 60, # 30 minutes
|
interval_s=30 * 60, # 30 minutes
|
||||||
enabled=True
|
enabled=True,
|
||||||
|
session_manager=session_manager, # Pass session manager
|
||||||
|
target_session_key="telegram:239824268", # Target session
|
||||||
|
idle_threshold_s=30 * 60, # 30 minutes idle
|
||||||
)
|
)
|
||||||
|
|
||||||
# Create channel manager
|
# Create channel manager
|
||||||
channels = ChannelManager(config, bus)
|
channels = ChannelManager(config, bus)
|
||||||
|
|
||||||
|
# Create hooks server
|
||||||
|
from nanobot.channels.hook import HookChannel
|
||||||
|
from nanobot.hooks.server import HooksServer
|
||||||
|
|
||||||
|
hooks_config = config.hooks if hasattr(config, 'hooks') else None
|
||||||
|
hooks_server = None
|
||||||
|
|
||||||
|
if hooks_config and hooks_config.enabled:
|
||||||
|
# Register hook channel
|
||||||
|
hook_channel = HookChannel(bus)
|
||||||
|
channels.register_channel("hook", hook_channel)
|
||||||
|
|
||||||
|
# Create hooks server (checks has_tokens internally)
|
||||||
|
hooks_server = HooksServer(
|
||||||
|
host=config.gateway.host,
|
||||||
|
port=config.gateway.port,
|
||||||
|
config=hooks_config,
|
||||||
|
bus=bus,
|
||||||
|
)
|
||||||
|
console.print(f"[green]✓[/green] Hooks: {hooks_config.path} on port {config.gateway.port}")
|
||||||
|
|
||||||
if channels.enabled_channels:
|
if channels.enabled_channels:
|
||||||
console.print(f"[green]✓[/green] Channels enabled: {', '.join(channels.enabled_channels)}")
|
console.print(f"[green]✓[/green] Channels enabled: {', '.join(channels.enabled_channels)}")
|
||||||
else:
|
else:
|
||||||
@@ -380,18 +412,22 @@ def gateway(
|
|||||||
if cron_status["jobs"] > 0:
|
if cron_status["jobs"] > 0:
|
||||||
console.print(f"[green]✓[/green] Cron: {cron_status['jobs']} scheduled jobs")
|
console.print(f"[green]✓[/green] Cron: {cron_status['jobs']} scheduled jobs")
|
||||||
|
|
||||||
console.print(f"[green]✓[/green] Heartbeat: every 30m")
|
console.print("[green]✓[/green] Heartbeat: every 30m")
|
||||||
|
|
||||||
async def run():
|
async def run():
|
||||||
try:
|
try:
|
||||||
await cron.start()
|
await cron.start()
|
||||||
await heartbeat.start()
|
await heartbeat.start()
|
||||||
|
if hooks_server:
|
||||||
|
await hooks_server.start()
|
||||||
await asyncio.gather(
|
await asyncio.gather(
|
||||||
agent.run(),
|
agent.run(),
|
||||||
channels.start_all(),
|
channels.start_all(),
|
||||||
)
|
)
|
||||||
except KeyboardInterrupt:
|
except KeyboardInterrupt:
|
||||||
console.print("\nShutting down...")
|
console.print("\nShutting down...")
|
||||||
|
if hooks_server:
|
||||||
|
await hooks_server.stop()
|
||||||
heartbeat.stop()
|
heartbeat.stop()
|
||||||
cron.stop()
|
cron.stop()
|
||||||
agent.stop()
|
agent.stop()
|
||||||
@@ -415,11 +451,12 @@ def agent(
|
|||||||
logs: bool = typer.Option(False, "--logs/--no-logs", help="Show nanobot runtime logs during chat"),
|
logs: bool = typer.Option(False, "--logs/--no-logs", help="Show nanobot runtime logs during chat"),
|
||||||
):
|
):
|
||||||
"""Interact with the agent directly."""
|
"""Interact with the agent directly."""
|
||||||
from nanobot.config.loader import load_config
|
|
||||||
from nanobot.bus.queue import MessageBus
|
|
||||||
from nanobot.agent.loop import AgentLoop
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
|
from nanobot.agent.loop import AgentLoop
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.config.loader import load_config
|
||||||
|
|
||||||
config = load_config()
|
config = load_config()
|
||||||
|
|
||||||
bus = MessageBus()
|
bus = MessageBus()
|
||||||
@@ -507,6 +544,8 @@ def agent(
|
|||||||
channels_app = typer.Typer(help="Manage channels")
|
channels_app = typer.Typer(help="Manage channels")
|
||||||
app.add_typer(channels_app, name="channels")
|
app.add_typer(channels_app, name="channels")
|
||||||
|
|
||||||
|
app.add_typer(oauth_app, name="oauth")
|
||||||
|
|
||||||
|
|
||||||
@channels_app.command("status")
|
@channels_app.command("status")
|
||||||
def channels_status():
|
def channels_status():
|
||||||
@@ -636,6 +675,7 @@ def _get_bridge_dir() -> Path:
|
|||||||
def channels_login():
|
def channels_login():
|
||||||
"""Link device via QR code."""
|
"""Link device via QR code."""
|
||||||
import subprocess
|
import subprocess
|
||||||
|
|
||||||
from nanobot.config.loader import load_config
|
from nanobot.config.loader import load_config
|
||||||
|
|
||||||
config = load_config()
|
config = load_config()
|
||||||
@@ -808,7 +848,7 @@ def cron_run(
|
|||||||
return await service.run_job(job_id, force=force)
|
return await service.run_job(job_id, force=force)
|
||||||
|
|
||||||
if asyncio.run(run()):
|
if asyncio.run(run()):
|
||||||
console.print(f"[green]✓[/green] Job executed")
|
console.print("[green]✓[/green] Job executed")
|
||||||
else:
|
else:
|
||||||
console.print(f"[red]Failed to run job {job_id}[/red]")
|
console.print(f"[red]Failed to run job {job_id}[/red]")
|
||||||
|
|
||||||
@@ -821,7 +861,7 @@ def cron_run(
|
|||||||
@app.command()
|
@app.command()
|
||||||
def status():
|
def status():
|
||||||
"""Show nanobot status."""
|
"""Show nanobot status."""
|
||||||
from nanobot.config.loader import load_config, get_config_path
|
from nanobot.config.loader import get_config_path, load_config
|
||||||
|
|
||||||
config_path = get_config_path()
|
config_path = get_config_path()
|
||||||
config = load_config()
|
config = load_config()
|
||||||
|
|||||||
@@ -0,0 +1,93 @@
|
|||||||
|
"""OAuth CLI commands for subscription authentication."""
|
||||||
|
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
import typer
|
||||||
|
from rich.console import Console
|
||||||
|
|
||||||
|
oauth_app = typer.Typer(help="Manage OAuth authentication for subscription-based providers")
|
||||||
|
console = Console()
|
||||||
|
|
||||||
|
|
||||||
|
@oauth_app.command("login")
|
||||||
|
def login(
|
||||||
|
provider: str = typer.Argument("anthropic", help="Provider name"),
|
||||||
|
token: Optional[str] = typer.Option(None, "--token", "-t", help="OAuth token (from claude setup-token)"),
|
||||||
|
):
|
||||||
|
"""Login to a provider using OAuth.
|
||||||
|
|
||||||
|
For Anthropic Claude Max/Pro, run 'claude setup-token' and paste the token here.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
nanobot oauth login anthropic --token sk-ant-oat01-xxx
|
||||||
|
"""
|
||||||
|
from nanobot.config.oauth_store import OAuthStore
|
||||||
|
from nanobot.config.schema import OAuthCredentials
|
||||||
|
|
||||||
|
if provider != "anthropic":
|
||||||
|
console.print(f"[red]OAuth login for {provider} not yet supported[/red]")
|
||||||
|
return
|
||||||
|
|
||||||
|
if not token:
|
||||||
|
console.print("Please provide your OAuth token:")
|
||||||
|
console.print(" 1. Run: claude setup-token")
|
||||||
|
console.print(" 2. Copy the sk-ant-oat01-... token")
|
||||||
|
console.print(" 3. Run: nanobot oauth login anthropic --token <your-token>")
|
||||||
|
console.print()
|
||||||
|
token = typer.prompt("Token", hide_input=True)
|
||||||
|
|
||||||
|
if not token or "sk-ant-oat" not in token:
|
||||||
|
console.print("[red]Invalid token. Must contain sk-ant-oat[/red]")
|
||||||
|
return
|
||||||
|
|
||||||
|
store = OAuthStore(Path.home() / ".nanobot")
|
||||||
|
creds = OAuthCredentials(
|
||||||
|
access_token=token,
|
||||||
|
token_type="token" # setup-token doesn't expire
|
||||||
|
)
|
||||||
|
store.save(provider, creds)
|
||||||
|
|
||||||
|
console.print(f"[green]Successfully saved {provider} OAuth credentials![/green]")
|
||||||
|
|
||||||
|
|
||||||
|
@oauth_app.command("status")
|
||||||
|
def status():
|
||||||
|
"""Show OAuth credential status."""
|
||||||
|
from nanobot.config.oauth_store import OAuthStore
|
||||||
|
|
||||||
|
store = OAuthStore(Path.home() / ".nanobot")
|
||||||
|
|
||||||
|
providers = ["anthropic"]
|
||||||
|
found_any = False
|
||||||
|
|
||||||
|
for provider in providers:
|
||||||
|
creds = store.load(provider)
|
||||||
|
if creds:
|
||||||
|
found_any = True
|
||||||
|
st = "valid"
|
||||||
|
if creds.is_expired:
|
||||||
|
st = "EXPIRED"
|
||||||
|
elif creds.expires_soon:
|
||||||
|
st = "expires soon"
|
||||||
|
|
||||||
|
token_preview = creds.access_token[:20] + "..."
|
||||||
|
console.print(f" {provider}: {token_preview} ({st})")
|
||||||
|
|
||||||
|
if not found_any:
|
||||||
|
console.print("No OAuth credentials configured.")
|
||||||
|
console.print("Run: nanobot oauth login anthropic --token <token>")
|
||||||
|
|
||||||
|
|
||||||
|
@oauth_app.command("logout")
|
||||||
|
def logout(
|
||||||
|
provider: str = typer.Argument("anthropic", help="Provider name"),
|
||||||
|
):
|
||||||
|
"""Remove OAuth credentials for a provider."""
|
||||||
|
from nanobot.config.oauth_store import OAuthStore
|
||||||
|
|
||||||
|
store = OAuthStore(Path.home() / ".nanobot")
|
||||||
|
if store.delete(provider):
|
||||||
|
console.print(f"[green]Removed {provider} OAuth credentials[/green]")
|
||||||
|
else:
|
||||||
|
console.print(f"No credentials found for {provider}")
|
||||||
@@ -12,12 +12,29 @@ def get_config_path() -> Path:
|
|||||||
return Path.home() / ".nanobot" / "config.json"
|
return Path.home() / ".nanobot" / "config.json"
|
||||||
|
|
||||||
|
|
||||||
|
def _get_oauth_store_dir() -> Path:
|
||||||
|
"""Get the OAuth store directory."""
|
||||||
|
return Path.home() / ".nanobot"
|
||||||
|
|
||||||
|
|
||||||
def get_data_dir() -> Path:
|
def get_data_dir() -> Path:
|
||||||
"""Get the nanobot data directory."""
|
"""Get the nanobot data directory."""
|
||||||
from nanobot.utils.helpers import get_data_path
|
from nanobot.utils.helpers import get_data_path
|
||||||
return get_data_path()
|
return get_data_path()
|
||||||
|
|
||||||
|
|
||||||
|
def _inject_oauth_credentials(config: Config) -> Config:
|
||||||
|
"""Inject OAuth credentials from store into config if available."""
|
||||||
|
from nanobot.config.oauth_store import OAuthStore
|
||||||
|
|
||||||
|
store = OAuthStore(_get_oauth_store_dir())
|
||||||
|
creds = store.load("anthropic")
|
||||||
|
if creds and creds.access_token and not creds.is_expired:
|
||||||
|
config.providers.anthropic.api_key = creds.access_token
|
||||||
|
|
||||||
|
return config
|
||||||
|
|
||||||
|
|
||||||
def load_config(config_path: Path | None = None) -> Config:
|
def load_config(config_path: Path | None = None) -> Config:
|
||||||
"""
|
"""
|
||||||
Load configuration from file or create default.
|
Load configuration from file or create default.
|
||||||
@@ -35,12 +52,13 @@ def load_config(config_path: Path | None = None) -> Config:
|
|||||||
with open(path) as f:
|
with open(path) as f:
|
||||||
data = json.load(f)
|
data = json.load(f)
|
||||||
data = _migrate_config(data)
|
data = _migrate_config(data)
|
||||||
return Config.model_validate(convert_keys(data))
|
config = Config.model_validate(convert_keys(data))
|
||||||
|
return _inject_oauth_credentials(config)
|
||||||
except (json.JSONDecodeError, ValueError) as e:
|
except (json.JSONDecodeError, ValueError) as e:
|
||||||
print(f"Warning: Failed to load config from {path}: {e}")
|
print(f"Warning: Failed to load config from {path}: {e}")
|
||||||
print("Using default configuration.")
|
print("Using default configuration.")
|
||||||
|
|
||||||
return Config()
|
return _inject_oauth_credentials(Config())
|
||||||
|
|
||||||
|
|
||||||
def save_config(config: Config, config_path: Path | None = None) -> None:
|
def save_config(config: Config, config_path: Path | None = None) -> None:
|
||||||
|
|||||||
@@ -0,0 +1,59 @@
|
|||||||
|
"""OAuth credential storage."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
from nanobot.config.schema import OAuthCredentials
|
||||||
|
|
||||||
|
|
||||||
|
class OAuthStore:
|
||||||
|
"""Stores OAuth credentials in a JSON file."""
|
||||||
|
|
||||||
|
FILENAME = "oauth-credentials.json"
|
||||||
|
|
||||||
|
def __init__(self, config_dir: Path):
|
||||||
|
self.config_dir = config_dir
|
||||||
|
self.file_path = config_dir / self.FILENAME
|
||||||
|
|
||||||
|
def _load_all(self) -> dict[str, Any]:
|
||||||
|
"""Load all credentials from file."""
|
||||||
|
if not self.file_path.exists():
|
||||||
|
return {}
|
||||||
|
|
||||||
|
with open(self.file_path, "r") as f:
|
||||||
|
return json.load(f)
|
||||||
|
|
||||||
|
def _save_all(self, data: dict[str, Any]) -> None:
|
||||||
|
"""Save all credentials to file."""
|
||||||
|
self.config_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
with open(self.file_path, "w") as f:
|
||||||
|
json.dump(data, f, indent=2)
|
||||||
|
|
||||||
|
# Secure permissions
|
||||||
|
self.file_path.chmod(0o600)
|
||||||
|
|
||||||
|
def save(self, provider: str, credentials: OAuthCredentials) -> None:
|
||||||
|
"""Save credentials for a provider."""
|
||||||
|
data = self._load_all()
|
||||||
|
data[provider] = credentials.model_dump()
|
||||||
|
self._save_all(data)
|
||||||
|
|
||||||
|
def load(self, provider: str) -> OAuthCredentials | None:
|
||||||
|
"""Load credentials for a provider."""
|
||||||
|
data = self._load_all()
|
||||||
|
if provider not in data:
|
||||||
|
return None
|
||||||
|
|
||||||
|
return OAuthCredentials(**data[provider])
|
||||||
|
|
||||||
|
def delete(self, provider: str) -> bool:
|
||||||
|
"""Delete credentials for a provider."""
|
||||||
|
data = self._load_all()
|
||||||
|
if provider not in data:
|
||||||
|
return False
|
||||||
|
|
||||||
|
del data[provider]
|
||||||
|
self._save_all(data)
|
||||||
|
return True
|
||||||
@@ -163,6 +163,7 @@ class AgentDefaults(BaseModel):
|
|||||||
temperature: float = 0.7
|
temperature: float = 0.7
|
||||||
max_tool_iterations: int = 20
|
max_tool_iterations: int = 20
|
||||||
memory_window: int = 50
|
memory_window: int = 50
|
||||||
|
thinking_budget: int = 0 # 0 = disabled; >0 = token budget for extended thinking
|
||||||
|
|
||||||
|
|
||||||
class AgentsConfig(BaseModel):
|
class AgentsConfig(BaseModel):
|
||||||
@@ -170,11 +171,41 @@ class AgentsConfig(BaseModel):
|
|||||||
defaults: AgentDefaults = Field(default_factory=AgentDefaults)
|
defaults: AgentDefaults = Field(default_factory=AgentDefaults)
|
||||||
|
|
||||||
|
|
||||||
|
class OAuthCredentials(BaseModel):
|
||||||
|
"""OAuth token credentials for subscription-based auth."""
|
||||||
|
access_token: str = ""
|
||||||
|
refresh_token: str = ""
|
||||||
|
expires_at: int = 0 # Unix timestamp
|
||||||
|
token_type: str = "oauth" # "oauth" or "token" (setup-token)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_oauth_token(self) -> bool:
|
||||||
|
"""Check if this is an OAuth token (vs regular API key)."""
|
||||||
|
return "sk-ant-oat" in self.access_token
|
||||||
|
|
||||||
|
@property
|
||||||
|
def is_expired(self) -> bool:
|
||||||
|
"""Check if token has expired."""
|
||||||
|
import time
|
||||||
|
if self.expires_at == 0:
|
||||||
|
return False # No expiry set (setup-token)
|
||||||
|
return time.time() > self.expires_at
|
||||||
|
|
||||||
|
@property
|
||||||
|
def expires_soon(self) -> bool:
|
||||||
|
"""Check if token expires within 10 minutes."""
|
||||||
|
import time
|
||||||
|
if self.expires_at == 0:
|
||||||
|
return False
|
||||||
|
return time.time() > (self.expires_at - 600)
|
||||||
|
|
||||||
|
|
||||||
class ProviderConfig(BaseModel):
|
class ProviderConfig(BaseModel):
|
||||||
"""LLM provider configuration."""
|
"""LLM provider configuration."""
|
||||||
api_key: str = ""
|
api_key: str = ""
|
||||||
api_base: str | None = None
|
api_base: str | None = None
|
||||||
extra_headers: dict[str, str] | None = None # Custom headers (e.g. APP-Code for AiHubMix)
|
extra_headers: dict[str, str] | None = None # Custom headers (e.g. APP-Code for AiHubMix)
|
||||||
|
oauth_credentials: OAuthCredentials | None = None
|
||||||
|
|
||||||
|
|
||||||
class ProvidersConfig(BaseModel):
|
class ProvidersConfig(BaseModel):
|
||||||
@@ -199,6 +230,26 @@ class GatewayConfig(BaseModel):
|
|||||||
port: int = 18790
|
port: int = 18790
|
||||||
|
|
||||||
|
|
||||||
|
class HooksConfig(BaseModel):
|
||||||
|
"""Webhook endpoint configuration."""
|
||||||
|
enabled: bool = False
|
||||||
|
tokens: dict[str, str] = Field(default_factory=dict) # Named tokens: {name: secret}
|
||||||
|
path: str = "/hooks" # URL path for the endpoint
|
||||||
|
timeout_seconds: int = 120 # Max time to wait for agent response
|
||||||
|
|
||||||
|
def resolve_token(self, provided: str) -> str | None:
|
||||||
|
"""Return token name if provided secret matches, else None."""
|
||||||
|
for name, secret in self.tokens.items():
|
||||||
|
if secret == provided:
|
||||||
|
return name
|
||||||
|
return None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def has_tokens(self) -> bool:
|
||||||
|
"""True if at least one token is configured."""
|
||||||
|
return bool(self.tokens)
|
||||||
|
|
||||||
|
|
||||||
class WebSearchConfig(BaseModel):
|
class WebSearchConfig(BaseModel):
|
||||||
"""Web search tool configuration."""
|
"""Web search tool configuration."""
|
||||||
api_key: str = "" # Brave Search API key
|
api_key: str = "" # Brave Search API key
|
||||||
@@ -228,6 +279,7 @@ class Config(BaseSettings):
|
|||||||
channels: ChannelsConfig = Field(default_factory=ChannelsConfig)
|
channels: ChannelsConfig = Field(default_factory=ChannelsConfig)
|
||||||
providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
|
providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
|
||||||
gateway: GatewayConfig = Field(default_factory=GatewayConfig)
|
gateway: GatewayConfig = Field(default_factory=GatewayConfig)
|
||||||
|
hooks: HooksConfig = Field(default_factory=HooksConfig)
|
||||||
tools: ToolsConfig = Field(default_factory=ToolsConfig)
|
tools: ToolsConfig = Field(default_factory=ToolsConfig)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
|
|||||||
@@ -1,11 +1,16 @@
|
|||||||
"""Heartbeat service - periodic agent wake-up to check for tasks."""
|
"""Heartbeat service - periodic agent wake-up to check for tasks."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Callable, Coroutine
|
from typing import TYPE_CHECKING, Any, Callable, Coroutine
|
||||||
|
|
||||||
from loguru import logger
|
from loguru import logger
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from nanobot.session.manager import SessionManager
|
||||||
|
|
||||||
# Default interval: 30 minutes
|
# Default interval: 30 minutes
|
||||||
DEFAULT_HEARTBEAT_INTERVAL_S = 30 * 60
|
DEFAULT_HEARTBEAT_INTERVAL_S = 30 * 60
|
||||||
|
|
||||||
@@ -46,14 +51,20 @@ class HeartbeatService:
|
|||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
workspace: Path,
|
workspace: Path,
|
||||||
on_heartbeat: Callable[[str], Coroutine[Any, Any, str]] | None = None,
|
on_heartbeat: Callable[[str, dict[str, Any] | None], Coroutine[Any, Any, str]] | None = None,
|
||||||
interval_s: int = DEFAULT_HEARTBEAT_INTERVAL_S,
|
interval_s: int = DEFAULT_HEARTBEAT_INTERVAL_S,
|
||||||
enabled: bool = True,
|
enabled: bool = True,
|
||||||
|
session_manager: SessionManager | None = None,
|
||||||
|
target_session_key: str = "telegram:239824268",
|
||||||
|
idle_threshold_s: int = 30 * 60, # 30 minutes
|
||||||
):
|
):
|
||||||
self.workspace = workspace
|
self.workspace = workspace
|
||||||
self.on_heartbeat = on_heartbeat
|
self.on_heartbeat = on_heartbeat
|
||||||
self.interval_s = interval_s
|
self.interval_s = interval_s
|
||||||
self.enabled = enabled
|
self.enabled = enabled
|
||||||
|
self.session_manager = session_manager
|
||||||
|
self.target_session_key = target_session_key
|
||||||
|
self.idle_threshold_s = idle_threshold_s
|
||||||
self._running = False
|
self._running = False
|
||||||
self._task: asyncio.Task | None = None
|
self._task: asyncio.Task | None = None
|
||||||
|
|
||||||
@@ -101,6 +112,32 @@ class HeartbeatService:
|
|||||||
|
|
||||||
async def _tick(self) -> None:
|
async def _tick(self) -> None:
|
||||||
"""Execute a single heartbeat tick."""
|
"""Execute a single heartbeat tick."""
|
||||||
|
|
||||||
|
# Check if user is idle (if session manager provided)
|
||||||
|
if self.session_manager and self.target_session_key:
|
||||||
|
try:
|
||||||
|
session = self.session_manager.get_or_create(self.target_session_key)
|
||||||
|
|
||||||
|
# Find last user message timestamp
|
||||||
|
last_user_timestamp = None
|
||||||
|
for msg in reversed(session.messages):
|
||||||
|
if msg.get("role") == "user":
|
||||||
|
last_user_timestamp = msg.get("timestamp")
|
||||||
|
break
|
||||||
|
|
||||||
|
if last_user_timestamp:
|
||||||
|
from datetime import datetime
|
||||||
|
last_dt = datetime.fromisoformat(last_user_timestamp)
|
||||||
|
elapsed = (datetime.now() - last_dt).total_seconds()
|
||||||
|
|
||||||
|
if elapsed < self.idle_threshold_s:
|
||||||
|
logger.debug(f"Heartbeat: user active {int(elapsed)}s ago, skipping")
|
||||||
|
return # User is active, don't trigger heartbeat
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Heartbeat: error checking idle state: {e}")
|
||||||
|
# Continue with heartbeat on error (fail open)
|
||||||
|
|
||||||
|
# Original heartbeat logic
|
||||||
content = self._read_heartbeat_file()
|
content = self._read_heartbeat_file()
|
||||||
|
|
||||||
# Skip if HEARTBEAT.md is empty or doesn't exist
|
# Skip if HEARTBEAT.md is empty or doesn't exist
|
||||||
@@ -108,17 +145,18 @@ class HeartbeatService:
|
|||||||
logger.debug("Heartbeat: no tasks (HEARTBEAT.md empty)")
|
logger.debug("Heartbeat: no tasks (HEARTBEAT.md empty)")
|
||||||
return
|
return
|
||||||
|
|
||||||
logger.info("Heartbeat: checking for tasks...")
|
logger.info("Heartbeat: user idle, checking for tasks...")
|
||||||
|
|
||||||
if self.on_heartbeat:
|
if self.on_heartbeat:
|
||||||
try:
|
try:
|
||||||
response = await self.on_heartbeat(HEARTBEAT_PROMPT)
|
# Call with suppress_output metadata
|
||||||
|
await self.on_heartbeat(
|
||||||
|
HEARTBEAT_PROMPT,
|
||||||
|
metadata={"suppress_output": True}
|
||||||
|
)
|
||||||
|
|
||||||
# Check if agent said "nothing to do"
|
# Note: HEARTBEAT_OK check removed - suppress mode makes it unnecessary
|
||||||
if HEARTBEAT_OK_TOKEN.replace("_", "") in response.upper().replace("_", ""):
|
logger.info("Heartbeat: completed")
|
||||||
logger.info("Heartbeat: OK (no action needed)")
|
|
||||||
else:
|
|
||||||
logger.info(f"Heartbeat: completed task")
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Heartbeat execution failed: {e}")
|
logger.error(f"Heartbeat execution failed: {e}")
|
||||||
@@ -126,5 +164,5 @@ class HeartbeatService:
|
|||||||
async def trigger_now(self) -> str | None:
|
async def trigger_now(self) -> str | None:
|
||||||
"""Manually trigger a heartbeat."""
|
"""Manually trigger a heartbeat."""
|
||||||
if self.on_heartbeat:
|
if self.on_heartbeat:
|
||||||
return await self.on_heartbeat(HEARTBEAT_PROMPT)
|
return await self.on_heartbeat(HEARTBEAT_PROMPT, metadata={"suppress_output": True})
|
||||||
return None
|
return None
|
||||||
|
|||||||
@@ -0,0 +1,139 @@
|
|||||||
|
"""HTTP hooks server for external service integration."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import uuid
|
||||||
|
|
||||||
|
from aiohttp import web
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from nanobot.bus.events import InboundMessage
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.config.schema import HooksConfig
|
||||||
|
|
||||||
|
|
||||||
|
class HooksServer:
|
||||||
|
"""
|
||||||
|
HTTP server exposing a /hooks endpoint.
|
||||||
|
|
||||||
|
External services POST JSON messages. The server publishes them
|
||||||
|
to the bus as InboundMessages and uses bus-level correlation
|
||||||
|
to return the agent's response synchronously.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
host: str,
|
||||||
|
port: int,
|
||||||
|
config: HooksConfig,
|
||||||
|
bus: MessageBus,
|
||||||
|
):
|
||||||
|
self.host = host
|
||||||
|
self.port = port
|
||||||
|
self.config = config
|
||||||
|
self.bus = bus
|
||||||
|
self._app = web.Application()
|
||||||
|
self._app.router.add_post(self.config.path, self._handle_hook)
|
||||||
|
self._app.router.add_get("/health", self._handle_health)
|
||||||
|
self._runner: web.AppRunner | None = None
|
||||||
|
|
||||||
|
async def start(self) -> None:
|
||||||
|
"""Start the HTTP server."""
|
||||||
|
if not self.config.has_tokens:
|
||||||
|
logger.warning("Hooks server has no tokens configured — endpoint disabled for security")
|
||||||
|
return
|
||||||
|
|
||||||
|
self._runner = web.AppRunner(self._app)
|
||||||
|
await self._runner.setup()
|
||||||
|
site = web.TCPSite(self._runner, self.host, self.port)
|
||||||
|
await site.start()
|
||||||
|
logger.info(f"Hooks server listening on {self.host}:{self.port}{self.config.path}")
|
||||||
|
|
||||||
|
async def stop(self) -> None:
|
||||||
|
"""Stop the HTTP server."""
|
||||||
|
if self._runner:
|
||||||
|
await self._runner.cleanup()
|
||||||
|
self._runner = None
|
||||||
|
|
||||||
|
def _resolve_auth(self, request: web.Request) -> str | None:
|
||||||
|
"""
|
||||||
|
Validate auth and return token name if valid, None otherwise.
|
||||||
|
Checks Authorization: Bearer <token> and X-Hook-Token headers.
|
||||||
|
"""
|
||||||
|
# Try Authorization: Bearer <token>
|
||||||
|
auth = request.headers.get("Authorization", "")
|
||||||
|
if auth.startswith("Bearer "):
|
||||||
|
token = auth[7:]
|
||||||
|
else:
|
||||||
|
# Try X-Hook-Token header
|
||||||
|
token = request.headers.get("X-Hook-Token", "")
|
||||||
|
|
||||||
|
return self.config.resolve_token(token) if token else None
|
||||||
|
|
||||||
|
async def _handle_health(self, request: web.Request) -> web.Response:
|
||||||
|
"""Health check endpoint — no auth required."""
|
||||||
|
return web.json_response({"status": "ok"})
|
||||||
|
|
||||||
|
async def _handle_hook(self, request: web.Request) -> web.Response:
|
||||||
|
"""Handle incoming hook request."""
|
||||||
|
# Auth check — resolve token name
|
||||||
|
token_name = self._resolve_auth(request)
|
||||||
|
if not token_name:
|
||||||
|
return web.json_response({"error": "unauthorized"}, status=401)
|
||||||
|
|
||||||
|
# Parse body
|
||||||
|
try:
|
||||||
|
body = await request.json()
|
||||||
|
except (json.JSONDecodeError, Exception):
|
||||||
|
return web.json_response({"error": "invalid JSON body"}, status=400)
|
||||||
|
|
||||||
|
# Validate required fields
|
||||||
|
message = body.get("message")
|
||||||
|
if not message or not isinstance(message, str):
|
||||||
|
return web.json_response(
|
||||||
|
{"error": "missing or invalid 'message' field"}, status=400
|
||||||
|
)
|
||||||
|
|
||||||
|
# Optional fields
|
||||||
|
channel = body.get("channel", "hook")
|
||||||
|
chat_id = body.get("chat_id", token_name)
|
||||||
|
timeout = body.get("timeout", self.config.timeout_seconds)
|
||||||
|
|
||||||
|
# Create correlation
|
||||||
|
correlation_id = str(uuid.uuid4())
|
||||||
|
|
||||||
|
# Build InboundMessage
|
||||||
|
msg = InboundMessage(
|
||||||
|
channel=channel,
|
||||||
|
sender_id=f"hook:{token_name}",
|
||||||
|
chat_id=str(chat_id),
|
||||||
|
content=message,
|
||||||
|
metadata={
|
||||||
|
"correlation_id": correlation_id,
|
||||||
|
"hook_source": token_name,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
# Fire-and-forget mode
|
||||||
|
if timeout == 0:
|
||||||
|
await self.bus.publish_inbound(msg)
|
||||||
|
return web.json_response({"ok": True}, status=202)
|
||||||
|
|
||||||
|
# Request-response mode
|
||||||
|
future = self.bus.register_correlation(correlation_id)
|
||||||
|
await self.bus.publish_inbound(msg)
|
||||||
|
|
||||||
|
try:
|
||||||
|
response = await asyncio.wait_for(future, timeout=timeout)
|
||||||
|
return web.json_response({"ok": True, "response": response})
|
||||||
|
except asyncio.TimeoutError:
|
||||||
|
return web.json_response(
|
||||||
|
{"ok": False, "error": f"agent did not respond within {timeout}s"},
|
||||||
|
status=504,
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Hook processing error: {e}")
|
||||||
|
return web.json_response({"error": "internal error"}, status=500)
|
||||||
|
finally:
|
||||||
|
# Clean up correlation on any failure
|
||||||
|
self.bus.cancel_correlation(correlation_id)
|
||||||
@@ -1,6 +1,45 @@
|
|||||||
"""LLM provider abstraction module."""
|
"""Provider module exports."""
|
||||||
|
|
||||||
from nanobot.providers.base import LLMProvider, LLMResponse
|
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
|
||||||
from nanobot.providers.litellm_provider import LiteLLMProvider
|
from nanobot.providers.litellm_provider import LiteLLMProvider
|
||||||
|
from nanobot.providers.anthropic_oauth import AnthropicOAuthProvider
|
||||||
|
from nanobot.providers.registry import should_use_oauth_provider
|
||||||
|
|
||||||
__all__ = ["LLMProvider", "LLMResponse", "LiteLLMProvider"]
|
__all__ = [
|
||||||
|
"LLMProvider",
|
||||||
|
"LLMResponse",
|
||||||
|
"ToolCallRequest",
|
||||||
|
"LiteLLMProvider",
|
||||||
|
"AnthropicOAuthProvider",
|
||||||
|
"create_provider",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def create_provider(
|
||||||
|
api_key: str,
|
||||||
|
model: str,
|
||||||
|
api_base: str | None = None,
|
||||||
|
extra_headers: dict[str, str] | None = None,
|
||||||
|
provider_name: str | None = None,
|
||||||
|
thinking_budget: int = 0,
|
||||||
|
) -> LLMProvider:
|
||||||
|
"""Factory function to create appropriate provider.
|
||||||
|
|
||||||
|
Automatically selects AnthropicOAuthProvider for OAuth tokens,
|
||||||
|
LiteLLMProvider for everything else.
|
||||||
|
"""
|
||||||
|
if should_use_oauth_provider(api_key, model):
|
||||||
|
return AnthropicOAuthProvider(
|
||||||
|
oauth_token=api_key,
|
||||||
|
default_model=model,
|
||||||
|
api_base=api_base,
|
||||||
|
thinking_budget=thinking_budget,
|
||||||
|
)
|
||||||
|
|
||||||
|
return LiteLLMProvider(
|
||||||
|
api_key=api_key,
|
||||||
|
api_base=api_base,
|
||||||
|
default_model=model,
|
||||||
|
extra_headers=extra_headers,
|
||||||
|
provider_name=provider_name,
|
||||||
|
)
|
||||||
|
|||||||
@@ -0,0 +1,453 @@
|
|||||||
|
"""Anthropic OAuth provider - direct API calls with Bearer auth.
|
||||||
|
|
||||||
|
This provider bypasses litellm to properly handle OAuth tokens
|
||||||
|
which require Authorization: Bearer header instead of x-api-key.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
|
||||||
|
from nanobot.providers.oauth_utils import get_auth_headers
|
||||||
|
|
||||||
|
|
||||||
|
class AnthropicOAuthProvider(LLMProvider):
|
||||||
|
"""
|
||||||
|
Anthropic provider using OAuth token authentication.
|
||||||
|
|
||||||
|
Unlike the LiteLLM provider, this calls the Anthropic API directly
|
||||||
|
with proper Bearer token authentication for Claude Max/Pro subscriptions.
|
||||||
|
"""
|
||||||
|
|
||||||
|
ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
oauth_token: str,
|
||||||
|
default_model: str = "claude-opus-4-5",
|
||||||
|
api_base: str | None = None,
|
||||||
|
thinking_budget: int = 0,
|
||||||
|
):
|
||||||
|
super().__init__(api_key=None, api_base=api_base)
|
||||||
|
self.oauth_token = oauth_token
|
||||||
|
self.default_model = default_model
|
||||||
|
self.thinking_budget = thinking_budget
|
||||||
|
self._client: httpx.AsyncClient | None = None
|
||||||
|
|
||||||
|
def _get_headers(self) -> dict[str, str]:
|
||||||
|
"""Get request headers with Bearer auth."""
|
||||||
|
return get_auth_headers(self.oauth_token, is_oauth=True)
|
||||||
|
|
||||||
|
def _get_api_url(self) -> str:
|
||||||
|
"""Get API endpoint URL."""
|
||||||
|
if self.api_base:
|
||||||
|
return f"{self.api_base.rstrip('/')}/v1/messages"
|
||||||
|
return self.ANTHROPIC_API_URL
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _normalize_model(model: str) -> str:
|
||||||
|
"""Normalize model name for the Anthropic API.
|
||||||
|
|
||||||
|
Anthropic model IDs use hyphens (claude-sonnet-4-5), but users often
|
||||||
|
write dots (claude-sonnet-4.5). Normalize so both work.
|
||||||
|
"""
|
||||||
|
return model.replace(".", "-")
|
||||||
|
|
||||||
|
async def _get_client(self) -> httpx.AsyncClient:
|
||||||
|
"""Get or create async HTTP client."""
|
||||||
|
if self._client is None:
|
||||||
|
self._client = httpx.AsyncClient(timeout=300.0)
|
||||||
|
return self._client
|
||||||
|
|
||||||
|
def _prepare_messages(
|
||||||
|
self,
|
||||||
|
messages: list[dict[str, Any]]
|
||||||
|
) -> tuple[str | None, list[dict[str, Any]]]:
|
||||||
|
"""Prepare messages: extract system prompt and convert OpenAI format to Anthropic.
|
||||||
|
|
||||||
|
The agent loop produces messages in OpenAI format:
|
||||||
|
- assistant msgs with tool_calls [{type:"function", function:{name, arguments}}]
|
||||||
|
- tool role msgs with tool_call_id, name, content
|
||||||
|
|
||||||
|
Anthropic API expects:
|
||||||
|
- assistant msgs with content blocks [{type:"tool_use", id, name, input}]
|
||||||
|
- user msgs with content blocks [{type:"tool_result", tool_use_id, content}]
|
||||||
|
|
||||||
|
Returns (system_prompt, anthropic_messages)
|
||||||
|
"""
|
||||||
|
system_parts = []
|
||||||
|
converted: list[dict[str, Any]] = []
|
||||||
|
|
||||||
|
for msg in messages:
|
||||||
|
role = msg.get("role")
|
||||||
|
|
||||||
|
if role == "system":
|
||||||
|
system_parts.append(msg.get("content", ""))
|
||||||
|
continue
|
||||||
|
|
||||||
|
if role == "assistant" and msg.get("tool_calls"):
|
||||||
|
# Convert OpenAI tool_calls to Anthropic content blocks
|
||||||
|
content_blocks: list[dict[str, Any]] = []
|
||||||
|
# Preserve thinking blocks (list=raw API blocks with signatures, str=legacy)
|
||||||
|
rc = msg.get("reasoning_content")
|
||||||
|
if isinstance(rc, list):
|
||||||
|
content_blocks.extend(rc)
|
||||||
|
elif isinstance(rc, str) and rc:
|
||||||
|
content_blocks.append({"type": "thinking", "thinking": rc})
|
||||||
|
text = msg.get("content")
|
||||||
|
if text:
|
||||||
|
content_blocks.append({"type": "text", "text": text})
|
||||||
|
for tc in msg["tool_calls"]:
|
||||||
|
func = tc.get("function", {})
|
||||||
|
args = func.get("arguments", "{}")
|
||||||
|
if isinstance(args, str):
|
||||||
|
try:
|
||||||
|
args = json.loads(args)
|
||||||
|
except (json.JSONDecodeError, TypeError):
|
||||||
|
args = {}
|
||||||
|
content_blocks.append({
|
||||||
|
"type": "tool_use",
|
||||||
|
"id": tc.get("id", ""),
|
||||||
|
"name": func.get("name", ""),
|
||||||
|
"input": args,
|
||||||
|
})
|
||||||
|
converted.append({"role": "assistant", "content": content_blocks})
|
||||||
|
continue
|
||||||
|
|
||||||
|
if role == "assistant" and msg.get("reasoning_content"):
|
||||||
|
# Plain assistant message with thinking (no tool calls)
|
||||||
|
rc = msg["reasoning_content"]
|
||||||
|
if isinstance(rc, list):
|
||||||
|
content_blocks = list(rc)
|
||||||
|
else:
|
||||||
|
content_blocks = [{"type": "thinking", "thinking": rc}]
|
||||||
|
text = msg.get("content")
|
||||||
|
if text:
|
||||||
|
content_blocks.append({"type": "text", "text": text})
|
||||||
|
converted.append({"role": "assistant", "content": content_blocks})
|
||||||
|
continue
|
||||||
|
|
||||||
|
if role == "tool":
|
||||||
|
# Convert tool result to Anthropic user message with tool_result block
|
||||||
|
tool_result_block = {
|
||||||
|
"type": "tool_result",
|
||||||
|
"tool_use_id": msg.get("tool_call_id", ""),
|
||||||
|
"content": msg.get("content", ""),
|
||||||
|
}
|
||||||
|
# Merge into previous user message if it already has tool_result blocks
|
||||||
|
if converted and converted[-1].get("role") == "user":
|
||||||
|
prev_content = converted[-1].get("content")
|
||||||
|
if isinstance(prev_content, list):
|
||||||
|
prev_content.append(tool_result_block)
|
||||||
|
continue
|
||||||
|
converted.append({"role": "user", "content": [tool_result_block]})
|
||||||
|
continue
|
||||||
|
|
||||||
|
if role == "user":
|
||||||
|
content = msg.get("content", "")
|
||||||
|
# Convert OpenAI image_url blocks to Anthropic image blocks
|
||||||
|
if isinstance(content, list):
|
||||||
|
content = self._convert_image_blocks(content)
|
||||||
|
# Merge text into previous user message if it has tool_result blocks
|
||||||
|
# (handles the "Reflect on the results" interleaved message)
|
||||||
|
if converted and converted[-1].get("role") == "user":
|
||||||
|
prev_content = converted[-1].get("content")
|
||||||
|
if isinstance(prev_content, list):
|
||||||
|
if isinstance(content, str):
|
||||||
|
prev_content.append({"type": "text", "text": content})
|
||||||
|
elif isinstance(content, list):
|
||||||
|
prev_content.extend(content)
|
||||||
|
continue
|
||||||
|
converted.append({"role": role, "content": content})
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Pass through other messages (assistant without tool_calls, etc.)
|
||||||
|
converted.append(msg)
|
||||||
|
|
||||||
|
system_prompt = "\n\n".join(system_parts)
|
||||||
|
return system_prompt, converted
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _convert_image_blocks(content: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||||
|
"""Convert OpenAI image_url blocks to Anthropic image blocks.
|
||||||
|
|
||||||
|
OpenAI format: {"type": "image_url", "image_url": {"url": "data:mime;base64,DATA"}}
|
||||||
|
Anthropic format: {"type": "image", "source": {"type": "base64", "media_type": "mime", "data": "DATA"}}
|
||||||
|
"""
|
||||||
|
converted = []
|
||||||
|
for block in content:
|
||||||
|
if block.get("type") == "image_url":
|
||||||
|
url = block.get("image_url", {}).get("url", "")
|
||||||
|
if url.startswith("data:") and ";base64," in url:
|
||||||
|
header, data = url.split(";base64,", 1)
|
||||||
|
media_type = header.removeprefix("data:")
|
||||||
|
converted.append({
|
||||||
|
"type": "image",
|
||||||
|
"source": {"type": "base64", "media_type": media_type, "data": data},
|
||||||
|
})
|
||||||
|
else:
|
||||||
|
converted.append({
|
||||||
|
"type": "image",
|
||||||
|
"source": {"type": "url", "url": url},
|
||||||
|
})
|
||||||
|
else:
|
||||||
|
converted.append(block)
|
||||||
|
return converted
|
||||||
|
|
||||||
|
def _convert_tools_to_anthropic(
|
||||||
|
self,
|
||||||
|
tools: list[dict[str, Any]] | None
|
||||||
|
) -> list[dict[str, Any]] | None:
|
||||||
|
"""Convert OpenAI-format tools to Anthropic format."""
|
||||||
|
if not tools:
|
||||||
|
return None
|
||||||
|
|
||||||
|
anthropic_tools = []
|
||||||
|
for tool in tools:
|
||||||
|
if tool.get("type") == "function":
|
||||||
|
func = tool["function"]
|
||||||
|
anthropic_tools.append({
|
||||||
|
"name": func["name"],
|
||||||
|
"description": func.get("description", ""),
|
||||||
|
"input_schema": func.get("parameters", {"type": "object", "properties": {}})
|
||||||
|
})
|
||||||
|
|
||||||
|
return anthropic_tools if anthropic_tools else None
|
||||||
|
|
||||||
|
async def _make_request(
|
||||||
|
self,
|
||||||
|
messages: list[dict[str, Any]],
|
||||||
|
system: str | None = None,
|
||||||
|
model: str = "claude-opus-4-5",
|
||||||
|
max_tokens: int = 4096,
|
||||||
|
temperature: float = 0.7,
|
||||||
|
tools: list[dict[str, Any]] | None = None,
|
||||||
|
thinking_budget_override: int | None = None,
|
||||||
|
context_management: dict[str, Any] | None = None,
|
||||||
|
) -> dict[str, Any]:
|
||||||
|
"""Make request to Anthropic API."""
|
||||||
|
client = await self._get_client()
|
||||||
|
|
||||||
|
# Cache the last user message so conversation history is cached across turns
|
||||||
|
if messages:
|
||||||
|
last = messages[-1]
|
||||||
|
if last.get("role") == "user":
|
||||||
|
content = last["content"]
|
||||||
|
if isinstance(content, str):
|
||||||
|
last = {**last, "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"}}
|
||||||
|
last = {**last, "content": new_content}
|
||||||
|
messages = messages[:-1] + [last]
|
||||||
|
|
||||||
|
payload: dict[str, Any] = {
|
||||||
|
"model": model,
|
||||||
|
"messages": messages,
|
||||||
|
"max_tokens": max_tokens,
|
||||||
|
}
|
||||||
|
|
||||||
|
# Extended thinking: temperature must be 1 when enabled
|
||||||
|
effective_thinking = thinking_budget_override if thinking_budget_override is not None else self.thinking_budget
|
||||||
|
if effective_thinking > 0:
|
||||||
|
payload["temperature"] = 1
|
||||||
|
# max_tokens must exceed budget_tokens
|
||||||
|
if max_tokens <= effective_thinking:
|
||||||
|
payload["max_tokens"] = effective_thinking + 4096
|
||||||
|
payload["thinking"] = {
|
||||||
|
"type": "enabled",
|
||||||
|
"budget_tokens": effective_thinking,
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
payload["temperature"] = temperature
|
||||||
|
|
||||||
|
if system:
|
||||||
|
payload["system"] = [{"type": "text", "text": system, "cache_control": {"type": "ephemeral", "ttl": "1h"}}]
|
||||||
|
|
||||||
|
if tools:
|
||||||
|
cached_tools = list(tools)
|
||||||
|
cached_tools[-1] = {**cached_tools[-1], "cache_control": {"type": "ephemeral", "ttl": "1h"}}
|
||||||
|
payload["tools"] = cached_tools
|
||||||
|
|
||||||
|
if context_management:
|
||||||
|
payload["context_management"] = context_management
|
||||||
|
|
||||||
|
edit_types = [e.get("type") for e in (context_management or {}).get("edits", [])]
|
||||||
|
logger.info(
|
||||||
|
"Anthropic request: model={} max_tokens={} thinking={} tools={} context_mgmt={}",
|
||||||
|
payload.get("model"), payload.get("max_tokens"),
|
||||||
|
payload.get("thinking", "disabled"),
|
||||||
|
len(payload.get("tools", [])),
|
||||||
|
edit_types or "none",
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await client.post(
|
||||||
|
self._get_api_url(),
|
||||||
|
headers=self._get_headers(),
|
||||||
|
json=payload,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Dump rate limit headers for analysis
|
||||||
|
try:
|
||||||
|
import datetime
|
||||||
|
import os
|
||||||
|
header_dump = {
|
||||||
|
"timestamp": datetime.datetime.utcnow().isoformat(),
|
||||||
|
"status_code": response.status_code,
|
||||||
|
"model": payload.get("model"),
|
||||||
|
"headers": dict(response.headers),
|
||||||
|
}
|
||||||
|
dump_path = "/root/.nanobot/workspace/api_headers.jsonl"
|
||||||
|
with open(dump_path, "a") as f:
|
||||||
|
f.write(json.dumps(header_dump) + "\n")
|
||||||
|
# Capture rate limit state for quota-based model switching
|
||||||
|
hdrs = response.headers
|
||||||
|
rate_limit_state = {
|
||||||
|
"updated_at": datetime.datetime.utcnow().isoformat(),
|
||||||
|
"model": payload.get("model"),
|
||||||
|
"weekly_all_models": float(hdrs["anthropic-ratelimit-unified-7d-utilization"]) if hdrs.get("anthropic-ratelimit-unified-7d-utilization") else None,
|
||||||
|
"weekly_sonnet": float(hdrs["anthropic-ratelimit-unified-7d_sonnet-utilization"]) if hdrs.get("anthropic-ratelimit-unified-7d_sonnet-utilization") else None,
|
||||||
|
"session_5h": float(hdrs["anthropic-ratelimit-unified-5h-utilization"]) if hdrs.get("anthropic-ratelimit-unified-5h-utilization") else None,
|
||||||
|
"weekly_reset": int(hdrs["anthropic-ratelimit-unified-7d-reset"]) if hdrs.get("anthropic-ratelimit-unified-7d-reset") else None,
|
||||||
|
"session_reset": int(hdrs["anthropic-ratelimit-unified-5h-reset"]) if hdrs.get("anthropic-ratelimit-unified-5h-reset") else None,
|
||||||
|
"binding_limit": hdrs.get("anthropic-ratelimit-unified-representative-claim"),
|
||||||
|
"sonnet_fallback": hdrs.get("anthropic-ratelimit-unified-fallback"),
|
||||||
|
}
|
||||||
|
state_path = "/root/.nanobot/workspace/memory/rate_limits.json"
|
||||||
|
os.makedirs(os.path.dirname(state_path), exist_ok=True)
|
||||||
|
with open(state_path, "w") as f:
|
||||||
|
json.dump(rate_limit_state, f, indent=2)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning("Rate limit header capture failed: {}", e)
|
||||||
|
|
||||||
|
if response.status_code != 200:
|
||||||
|
error_text = response.text
|
||||||
|
raise Exception(f"Anthropic API error {response.status_code}: {error_text}")
|
||||||
|
|
||||||
|
return response.json()
|
||||||
|
|
||||||
|
async def chat(
|
||||||
|
self,
|
||||||
|
messages: list[dict[str, Any]],
|
||||||
|
tools: list[dict[str, Any]] | None = None,
|
||||||
|
model: str | None = None,
|
||||||
|
max_tokens: int = 4096,
|
||||||
|
temperature: float = 0.7,
|
||||||
|
thinking_budget: int | None = None,
|
||||||
|
context_management: dict[str, Any] | None = None,
|
||||||
|
) -> LLMResponse:
|
||||||
|
"""Send chat completion request to Anthropic API."""
|
||||||
|
model = model or self.default_model
|
||||||
|
|
||||||
|
# Strip provider prefix if present (e.g. "anthropic/claude-opus-4-5" -> "claude-opus-4-5")
|
||||||
|
if "/" in model:
|
||||||
|
model = model.split("/")[-1]
|
||||||
|
|
||||||
|
# Normalize dots to hyphens (claude-sonnet-4.5 -> claude-sonnet-4-5)
|
||||||
|
model = self._normalize_model(model)
|
||||||
|
|
||||||
|
system, prepared_messages = self._prepare_messages(messages)
|
||||||
|
anthropic_tools = self._convert_tools_to_anthropic(tools)
|
||||||
|
|
||||||
|
# Per-call thinking override (None = use instance default)
|
||||||
|
effective_thinking = self.thinking_budget if thinking_budget is None else thinking_budget
|
||||||
|
|
||||||
|
try:
|
||||||
|
response = await self._make_request(
|
||||||
|
messages=prepared_messages,
|
||||||
|
system=system,
|
||||||
|
model=model,
|
||||||
|
max_tokens=max_tokens,
|
||||||
|
temperature=temperature,
|
||||||
|
tools=anthropic_tools,
|
||||||
|
thinking_budget_override=effective_thinking,
|
||||||
|
context_management=context_management,
|
||||||
|
)
|
||||||
|
return self._parse_response(response)
|
||||||
|
except Exception as e:
|
||||||
|
return LLMResponse(
|
||||||
|
content=f"Error calling LLM: {str(e)}",
|
||||||
|
finish_reason="error",
|
||||||
|
)
|
||||||
|
|
||||||
|
def _parse_response(self, response: dict[str, Any]) -> LLMResponse:
|
||||||
|
"""Parse Anthropic API response."""
|
||||||
|
content_blocks = response.get("content", [])
|
||||||
|
|
||||||
|
text_content = ""
|
||||||
|
thinking_blocks: list[dict[str, Any]] = []
|
||||||
|
tool_calls = []
|
||||||
|
|
||||||
|
for block in content_blocks:
|
||||||
|
if block.get("type") == "thinking":
|
||||||
|
# Preserve full block including signature for multi-turn replay
|
||||||
|
thinking_blocks.append(block)
|
||||||
|
elif block.get("type") == "text":
|
||||||
|
text_content += block.get("text", "")
|
||||||
|
elif block.get("type") == "tool_use":
|
||||||
|
tool_calls.append(ToolCallRequest(
|
||||||
|
id=block.get("id", ""),
|
||||||
|
name=block.get("name", ""),
|
||||||
|
arguments=block.get("input", {}),
|
||||||
|
))
|
||||||
|
|
||||||
|
usage = {}
|
||||||
|
if "usage" in response:
|
||||||
|
usage = {
|
||||||
|
"prompt_tokens": response["usage"].get("input_tokens", 0),
|
||||||
|
"completion_tokens": response["usage"].get("output_tokens", 0),
|
||||||
|
"total_tokens": (
|
||||||
|
response["usage"].get("input_tokens", 0) +
|
||||||
|
response["usage"].get("output_tokens", 0)
|
||||||
|
),
|
||||||
|
}
|
||||||
|
|
||||||
|
stop_reason = response.get("stop_reason", "end_turn")
|
||||||
|
thinking_chars = sum(len(b.get("thinking", "")) for b in thinking_blocks) if thinking_blocks else 0
|
||||||
|
raw_usage = response.get("usage", {})
|
||||||
|
cache_write = raw_usage.get("cache_creation_input_tokens", 0)
|
||||||
|
cache_read = raw_usage.get("cache_read_input_tokens", 0)
|
||||||
|
logger.info(
|
||||||
|
"Anthropic response: stop={} tool_calls={} thinking={} chars, "
|
||||||
|
"input={} output={} cache_write={} cache_read={} tokens",
|
||||||
|
stop_reason, len(tool_calls), thinking_chars,
|
||||||
|
usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0),
|
||||||
|
cache_write, cache_read,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log context editing activity if any edits were applied
|
||||||
|
if applied_edits := response.get("context_management", {}).get("applied_edits"):
|
||||||
|
for edit in applied_edits:
|
||||||
|
edit_type = edit.get("type", "?")
|
||||||
|
cleared_tokens = edit.get("cleared_input_tokens", 0)
|
||||||
|
if edit_type == "clear_tool_uses_20250919":
|
||||||
|
logger.info(
|
||||||
|
"Context edit: cleared {} tool uses ({} tokens)",
|
||||||
|
edit.get("cleared_tool_uses", 0), cleared_tokens,
|
||||||
|
)
|
||||||
|
elif edit_type == "clear_thinking_20251015":
|
||||||
|
logger.info(
|
||||||
|
"Context edit: cleared {} thinking turns ({} tokens)",
|
||||||
|
edit.get("cleared_thinking_turns", 0), cleared_tokens,
|
||||||
|
)
|
||||||
|
|
||||||
|
return LLMResponse(
|
||||||
|
content=text_content or None,
|
||||||
|
tool_calls=tool_calls,
|
||||||
|
finish_reason=stop_reason,
|
||||||
|
usage=usage,
|
||||||
|
reasoning_content=thinking_blocks or None,
|
||||||
|
)
|
||||||
|
|
||||||
|
def get_default_model(self) -> str:
|
||||||
|
"""Get the default model."""
|
||||||
|
return self.default_model
|
||||||
|
|
||||||
|
async def close(self):
|
||||||
|
"""Close the HTTP client."""
|
||||||
|
if self._client:
|
||||||
|
await self._client.aclose()
|
||||||
|
self._client = None
|
||||||
@@ -20,7 +20,7 @@ class LLMResponse:
|
|||||||
tool_calls: list[ToolCallRequest] = field(default_factory=list)
|
tool_calls: list[ToolCallRequest] = field(default_factory=list)
|
||||||
finish_reason: str = "stop"
|
finish_reason: str = "stop"
|
||||||
usage: dict[str, int] = field(default_factory=dict)
|
usage: dict[str, int] = field(default_factory=dict)
|
||||||
reasoning_content: str | None = None # Kimi, DeepSeek-R1 etc.
|
reasoning_content: Any = None # str for Kimi/DeepSeek-R1; list[dict] for Anthropic thinking blocks
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def has_tool_calls(self) -> bool:
|
def has_tool_calls(self) -> bool:
|
||||||
@@ -48,6 +48,8 @@ class LLMProvider(ABC):
|
|||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
max_tokens: int = 4096,
|
max_tokens: int = 4096,
|
||||||
temperature: float = 0.7,
|
temperature: float = 0.7,
|
||||||
|
thinking_budget: int | None = None,
|
||||||
|
context_management: dict[str, Any] | None = None,
|
||||||
) -> LLMResponse:
|
) -> LLMResponse:
|
||||||
"""
|
"""
|
||||||
Send a chat completion request.
|
Send a chat completion request.
|
||||||
|
|||||||
@@ -106,6 +106,8 @@ class LiteLLMProvider(LLMProvider):
|
|||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
max_tokens: int = 4096,
|
max_tokens: int = 4096,
|
||||||
temperature: float = 0.7,
|
temperature: float = 0.7,
|
||||||
|
thinking_budget: int | None = None,
|
||||||
|
context_management: dict[str, Any] | None = None, # Anthropic-only, ignored here
|
||||||
) -> LLMResponse:
|
) -> LLMResponse:
|
||||||
"""
|
"""
|
||||||
Send a chat completion request via LiteLLM.
|
Send a chat completion request via LiteLLM.
|
||||||
|
|||||||
@@ -0,0 +1,38 @@
|
|||||||
|
"""OAuth utility functions for Anthropic subscription auth."""
|
||||||
|
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
|
||||||
|
def is_oauth_token(token: str | None) -> bool:
|
||||||
|
"""Check if token is an OAuth token (vs regular API key).
|
||||||
|
|
||||||
|
OAuth tokens from Claude Max/Pro contain 'sk-ant-oat' prefix.
|
||||||
|
Regular API keys use 'sk-ant-api03' or similar.
|
||||||
|
"""
|
||||||
|
if not token:
|
||||||
|
return False
|
||||||
|
return "sk-ant-oat" in token
|
||||||
|
|
||||||
|
|
||||||
|
def get_auth_headers(token: str, is_oauth: bool = False) -> dict[str, str]:
|
||||||
|
"""Get authentication headers for Anthropic API.
|
||||||
|
|
||||||
|
OAuth tokens require Authorization: Bearer header.
|
||||||
|
Regular API keys use x-api-key header.
|
||||||
|
"""
|
||||||
|
headers: dict[str, str] = {
|
||||||
|
"anthropic-version": "2023-06-01",
|
||||||
|
"content-type": "application/json",
|
||||||
|
}
|
||||||
|
|
||||||
|
if is_oauth:
|
||||||
|
headers["Authorization"] = f"Bearer {token}"
|
||||||
|
# Required headers to mimic Claude Code client
|
||||||
|
headers["anthropic-beta"] = "claude-code-20250219,oauth-2025-04-20,context-management-2025-06-27"
|
||||||
|
headers["anthropic-dangerous-direct-browser-access"] = "true"
|
||||||
|
headers["user-agent"] = "claude-cli/2.1.2 (external, cli)"
|
||||||
|
headers["x-app"] = "cli"
|
||||||
|
else:
|
||||||
|
headers["x-api-key"] = token
|
||||||
|
|
||||||
|
return headers
|
||||||
@@ -357,3 +357,24 @@ def find_by_name(name: str) -> ProviderSpec | None:
|
|||||||
if spec.name == name:
|
if spec.name == name:
|
||||||
return spec
|
return spec
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def should_use_oauth_provider(api_key: str | None, model: str) -> bool:
|
||||||
|
"""Determine if OAuth provider should be used.
|
||||||
|
|
||||||
|
OAuth provider is used when:
|
||||||
|
1. API key is an OAuth token (contains 'sk-ant-oat')
|
||||||
|
2. Model is an Anthropic model (contains 'claude' or 'anthropic')
|
||||||
|
"""
|
||||||
|
if not api_key:
|
||||||
|
return False
|
||||||
|
|
||||||
|
if "sk-ant-oat" not in api_key:
|
||||||
|
return False
|
||||||
|
|
||||||
|
model_lower = model.lower()
|
||||||
|
anthropic_spec = find_by_name("anthropic")
|
||||||
|
if anthropic_spec:
|
||||||
|
return any(kw in model_lower for kw in anthropic_spec.keywords)
|
||||||
|
|
||||||
|
return False
|
||||||
|
|||||||
+23
-11
@@ -36,21 +36,33 @@ class Session:
|
|||||||
self.messages.append(msg)
|
self.messages.append(msg)
|
||||||
self.updated_at = datetime.now()
|
self.updated_at = datetime.now()
|
||||||
|
|
||||||
def get_history(self, max_messages: int = 50) -> list[dict[str, Any]]:
|
def add_raw_message(self, msg: dict[str, Any]) -> None:
|
||||||
"""
|
"""Add a pre-formed message dict to the session, preserving all fields."""
|
||||||
Get message history for LLM context.
|
stored = dict(msg)
|
||||||
|
if "timestamp" not in stored:
|
||||||
|
stored["timestamp"] = datetime.now().isoformat()
|
||||||
|
self.messages.append(stored)
|
||||||
|
self.updated_at = datetime.now()
|
||||||
|
|
||||||
Args:
|
# Fields that are valid in the Anthropic/OpenAI messages API.
|
||||||
max_messages: Maximum messages to return.
|
# Everything else (timestamp, tools_used, etc.) is internal metadata.
|
||||||
|
_API_FIELDS = {"role", "content", "tool_calls", "tool_call_id", "name", "reasoning_content"}
|
||||||
|
|
||||||
|
def get_history(self) -> list[dict[str, Any]]:
|
||||||
|
"""
|
||||||
|
Get full message history for LLM context.
|
||||||
|
|
||||||
|
The server-side context editing API (clear_tool_uses_20250919) handles
|
||||||
|
trimming old tool chains safely at token thresholds, so we send the full
|
||||||
|
history and let the server decide what to drop.
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
List of messages in LLM format.
|
List of messages in LLM format (API-relevant fields only).
|
||||||
"""
|
"""
|
||||||
# Get recent messages
|
return [
|
||||||
recent = self.messages[-max_messages:] if len(self.messages) > max_messages else self.messages
|
{k: v for k, v in m.items() if k in self._API_FIELDS and v is not None}
|
||||||
|
for m in self.messages
|
||||||
# Convert to LLM format (just role and content)
|
]
|
||||||
return [{"role": m["role"], "content": m["content"]} for m in recent]
|
|
||||||
|
|
||||||
def clear(self) -> None:
|
def clear(self) -> None:
|
||||||
"""Clear all messages in the session."""
|
"""Clear all messages in the session."""
|
||||||
|
|||||||
@@ -0,0 +1,97 @@
|
|||||||
|
# tests/test_agent_loop_metadata.py
|
||||||
|
import pytest
|
||||||
|
from pathlib import Path
|
||||||
|
from nanobot.agent.loop import AgentLoop
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.providers.base import LLMProvider, LLMResponse
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_process_direct_passes_metadata():
|
||||||
|
"""Test that process_direct passes metadata to InboundMessage."""
|
||||||
|
bus = MessageBus()
|
||||||
|
provider = MagicMock(spec=LLMProvider)
|
||||||
|
provider.chat = AsyncMock(return_value=LLMResponse(
|
||||||
|
content="test response",
|
||||||
|
tool_calls=[]
|
||||||
|
))
|
||||||
|
provider.get_default_model = MagicMock(return_value="test-model")
|
||||||
|
provider.thinking_budget = 0
|
||||||
|
|
||||||
|
workspace = Path("/tmp/test-workspace")
|
||||||
|
workspace.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
|
||||||
|
|
||||||
|
# Call with metadata
|
||||||
|
test_metadata = {"suppress_output": True, "test_key": "test_value"}
|
||||||
|
await loop.process_direct(
|
||||||
|
content="test message",
|
||||||
|
metadata=test_metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify provider.chat was called
|
||||||
|
assert provider.chat.called
|
||||||
|
call_args = provider.chat.call_args
|
||||||
|
messages = call_args.kwargs["messages"]
|
||||||
|
|
||||||
|
# The user message should contain the content
|
||||||
|
# (We can't easily check InboundMessage directly, but we verify
|
||||||
|
# the flow worked by checking the session was created)
|
||||||
|
session = loop.sessions.get_or_create("cli:direct")
|
||||||
|
assert len(session.messages) > 0
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_suppress_mode_adds_hidden_prefix():
|
||||||
|
"""Test that suppress_output metadata adds signed [HIDDEN:{sig}] prefix."""
|
||||||
|
bus = MessageBus()
|
||||||
|
provider = MagicMock(spec=LLMProvider)
|
||||||
|
provider.chat = AsyncMock(return_value=LLMResponse(
|
||||||
|
content="This is the agent response",
|
||||||
|
tool_calls=[]
|
||||||
|
))
|
||||||
|
provider.get_default_model = MagicMock(return_value="test-model")
|
||||||
|
provider.thinking_budget = 0
|
||||||
|
|
||||||
|
workspace = Path("/tmp/test-workspace")
|
||||||
|
workspace.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
|
||||||
|
|
||||||
|
# Call with suppress_output=True
|
||||||
|
response = await loop.process_direct(
|
||||||
|
content="test message",
|
||||||
|
metadata={"suppress_output": True}
|
||||||
|
)
|
||||||
|
|
||||||
|
# Response content should have signed [HIDDEN:{sig}] prefix
|
||||||
|
assert response.startswith("[HIDDEN:")
|
||||||
|
assert "] " in response # Check for signature end
|
||||||
|
assert "This is the agent response" in response
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_normal_mode_no_hidden_prefix():
|
||||||
|
"""Test that normal messages don't get [HIDDEN:*] prefix."""
|
||||||
|
bus = MessageBus()
|
||||||
|
provider = MagicMock(spec=LLMProvider)
|
||||||
|
provider.chat = AsyncMock(return_value=LLMResponse(
|
||||||
|
content="Normal response",
|
||||||
|
tool_calls=[]
|
||||||
|
))
|
||||||
|
provider.get_default_model = MagicMock(return_value="test-model")
|
||||||
|
provider.thinking_budget = 0
|
||||||
|
|
||||||
|
workspace = Path("/tmp/test-workspace")
|
||||||
|
workspace.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
loop = AgentLoop(bus=bus, provider=provider, workspace=workspace)
|
||||||
|
|
||||||
|
# Call without suppress_output
|
||||||
|
response = await loop.process_direct(content="test message")
|
||||||
|
|
||||||
|
# Response should NOT have [HIDDEN:*] prefix
|
||||||
|
assert not response.startswith("[HIDDEN:")
|
||||||
|
assert response == "Normal response"
|
||||||
@@ -0,0 +1,88 @@
|
|||||||
|
"""Test Anthropic OAuth provider."""
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import AsyncMock, patch, MagicMock
|
||||||
|
from nanobot.providers.anthropic_oauth import AnthropicOAuthProvider
|
||||||
|
from nanobot.providers.base import LLMResponse
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def provider():
|
||||||
|
"""Create provider with test OAuth token."""
|
||||||
|
return AnthropicOAuthProvider(
|
||||||
|
oauth_token="sk-ant-oat01-test-token",
|
||||||
|
default_model="claude-opus-4-5"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_provider_init(provider):
|
||||||
|
"""Provider should initialize with OAuth token."""
|
||||||
|
assert provider.oauth_token == "sk-ant-oat01-test-token"
|
||||||
|
assert provider.default_model == "claude-opus-4-5"
|
||||||
|
|
||||||
|
|
||||||
|
def test_provider_uses_bearer_auth(provider):
|
||||||
|
"""Provider should use Bearer auth, not x-api-key."""
|
||||||
|
headers = provider._get_headers()
|
||||||
|
assert "Authorization" in headers
|
||||||
|
assert headers["Authorization"].startswith("Bearer ")
|
||||||
|
assert "x-api-key" not in headers
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_chat_prepends_system_prompt(provider):
|
||||||
|
"""Chat should prepend Claude Code identity to system prompt."""
|
||||||
|
messages = [{"role": "user", "content": "Hello"}]
|
||||||
|
|
||||||
|
with patch.object(provider, "_make_request", new_callable=AsyncMock) as mock:
|
||||||
|
mock.return_value = {"content": [{"type": "text", "text": "Hi"}], "stop_reason": "end_turn"}
|
||||||
|
await provider.chat(messages)
|
||||||
|
|
||||||
|
call_args = mock.call_args
|
||||||
|
system = call_args[1]["system"]
|
||||||
|
assert "Claude Code" in system
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_response_text(provider):
|
||||||
|
"""Should parse text response correctly."""
|
||||||
|
response = {
|
||||||
|
"content": [{"type": "text", "text": "Hello world"}],
|
||||||
|
"stop_reason": "end_turn",
|
||||||
|
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||||
|
}
|
||||||
|
result = provider._parse_response(response)
|
||||||
|
assert result.content == "Hello world"
|
||||||
|
assert result.finish_reason == "end_turn"
|
||||||
|
assert result.usage["prompt_tokens"] == 10
|
||||||
|
|
||||||
|
|
||||||
|
def test_parse_response_tool_calls(provider):
|
||||||
|
"""Should parse tool call response correctly."""
|
||||||
|
response = {
|
||||||
|
"content": [
|
||||||
|
{"type": "tool_use", "id": "call_1", "name": "read_file", "input": {"path": "/tmp/test"}}
|
||||||
|
],
|
||||||
|
"stop_reason": "tool_use",
|
||||||
|
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||||
|
}
|
||||||
|
result = provider._parse_response(response)
|
||||||
|
assert len(result.tool_calls) == 1
|
||||||
|
assert result.tool_calls[0].name == "read_file"
|
||||||
|
assert result.tool_calls[0].arguments == {"path": "/tmp/test"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_convert_tools_to_anthropic(provider):
|
||||||
|
"""Should convert OpenAI-format tools to Anthropic format."""
|
||||||
|
openai_tools = [
|
||||||
|
{
|
||||||
|
"type": "function",
|
||||||
|
"function": {
|
||||||
|
"name": "read_file",
|
||||||
|
"description": "Read a file",
|
||||||
|
"parameters": {"type": "object", "properties": {"path": {"type": "string"}}}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
]
|
||||||
|
anthropic_tools = provider._convert_tools_to_anthropic(openai_tools)
|
||||||
|
assert len(anthropic_tools) == 1
|
||||||
|
assert anthropic_tools[0]["name"] == "read_file"
|
||||||
|
assert "input_schema" in anthropic_tools[0]
|
||||||
@@ -0,0 +1,59 @@
|
|||||||
|
"""Tests for bus-level correlation (request-response via Futures)."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import pytest
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.bus.events import OutboundMessage
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def bus():
|
||||||
|
return MessageBus()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_register_correlation_returns_future(bus):
|
||||||
|
future = bus.register_correlation("test-id-1")
|
||||||
|
assert isinstance(future, asyncio.Future)
|
||||||
|
assert not future.done()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_resolve_correlation_sets_future_result(bus):
|
||||||
|
future = bus.register_correlation("test-id-1")
|
||||||
|
msg = OutboundMessage(channel="hook", chat_id="test", content="hello", metadata={"correlation_id": "test-id-1"})
|
||||||
|
bus.resolve_correlation(msg)
|
||||||
|
assert future.done()
|
||||||
|
assert future.result() == "hello"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_resolve_correlation_no_match_is_noop(bus):
|
||||||
|
future = bus.register_correlation("test-id-1")
|
||||||
|
msg = OutboundMessage(channel="hook", chat_id="test", content="hello", metadata={"correlation_id": "other-id"})
|
||||||
|
bus.resolve_correlation(msg)
|
||||||
|
assert not future.done()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_resolve_correlation_no_metadata_is_noop(bus):
|
||||||
|
future = bus.register_correlation("test-id-1")
|
||||||
|
msg = OutboundMessage(channel="hook", chat_id="test", content="hello")
|
||||||
|
bus.resolve_correlation(msg)
|
||||||
|
assert not future.done()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_resolve_correlation_cleans_up_store(bus):
|
||||||
|
future = bus.register_correlation("test-id-1")
|
||||||
|
msg = OutboundMessage(channel="hook", chat_id="test", content="hello", metadata={"correlation_id": "test-id-1"})
|
||||||
|
bus.resolve_correlation(msg)
|
||||||
|
assert "test-id-1" not in bus._correlation_store
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_cancel_correlation(bus):
|
||||||
|
future = bus.register_correlation("test-id-1")
|
||||||
|
bus.cancel_correlation("test-id-1")
|
||||||
|
assert "test-id-1" not in bus._correlation_store
|
||||||
|
assert future.cancelled()
|
||||||
@@ -0,0 +1,55 @@
|
|||||||
|
"""Test OAuth CLI commands."""
|
||||||
|
import pytest
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from typer.testing import CliRunner
|
||||||
|
from nanobot.cli.oauth import oauth_app
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def runner():
|
||||||
|
return CliRunner()
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_login_help(runner):
|
||||||
|
"""Login command should have help text."""
|
||||||
|
result = runner.invoke(oauth_app, ["login", "--help"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "token" in result.output.lower()
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_status_no_credentials(runner, tmp_path, monkeypatch):
|
||||||
|
"""Status should show no credentials when none exist."""
|
||||||
|
monkeypatch.setenv("HOME", str(tmp_path))
|
||||||
|
result = runner.invoke(oauth_app, ["status"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "No OAuth credentials" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_login_and_status(runner, tmp_path, monkeypatch):
|
||||||
|
"""Login should save credentials, status should show them."""
|
||||||
|
monkeypatch.setenv("HOME", str(tmp_path))
|
||||||
|
result = runner.invoke(oauth_app, ["login", "--token", "sk-ant-oat01-test-xxx"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "Successfully saved" in result.output
|
||||||
|
|
||||||
|
result = runner.invoke(oauth_app, ["status"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "sk-ant-oat01-test-x" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_logout(runner, tmp_path, monkeypatch):
|
||||||
|
"""Logout should remove credentials."""
|
||||||
|
monkeypatch.setenv("HOME", str(tmp_path))
|
||||||
|
runner.invoke(oauth_app, ["login", "--token", "sk-ant-oat01-test-xxx"])
|
||||||
|
result = runner.invoke(oauth_app, ["logout"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "Removed" in result.output
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_login_invalid_token(runner, tmp_path, monkeypatch):
|
||||||
|
"""Login should reject non-OAuth tokens."""
|
||||||
|
monkeypatch.setenv("HOME", str(tmp_path))
|
||||||
|
result = runner.invoke(oauth_app, ["login", "--token", "sk-ant-api03-regular"])
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "Invalid token" in result.output
|
||||||
@@ -0,0 +1,65 @@
|
|||||||
|
"""Test OAuth store integration with config loading."""
|
||||||
|
import json
|
||||||
|
import pytest
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from nanobot.config.loader import load_config
|
||||||
|
from nanobot.config.oauth_store import OAuthStore
|
||||||
|
from nanobot.config.schema import OAuthCredentials
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_token_injected_into_config(tmp_path, monkeypatch):
|
||||||
|
"""OAuth token from store should be injected into provider api_key."""
|
||||||
|
# Create a minimal config file (no api key set)
|
||||||
|
config_path = tmp_path / "config.json"
|
||||||
|
config_path.write_text(json.dumps({
|
||||||
|
"agents": {"defaults": {"model": "anthropic/claude-opus-4-5"}},
|
||||||
|
"providers": {"anthropic": {"apiKey": ""}}
|
||||||
|
}))
|
||||||
|
|
||||||
|
# Save OAuth credentials
|
||||||
|
store = OAuthStore(tmp_path)
|
||||||
|
creds = OAuthCredentials(access_token="sk-ant-oat01-test-inject")
|
||||||
|
store.save("anthropic", creds)
|
||||||
|
|
||||||
|
# Monkeypatch get_config_path to use our tmp dir
|
||||||
|
monkeypatch.setattr("nanobot.config.loader.get_config_path", lambda: config_path)
|
||||||
|
# Monkeypatch the OAuth store path
|
||||||
|
monkeypatch.setattr("nanobot.config.loader._get_oauth_store_dir", lambda: tmp_path)
|
||||||
|
|
||||||
|
config = load_config(config_path)
|
||||||
|
|
||||||
|
assert config.providers.anthropic.api_key == "sk-ant-oat01-test-inject"
|
||||||
|
|
||||||
|
|
||||||
|
def test_config_without_oauth_unchanged(tmp_path, monkeypatch):
|
||||||
|
"""Config without OAuth store should load normally."""
|
||||||
|
config_path = tmp_path / "config.json"
|
||||||
|
config_path.write_text(json.dumps({
|
||||||
|
"providers": {"anthropic": {"apiKey": "sk-ant-api03-regular"}}
|
||||||
|
}))
|
||||||
|
|
||||||
|
monkeypatch.setattr("nanobot.config.loader._get_oauth_store_dir", lambda: tmp_path / "nonexistent")
|
||||||
|
|
||||||
|
config = load_config(config_path)
|
||||||
|
|
||||||
|
assert config.providers.anthropic.api_key == "sk-ant-api03-regular"
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_does_not_overwrite_existing_key(tmp_path, monkeypatch):
|
||||||
|
"""If user already has an API key, OAuth should still override (OAuth takes priority)."""
|
||||||
|
config_path = tmp_path / "config.json"
|
||||||
|
config_path.write_text(json.dumps({
|
||||||
|
"providers": {"anthropic": {"apiKey": "sk-ant-api03-existing"}}
|
||||||
|
}))
|
||||||
|
|
||||||
|
store = OAuthStore(tmp_path)
|
||||||
|
creds = OAuthCredentials(access_token="sk-ant-oat01-oauth-wins")
|
||||||
|
store.save("anthropic", creds)
|
||||||
|
|
||||||
|
monkeypatch.setattr("nanobot.config.loader._get_oauth_store_dir", lambda: tmp_path)
|
||||||
|
|
||||||
|
config = load_config(config_path)
|
||||||
|
|
||||||
|
# OAuth token takes priority over existing API key
|
||||||
|
assert config.providers.anthropic.api_key == "sk-ant-oat01-oauth-wins"
|
||||||
@@ -0,0 +1,95 @@
|
|||||||
|
# tests/test_heartbeat_idle.py
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from nanobot.heartbeat.service import HeartbeatService
|
||||||
|
from nanobot.session.manager import SessionManager
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_heartbeat_skips_when_user_active():
|
||||||
|
"""Test that heartbeat doesn't trigger if user messaged recently."""
|
||||||
|
workspace = Path("/tmp/test-heartbeat")
|
||||||
|
workspace.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
# Create session with recent user message
|
||||||
|
sessions = SessionManager(workspace)
|
||||||
|
session = sessions.get_or_create("telegram:239824268")
|
||||||
|
session.add_message("user", "Recent message")
|
||||||
|
sessions.save(session)
|
||||||
|
|
||||||
|
# Create heartbeat callback
|
||||||
|
callback_called = False
|
||||||
|
async def on_heartbeat(prompt, metadata=None):
|
||||||
|
nonlocal callback_called
|
||||||
|
callback_called = True
|
||||||
|
return "response"
|
||||||
|
|
||||||
|
# Create heartbeat service
|
||||||
|
service = HeartbeatService(
|
||||||
|
workspace=workspace,
|
||||||
|
on_heartbeat=on_heartbeat,
|
||||||
|
interval_s=1, # Short interval for testing
|
||||||
|
enabled=True,
|
||||||
|
session_manager=sessions,
|
||||||
|
target_session_key="telegram:239824268"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Trigger heartbeat
|
||||||
|
await service._tick()
|
||||||
|
|
||||||
|
# Callback should NOT have been called (user was active recently)
|
||||||
|
assert not callback_called
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_heartbeat_triggers_when_user_idle():
|
||||||
|
"""Test that heartbeat triggers after 30min of inactivity."""
|
||||||
|
workspace = Path("/tmp/test-heartbeat")
|
||||||
|
workspace.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
# Create HEARTBEAT.md with content
|
||||||
|
heartbeat_file = workspace / "HEARTBEAT.md"
|
||||||
|
heartbeat_file.write_text("# Tasks\n- Check something\n")
|
||||||
|
|
||||||
|
# Create session with old user message (>30min ago)
|
||||||
|
sessions = SessionManager(workspace)
|
||||||
|
session = sessions.get_or_create("telegram:239824268")
|
||||||
|
|
||||||
|
# Manually set old timestamp
|
||||||
|
old_timestamp = (datetime.now() - timedelta(minutes=31)).isoformat()
|
||||||
|
session.messages.append({
|
||||||
|
"role": "user",
|
||||||
|
"content": "Old message",
|
||||||
|
"timestamp": old_timestamp
|
||||||
|
})
|
||||||
|
sessions.save(session)
|
||||||
|
|
||||||
|
# Create heartbeat callback
|
||||||
|
callback_called = False
|
||||||
|
callback_metadata = None
|
||||||
|
|
||||||
|
async def on_heartbeat(prompt, metadata=None):
|
||||||
|
nonlocal callback_called, callback_metadata
|
||||||
|
callback_called = True
|
||||||
|
callback_metadata = metadata
|
||||||
|
return "response"
|
||||||
|
|
||||||
|
# Create heartbeat service
|
||||||
|
service = HeartbeatService(
|
||||||
|
workspace=workspace,
|
||||||
|
on_heartbeat=on_heartbeat,
|
||||||
|
interval_s=1,
|
||||||
|
enabled=True,
|
||||||
|
session_manager=sessions,
|
||||||
|
target_session_key="telegram:239824268"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Trigger heartbeat
|
||||||
|
await service._tick()
|
||||||
|
|
||||||
|
# Callback SHOULD have been called (user idle for >30min)
|
||||||
|
assert callback_called
|
||||||
|
assert callback_metadata == {"suppress_output": True}
|
||||||
@@ -0,0 +1,37 @@
|
|||||||
|
"""Tests for the hook channel."""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import MagicMock
|
||||||
|
from nanobot.channels.hook import HookChannel
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.bus.events import OutboundMessage
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def bus():
|
||||||
|
return MessageBus()
|
||||||
|
|
||||||
|
|
||||||
|
def test_hook_channel_name():
|
||||||
|
bus = MessageBus()
|
||||||
|
channel = HookChannel(bus)
|
||||||
|
assert channel.name == "hook"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_hook_channel_send_is_noop():
|
||||||
|
"""send() should not raise and should not do anything."""
|
||||||
|
bus = MessageBus()
|
||||||
|
channel = HookChannel(bus)
|
||||||
|
msg = OutboundMessage(channel="hook", chat_id="test", content="hello")
|
||||||
|
await channel.send(msg) # Should not raise
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_hook_channel_start_stop():
|
||||||
|
bus = MessageBus()
|
||||||
|
channel = HookChannel(bus)
|
||||||
|
await channel.start()
|
||||||
|
assert channel.is_running
|
||||||
|
await channel.stop()
|
||||||
|
assert not channel.is_running
|
||||||
@@ -0,0 +1,24 @@
|
|||||||
|
"""Tests for HooksConfig with named tokens."""
|
||||||
|
|
||||||
|
from nanobot.config.schema import HooksConfig
|
||||||
|
|
||||||
|
|
||||||
|
def test_hooks_config_named_tokens():
|
||||||
|
"""Named tokens dict should work."""
|
||||||
|
config = HooksConfig(enabled=True, tokens={"gitea": "secret1", "ha": "secret2"})
|
||||||
|
assert config.tokens == {"gitea": "secret1", "ha": "secret2"}
|
||||||
|
|
||||||
|
|
||||||
|
def test_hooks_config_resolve_token():
|
||||||
|
"""resolve_token should return token name for a matching token."""
|
||||||
|
config = HooksConfig(enabled=True, tokens={"gitea": "secret1", "ha": "secret2"})
|
||||||
|
assert config.resolve_token("secret1") == "gitea"
|
||||||
|
assert config.resolve_token("secret2") == "ha"
|
||||||
|
assert config.resolve_token("unknown") is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_hooks_config_has_tokens():
|
||||||
|
"""has_tokens should be True if tokens dict is non-empty."""
|
||||||
|
assert HooksConfig(enabled=True, tokens={"webhook": "secret"}).has_tokens
|
||||||
|
assert not HooksConfig(enabled=True).has_tokens
|
||||||
|
assert not HooksConfig(enabled=True, tokens={}).has_tokens
|
||||||
@@ -0,0 +1,128 @@
|
|||||||
|
"""End-to-end integration test for hooks → bus → correlation → response."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import pytest
|
||||||
|
from aiohttp.test_utils import TestClient, TestServer
|
||||||
|
from nanobot.hooks.server import HooksServer
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||||
|
from nanobot.config.schema import HooksConfig
|
||||||
|
from nanobot.channels.hook import HookChannel
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def bus():
|
||||||
|
return MessageBus()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def config():
|
||||||
|
return HooksConfig(
|
||||||
|
enabled=True,
|
||||||
|
tokens={"gitea": "gitea-secret", "ha": "ha-secret"},
|
||||||
|
timeout_seconds=5,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def server(bus, config):
|
||||||
|
return HooksServer(host="127.0.0.1", port=0, config=config, bus=bus)
|
||||||
|
|
||||||
|
|
||||||
|
async def fake_agent_loop(bus: MessageBus):
|
||||||
|
"""Simulate agent loop: consume inbound, process, publish outbound."""
|
||||||
|
msg = await asyncio.wait_for(bus.consume_inbound(), timeout=3.0)
|
||||||
|
response_content = f"Processed: {msg.content}"
|
||||||
|
await bus.publish_outbound(OutboundMessage(
|
||||||
|
channel=msg.channel,
|
||||||
|
chat_id=msg.chat_id,
|
||||||
|
content=response_content,
|
||||||
|
metadata=msg.metadata or {},
|
||||||
|
))
|
||||||
|
|
||||||
|
|
||||||
|
async def fake_dispatch_loop(bus: MessageBus, hook_channel: HookChannel):
|
||||||
|
"""Simulate outbound dispatcher: consume outbound, resolve correlation, dispatch."""
|
||||||
|
msg = await asyncio.wait_for(bus.consume_outbound(), timeout=3.0)
|
||||||
|
bus.resolve_correlation(msg)
|
||||||
|
if msg.channel == "hook":
|
||||||
|
await hook_channel.send(msg)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_full_hook_flow_default_channel(server, bus):
|
||||||
|
"""Hook with default channel: message goes through bus, response returned to HTTP caller."""
|
||||||
|
hook_channel = HookChannel(bus)
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
async def do_request():
|
||||||
|
return await client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "deploy started"},
|
||||||
|
headers={"Authorization": "Bearer gitea-secret"},
|
||||||
|
)
|
||||||
|
|
||||||
|
# Run request + fake agent + fake dispatcher concurrently
|
||||||
|
request_task = asyncio.create_task(do_request())
|
||||||
|
agent_task = asyncio.create_task(fake_agent_loop(bus))
|
||||||
|
dispatch_task = asyncio.create_task(fake_dispatch_loop(bus, hook_channel))
|
||||||
|
|
||||||
|
resp = await asyncio.wait_for(request_task, timeout=5.0)
|
||||||
|
await agent_task
|
||||||
|
await dispatch_task
|
||||||
|
|
||||||
|
assert resp.status == 200
|
||||||
|
data = await resp.json()
|
||||||
|
assert data["ok"] is True
|
||||||
|
assert "deploy started" in data["response"]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_full_hook_flow_telegram_channel(server, bus):
|
||||||
|
"""Hook targeting telegram: uses telegram session, response still returned to HTTP caller."""
|
||||||
|
hook_channel = HookChannel(bus)
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
async def do_request():
|
||||||
|
return await client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "doorbell rang", "channel": "telegram", "chat_id": "239824268"},
|
||||||
|
headers={"Authorization": "Bearer ha-secret"},
|
||||||
|
)
|
||||||
|
|
||||||
|
request_task = asyncio.create_task(do_request())
|
||||||
|
agent_task = asyncio.create_task(fake_agent_loop(bus))
|
||||||
|
dispatch_task = asyncio.create_task(fake_dispatch_loop(bus, hook_channel))
|
||||||
|
|
||||||
|
resp = await asyncio.wait_for(request_task, timeout=5.0)
|
||||||
|
await agent_task
|
||||||
|
await dispatch_task
|
||||||
|
|
||||||
|
assert resp.status == 200
|
||||||
|
data = await resp.json()
|
||||||
|
assert data["ok"] is True
|
||||||
|
assert "doorbell rang" in data["response"]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_named_token_identification(server, bus):
|
||||||
|
"""Different tokens should produce different hook_source in metadata."""
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
asyncio.create_task(client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "from gitea"},
|
||||||
|
headers={"Authorization": "Bearer gitea-secret"},
|
||||||
|
))
|
||||||
|
msg1 = await asyncio.wait_for(bus.consume_inbound(), timeout=2.0)
|
||||||
|
assert msg1.metadata["hook_source"] == "gitea"
|
||||||
|
assert msg1.chat_id == "gitea"
|
||||||
|
|
||||||
|
asyncio.create_task(client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "from ha"},
|
||||||
|
headers={"Authorization": "Bearer ha-secret"},
|
||||||
|
))
|
||||||
|
msg2 = await asyncio.wait_for(bus.consume_inbound(), timeout=2.0)
|
||||||
|
assert msg2.metadata["hook_source"] == "ha"
|
||||||
|
assert msg2.chat_id == "ha"
|
||||||
@@ -0,0 +1,176 @@
|
|||||||
|
"""Tests for the rewritten hooks server."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import json
|
||||||
|
import pytest
|
||||||
|
from aiohttp import web
|
||||||
|
from aiohttp.test_utils import AioHTTPTestCase, unittest_run_loop, TestClient, TestServer
|
||||||
|
from nanobot.hooks.server import HooksServer
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||||
|
from nanobot.config.schema import HooksConfig
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def bus():
|
||||||
|
return MessageBus()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def config():
|
||||||
|
return HooksConfig(
|
||||||
|
enabled=True,
|
||||||
|
tokens={"test-hook": "test-secret-123"},
|
||||||
|
timeout_seconds=5,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def server(bus, config):
|
||||||
|
return HooksServer(host="127.0.0.1", port=0, config=config, bus=bus)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_health_check(server):
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
resp = await client.get("/health")
|
||||||
|
assert resp.status == 200
|
||||||
|
data = await resp.json()
|
||||||
|
assert data["status"] == "ok"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_unauthorized_without_token(server):
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
resp = await client.post("/hooks", json={"message": "test"})
|
||||||
|
assert resp.status == 401
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_unauthorized_wrong_token(server):
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
resp = await client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "test"},
|
||||||
|
headers={"Authorization": "Bearer wrong-token"},
|
||||||
|
)
|
||||||
|
assert resp.status == 401
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_missing_message_field(server, bus):
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
resp = await client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"not_message": "test"},
|
||||||
|
headers={"Authorization": "Bearer test-secret-123"},
|
||||||
|
)
|
||||||
|
assert resp.status == 400
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_hook_publishes_to_bus(server, bus):
|
||||||
|
"""Hook should publish InboundMessage to bus and the message should contain hook prefix."""
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
# Send hook request in background (it will block waiting for correlation)
|
||||||
|
async def send_request():
|
||||||
|
return await client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "hello from webhook"},
|
||||||
|
headers={"Authorization": "Bearer test-secret-123"},
|
||||||
|
)
|
||||||
|
|
||||||
|
task = asyncio.create_task(send_request())
|
||||||
|
|
||||||
|
# Consume the inbound message
|
||||||
|
msg = await asyncio.wait_for(bus.consume_inbound(), timeout=2.0)
|
||||||
|
|
||||||
|
assert msg.channel == "hook"
|
||||||
|
assert msg.chat_id == "test-hook" # defaults to token name
|
||||||
|
assert msg.metadata.get("hook_source") == "test-hook"
|
||||||
|
assert msg.metadata.get("correlation_id") is not None
|
||||||
|
|
||||||
|
# Simulate agent response by resolving correlation
|
||||||
|
bus.resolve_correlation(OutboundMessage(
|
||||||
|
channel="hook",
|
||||||
|
chat_id="test-hook",
|
||||||
|
content="agent says hi",
|
||||||
|
metadata={"correlation_id": msg.metadata["correlation_id"]},
|
||||||
|
))
|
||||||
|
|
||||||
|
resp = await asyncio.wait_for(task, timeout=2.0)
|
||||||
|
assert resp.status == 200
|
||||||
|
data = await resp.json()
|
||||||
|
assert data["ok"] is True
|
||||||
|
assert data["response"] == "agent says hi"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_hook_with_custom_channel(server, bus):
|
||||||
|
"""Hook targeting telegram should use telegram channel in InboundMessage."""
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
async def send_request():
|
||||||
|
return await client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "notify user", "channel": "telegram", "chat_id": "239824268"},
|
||||||
|
headers={"Authorization": "Bearer test-secret-123"},
|
||||||
|
)
|
||||||
|
|
||||||
|
task = asyncio.create_task(send_request())
|
||||||
|
|
||||||
|
msg = await asyncio.wait_for(bus.consume_inbound(), timeout=2.0)
|
||||||
|
assert msg.channel == "telegram"
|
||||||
|
assert msg.chat_id == "239824268"
|
||||||
|
assert msg.session_key == "telegram:239824268"
|
||||||
|
|
||||||
|
bus.resolve_correlation(OutboundMessage(
|
||||||
|
channel="telegram",
|
||||||
|
chat_id="239824268",
|
||||||
|
content="done",
|
||||||
|
metadata={"correlation_id": msg.metadata["correlation_id"]},
|
||||||
|
))
|
||||||
|
|
||||||
|
resp = await asyncio.wait_for(task, timeout=2.0)
|
||||||
|
assert resp.status == 200
|
||||||
|
data = await resp.json()
|
||||||
|
assert data["response"] == "done"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_hook_timeout_returns_504(bus):
|
||||||
|
"""If agent doesn't respond in time, return 504."""
|
||||||
|
config = HooksConfig(enabled=True, tokens={"test-hook": "test-secret-123"}, timeout_seconds=1)
|
||||||
|
server = HooksServer(host="127.0.0.1", port=0, config=config, bus=bus)
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
resp = await client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "slow request"},
|
||||||
|
headers={"Authorization": "Bearer test-secret-123"},
|
||||||
|
)
|
||||||
|
assert resp.status == 504
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_hook_timeout_zero_returns_202(server, bus):
|
||||||
|
"""timeout=0 should return 202 immediately without waiting."""
|
||||||
|
client = TestClient(TestServer(server._app))
|
||||||
|
async with client:
|
||||||
|
resp = await client.post(
|
||||||
|
"/hooks",
|
||||||
|
json={"message": "fire and forget", "timeout": 0},
|
||||||
|
headers={"Authorization": "Bearer test-secret-123"},
|
||||||
|
)
|
||||||
|
assert resp.status == 202
|
||||||
|
data = await resp.json()
|
||||||
|
assert data["ok"] is True
|
||||||
|
|
||||||
|
# Message should still be on the bus
|
||||||
|
msg = await asyncio.wait_for(bus.consume_inbound(), timeout=1.0)
|
||||||
|
assert msg.content == "fire and forget"
|
||||||
@@ -0,0 +1,105 @@
|
|||||||
|
# tests/test_idle_heartbeat_integration.py
|
||||||
|
import pytest
|
||||||
|
import asyncio
|
||||||
|
from pathlib import Path
|
||||||
|
from datetime import datetime, timedelta
|
||||||
|
from nanobot.agent.loop import AgentLoop
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.heartbeat.service import HeartbeatService
|
||||||
|
from nanobot.session.manager import SessionManager
|
||||||
|
from nanobot.providers.base import LLMProvider, LLMResponse
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_idle_heartbeat_end_to_end(tmp_path):
|
||||||
|
"""
|
||||||
|
Integration test: heartbeat triggers when idle, runs in main session,
|
||||||
|
output is suppressed, session contains [HIDDEN] content.
|
||||||
|
"""
|
||||||
|
workspace = tmp_path / "test-integration"
|
||||||
|
workspace.mkdir()
|
||||||
|
|
||||||
|
# Use test-specific session key
|
||||||
|
test_session_key = "telegram:test_integration"
|
||||||
|
|
||||||
|
# Create HEARTBEAT.md with content
|
||||||
|
heartbeat_file = workspace / "HEARTBEAT.md"
|
||||||
|
heartbeat_file.write_text("# Test Task\n- Check something")
|
||||||
|
|
||||||
|
# Create mock provider
|
||||||
|
provider = MagicMock(spec=LLMProvider)
|
||||||
|
provider.chat = AsyncMock(return_value=LLMResponse(
|
||||||
|
content="Heartbeat executed successfully",
|
||||||
|
tool_calls=[] # has_tool_calls is a property, not a parameter
|
||||||
|
))
|
||||||
|
provider.get_default_model = MagicMock(return_value="test-model")
|
||||||
|
provider.thinking_budget = 0
|
||||||
|
|
||||||
|
# Create components
|
||||||
|
bus = MessageBus()
|
||||||
|
sessions = SessionManager(workspace)
|
||||||
|
# Override sessions_dir to use tmp_path for test isolation
|
||||||
|
sessions.sessions_dir = tmp_path / "sessions"
|
||||||
|
sessions.sessions_dir.mkdir()
|
||||||
|
loop = AgentLoop(
|
||||||
|
bus=bus,
|
||||||
|
provider=provider,
|
||||||
|
workspace=workspace,
|
||||||
|
session_manager=sessions
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create session with old user message
|
||||||
|
session = sessions.get_or_create(test_session_key)
|
||||||
|
old_timestamp = (datetime.now() - timedelta(minutes=31)).isoformat()
|
||||||
|
session.messages.append({
|
||||||
|
"role": "user",
|
||||||
|
"content": "Old user message",
|
||||||
|
"timestamp": old_timestamp
|
||||||
|
})
|
||||||
|
sessions.save(session)
|
||||||
|
|
||||||
|
# Create heartbeat callback
|
||||||
|
async def on_heartbeat(prompt: str, metadata=None):
|
||||||
|
return await loop.process_direct(
|
||||||
|
prompt,
|
||||||
|
session_key=test_session_key,
|
||||||
|
channel="telegram",
|
||||||
|
chat_id="test_integration",
|
||||||
|
metadata=metadata,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Create heartbeat service
|
||||||
|
heartbeat = HeartbeatService(
|
||||||
|
workspace=workspace,
|
||||||
|
on_heartbeat=on_heartbeat,
|
||||||
|
interval_s=1,
|
||||||
|
enabled=True,
|
||||||
|
session_manager=sessions,
|
||||||
|
target_session_key=test_session_key,
|
||||||
|
idle_threshold_s=30 * 60,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Trigger heartbeat
|
||||||
|
await heartbeat._tick()
|
||||||
|
|
||||||
|
# Reload session from disk
|
||||||
|
sessions._cache.clear() # Clear cache to force reload
|
||||||
|
session = sessions.get_or_create(test_session_key)
|
||||||
|
|
||||||
|
# Verify:
|
||||||
|
# 1. Session has new messages
|
||||||
|
assert len(session.messages) > 1
|
||||||
|
|
||||||
|
# 2. Find the heartbeat response (assistant message with signed marker)
|
||||||
|
heartbeat_messages = [
|
||||||
|
m for m in session.messages
|
||||||
|
if m.get("role") == "assistant" and "[HIDDEN:" in m.get("content", "")
|
||||||
|
]
|
||||||
|
assert len(heartbeat_messages) == 1, "Expected exactly 1 signed [HIDDEN:*] heartbeat message"
|
||||||
|
|
||||||
|
# 3. Verify content is prefixed with signed [HIDDEN:{sig}] marker
|
||||||
|
heartbeat_msg = heartbeat_messages[0]
|
||||||
|
assert heartbeat_msg["content"].startswith("[HIDDEN:")
|
||||||
|
assert "] " in heartbeat_msg["content"] # Check for signature end
|
||||||
|
assert "Heartbeat executed successfully" in heartbeat_msg["content"]
|
||||||
@@ -0,0 +1,48 @@
|
|||||||
|
"""Test OAuth configuration schema."""
|
||||||
|
import pytest
|
||||||
|
from nanobot.config.schema import ProviderConfig, OAuthCredentials
|
||||||
|
|
||||||
|
|
||||||
|
def test_provider_config_has_oauth_fields():
|
||||||
|
"""ProviderConfig should have oauth_credentials field."""
|
||||||
|
config = ProviderConfig(api_key="test")
|
||||||
|
assert hasattr(config, "oauth_credentials")
|
||||||
|
assert config.oauth_credentials is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_credentials_model():
|
||||||
|
"""OAuthCredentials should store token, refresh, expiry."""
|
||||||
|
creds = OAuthCredentials(
|
||||||
|
access_token="sk-ant-oat01-xxx",
|
||||||
|
refresh_token="rt_xxx",
|
||||||
|
expires_at=1234567890,
|
||||||
|
token_type="oauth"
|
||||||
|
)
|
||||||
|
assert creds.access_token.startswith("sk-ant-oat")
|
||||||
|
assert creds.is_oauth_token is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_credentials_expiry_check():
|
||||||
|
"""OAuthCredentials should detect expired tokens."""
|
||||||
|
import time
|
||||||
|
expired = OAuthCredentials(
|
||||||
|
access_token="sk-ant-oat01-xxx",
|
||||||
|
expires_at=int(time.time()) - 3600 # 1 hour ago
|
||||||
|
)
|
||||||
|
assert expired.is_expired is True
|
||||||
|
|
||||||
|
valid = OAuthCredentials(
|
||||||
|
access_token="sk-ant-oat01-xxx",
|
||||||
|
expires_at=int(time.time()) + 3600 # 1 hour from now
|
||||||
|
)
|
||||||
|
assert valid.is_expired is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_oauth_credentials_no_expiry():
|
||||||
|
"""Setup tokens with expires_at=0 should never be expired."""
|
||||||
|
creds = OAuthCredentials(
|
||||||
|
access_token="sk-ant-oat01-xxx",
|
||||||
|
expires_at=0 # No expiry (setup-token)
|
||||||
|
)
|
||||||
|
assert creds.is_expired is False
|
||||||
|
assert creds.expires_soon is False
|
||||||
@@ -0,0 +1,55 @@
|
|||||||
|
"""Test OAuth credential storage."""
|
||||||
|
import pytest
|
||||||
|
import tempfile
|
||||||
|
from pathlib import Path
|
||||||
|
from nanobot.config.oauth_store import OAuthStore
|
||||||
|
from nanobot.config.schema import OAuthCredentials
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def temp_store():
|
||||||
|
"""Create store with temp directory."""
|
||||||
|
with tempfile.TemporaryDirectory() as tmpdir:
|
||||||
|
yield OAuthStore(Path(tmpdir) / ".nanobot")
|
||||||
|
|
||||||
|
|
||||||
|
def test_save_and_load_credentials(temp_store):
|
||||||
|
"""Should save and load OAuth credentials."""
|
||||||
|
creds = OAuthCredentials(
|
||||||
|
access_token="sk-ant-oat01-xxx",
|
||||||
|
refresh_token="rt_xxx",
|
||||||
|
expires_at=1234567890
|
||||||
|
)
|
||||||
|
|
||||||
|
temp_store.save("anthropic", creds)
|
||||||
|
loaded = temp_store.load("anthropic")
|
||||||
|
|
||||||
|
assert loaded is not None
|
||||||
|
assert loaded.access_token == creds.access_token
|
||||||
|
assert loaded.refresh_token == creds.refresh_token
|
||||||
|
|
||||||
|
|
||||||
|
def test_load_nonexistent_returns_none(temp_store):
|
||||||
|
"""Should return None for missing credentials."""
|
||||||
|
assert temp_store.load("nonexistent") is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_delete_credentials(temp_store):
|
||||||
|
"""Should delete saved credentials."""
|
||||||
|
creds = OAuthCredentials(access_token="sk-ant-oat01-xxx")
|
||||||
|
temp_store.save("anthropic", creds)
|
||||||
|
assert temp_store.delete("anthropic") is True
|
||||||
|
assert temp_store.load("anthropic") is None
|
||||||
|
|
||||||
|
|
||||||
|
def test_delete_nonexistent_returns_false(temp_store):
|
||||||
|
"""Should return False when deleting missing credentials."""
|
||||||
|
assert temp_store.delete("nonexistent") is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_file_permissions(temp_store):
|
||||||
|
"""Credentials file should have restricted permissions."""
|
||||||
|
creds = OAuthCredentials(access_token="sk-ant-oat01-xxx")
|
||||||
|
temp_store.save("anthropic", creds)
|
||||||
|
perms = oct(temp_store.file_path.stat().st_mode)[-3:]
|
||||||
|
assert perms == "600"
|
||||||
@@ -0,0 +1,28 @@
|
|||||||
|
"""Test OAuth utility functions."""
|
||||||
|
import pytest
|
||||||
|
from nanobot.providers.oauth_utils import is_oauth_token, get_auth_headers
|
||||||
|
|
||||||
|
|
||||||
|
def test_is_oauth_token_detects_oat():
|
||||||
|
"""Should detect sk-ant-oat tokens as OAuth."""
|
||||||
|
assert is_oauth_token("sk-ant-oat01-buSdhCH2XEkebW7ZQZTvGqH5EwAFh4u52LrdJhAP") is True
|
||||||
|
assert is_oauth_token("sk-ant-api03-regularkey") is False
|
||||||
|
assert is_oauth_token("") is False
|
||||||
|
assert is_oauth_token(None) is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_auth_headers_oauth():
|
||||||
|
"""OAuth tokens should use Authorization: Bearer."""
|
||||||
|
headers = get_auth_headers("sk-ant-oat01-xxx", is_oauth=True)
|
||||||
|
assert "Authorization" in headers
|
||||||
|
assert headers["Authorization"] == "Bearer sk-ant-oat01-xxx"
|
||||||
|
assert "x-api-key" not in headers
|
||||||
|
assert headers["anthropic-beta"] == "claude-code-20250219,oauth-2025-04-20"
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_auth_headers_api_key():
|
||||||
|
"""Regular API keys should use x-api-key."""
|
||||||
|
headers = get_auth_headers("sk-ant-api03-xxx", is_oauth=False)
|
||||||
|
assert "x-api-key" in headers
|
||||||
|
assert headers["x-api-key"] == "sk-ant-api03-xxx"
|
||||||
|
assert "Authorization" not in headers
|
||||||
@@ -0,0 +1,24 @@
|
|||||||
|
"""Tests for correlation resolution in outbound dispatch."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import pytest
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.bus.events import OutboundMessage
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_dispatch_resolves_correlation_before_channel_send():
|
||||||
|
"""Correlation Future should be resolved when outbound message is dispatched."""
|
||||||
|
bus = MessageBus()
|
||||||
|
future = bus.register_correlation("corr-1")
|
||||||
|
|
||||||
|
msg = OutboundMessage(channel="telegram", chat_id="123", content="response", metadata={"correlation_id": "corr-1"})
|
||||||
|
await bus.publish_outbound(msg)
|
||||||
|
|
||||||
|
# Simulate what _dispatch_outbound does: consume + resolve
|
||||||
|
consumed = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
|
||||||
|
bus.resolve_correlation(consumed)
|
||||||
|
|
||||||
|
assert future.done()
|
||||||
|
assert future.result() == "response"
|
||||||
@@ -0,0 +1,32 @@
|
|||||||
|
"""Test provider factory with OAuth support."""
|
||||||
|
import pytest
|
||||||
|
from nanobot.providers import create_provider
|
||||||
|
from nanobot.providers.anthropic_oauth import AnthropicOAuthProvider
|
||||||
|
from nanobot.providers.litellm_provider import LiteLLMProvider
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_provider_oauth_token():
|
||||||
|
"""OAuth tokens should create AnthropicOAuthProvider."""
|
||||||
|
provider = create_provider(
|
||||||
|
api_key="sk-ant-oat01-test-token",
|
||||||
|
model="anthropic/claude-opus-4-5"
|
||||||
|
)
|
||||||
|
assert isinstance(provider, AnthropicOAuthProvider)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_provider_regular_key():
|
||||||
|
"""Regular API keys should create LiteLLMProvider."""
|
||||||
|
provider = create_provider(
|
||||||
|
api_key="sk-ant-api03-regular-key",
|
||||||
|
model="anthropic/claude-opus-4-5"
|
||||||
|
)
|
||||||
|
assert isinstance(provider, LiteLLMProvider)
|
||||||
|
|
||||||
|
|
||||||
|
def test_create_provider_openrouter():
|
||||||
|
"""OpenRouter keys should create LiteLLMProvider."""
|
||||||
|
provider = create_provider(
|
||||||
|
api_key="sk-or-v1-xxx",
|
||||||
|
model="anthropic/claude-opus-4-5"
|
||||||
|
)
|
||||||
|
assert isinstance(provider, LiteLLMProvider)
|
||||||
@@ -0,0 +1,21 @@
|
|||||||
|
"""Test OAuth detection in provider registry."""
|
||||||
|
import pytest
|
||||||
|
from nanobot.providers.registry import should_use_oauth_provider
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_use_oauth_for_oat_token():
|
||||||
|
"""OAuth provider should be used for sk-ant-oat tokens."""
|
||||||
|
assert should_use_oauth_provider("sk-ant-oat01-xxx", "anthropic/claude-opus-4-5") is True
|
||||||
|
assert should_use_oauth_provider("sk-ant-oat01-xxx", "claude-sonnet-4") is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_not_use_oauth_for_regular_key():
|
||||||
|
"""Regular API keys should not use OAuth provider."""
|
||||||
|
assert should_use_oauth_provider("sk-ant-api03-xxx", "claude-opus-4-5") is False
|
||||||
|
assert should_use_oauth_provider("sk-or-v1-xxx", "anthropic/claude-opus-4-5") is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_should_not_use_oauth_for_non_anthropic():
|
||||||
|
"""Non-Anthropic models should not use OAuth provider."""
|
||||||
|
assert should_use_oauth_provider("sk-ant-oat01-xxx", "gpt-4") is False
|
||||||
|
assert should_use_oauth_provider("sk-ant-oat01-xxx", "deepseek/deepseek-chat") is False
|
||||||
@@ -0,0 +1,162 @@
|
|||||||
|
"""Test that subagent announcements respect suppress mode."""
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from nanobot.agent.loop import AgentLoop
|
||||||
|
from nanobot.bus.events import InboundMessage, OutboundMessage
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.providers.base import LLMProvider, LLMResponse
|
||||||
|
from nanobot.session.manager import SessionManager
|
||||||
|
|
||||||
|
|
||||||
|
class MockProvider(LLMProvider):
|
||||||
|
"""Mock provider that spawns a subagent."""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.call_count = 0
|
||||||
|
self.thinking_budget = 0
|
||||||
|
|
||||||
|
def get_default_model(self) -> str:
|
||||||
|
return "mock-model"
|
||||||
|
|
||||||
|
async def chat(self, messages, tools=None, **kwargs):
|
||||||
|
self.call_count += 1
|
||||||
|
|
||||||
|
if self.call_count == 1:
|
||||||
|
# First call: spawn a subagent
|
||||||
|
return LLMResponse(
|
||||||
|
content="",
|
||||||
|
tool_calls=[
|
||||||
|
type(
|
||||||
|
"ToolCall",
|
||||||
|
(),
|
||||||
|
{
|
||||||
|
"id": "test_tool_call",
|
||||||
|
"name": "spawn",
|
||||||
|
"arguments": {"task": "Test task", "label": "test"},
|
||||||
|
},
|
||||||
|
)()
|
||||||
|
],
|
||||||
|
)
|
||||||
|
elif self.call_count == 2:
|
||||||
|
# Subagent completes its task
|
||||||
|
return LLMResponse(content="Subagent completed task")
|
||||||
|
else:
|
||||||
|
# Main agent responds to subagent announcement
|
||||||
|
return LLMResponse(content="Acknowledged subagent")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_subagent_announcement_without_suppress(tmp_path: Path):
|
||||||
|
"""Verify that WITHOUT suppress mode, announcements ARE published (baseline test)."""
|
||||||
|
|
||||||
|
bus = MessageBus()
|
||||||
|
sessions = SessionManager(workspace=tmp_path)
|
||||||
|
sessions.sessions_dir = tmp_path / "sessions"
|
||||||
|
sessions.sessions_dir.mkdir(parents=True)
|
||||||
|
|
||||||
|
provider = MockProvider()
|
||||||
|
agent = AgentLoop(
|
||||||
|
bus=bus,
|
||||||
|
provider=provider,
|
||||||
|
session_manager=sessions,
|
||||||
|
workspace=tmp_path,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Track published messages
|
||||||
|
published_messages = []
|
||||||
|
|
||||||
|
async def track_publish(msg: OutboundMessage):
|
||||||
|
published_messages.append(msg)
|
||||||
|
|
||||||
|
# Override publish to track
|
||||||
|
original_publish = bus.publish_outbound
|
||||||
|
bus.publish_outbound = track_publish
|
||||||
|
|
||||||
|
# Start agent loop
|
||||||
|
agent_task = asyncio.create_task(agent.run())
|
||||||
|
|
||||||
|
# Send test message WITHOUT suppress
|
||||||
|
test_msg = InboundMessage(
|
||||||
|
channel="test",
|
||||||
|
sender_id="user",
|
||||||
|
chat_id="normal",
|
||||||
|
content="Test message",
|
||||||
|
metadata={}, # NO suppress_output
|
||||||
|
)
|
||||||
|
await bus.publish_inbound(test_msg)
|
||||||
|
|
||||||
|
# Wait for processing (longer to allow system message to complete)
|
||||||
|
await asyncio.sleep(2.0)
|
||||||
|
|
||||||
|
# Stop agent
|
||||||
|
agent.stop()
|
||||||
|
await agent_task
|
||||||
|
|
||||||
|
# Verify: Should have published messages (NOT suppressed)
|
||||||
|
assert len(published_messages) >= 1, "Expected published messages without suppress mode"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_subagent_announcement_with_suppress(tmp_path: Path):
|
||||||
|
"""Test that subagent announcements respect suppress_output metadata."""
|
||||||
|
|
||||||
|
bus = MessageBus()
|
||||||
|
sessions = SessionManager(workspace=tmp_path)
|
||||||
|
sessions.sessions_dir = tmp_path / "sessions"
|
||||||
|
sessions.sessions_dir.mkdir(parents=True)
|
||||||
|
|
||||||
|
provider = MockProvider()
|
||||||
|
agent = AgentLoop(
|
||||||
|
bus=bus,
|
||||||
|
provider=provider,
|
||||||
|
session_manager=sessions,
|
||||||
|
workspace=tmp_path,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Track published messages
|
||||||
|
published_messages = []
|
||||||
|
|
||||||
|
async def track_publish(msg: OutboundMessage):
|
||||||
|
published_messages.append(msg)
|
||||||
|
|
||||||
|
# Override publish to track
|
||||||
|
bus.publish_outbound = track_publish
|
||||||
|
|
||||||
|
# Start agent loop
|
||||||
|
agent_task = asyncio.create_task(agent.run())
|
||||||
|
|
||||||
|
# Send test message WITH suppress
|
||||||
|
test_msg = InboundMessage(
|
||||||
|
channel="test",
|
||||||
|
sender_id="user",
|
||||||
|
chat_id="suppress",
|
||||||
|
content="Test message",
|
||||||
|
metadata={"suppress_output": True},
|
||||||
|
)
|
||||||
|
await bus.publish_inbound(test_msg)
|
||||||
|
|
||||||
|
# Wait for processing (longer to allow system message to complete)
|
||||||
|
await asyncio.sleep(2.0)
|
||||||
|
|
||||||
|
# Stop agent
|
||||||
|
agent.stop()
|
||||||
|
await agent_task
|
||||||
|
|
||||||
|
# Verify: NO messages should be published (all suppressed)
|
||||||
|
assert len(published_messages) == 0, (
|
||||||
|
f"Expected 0 published messages (all suppressed), "
|
||||||
|
f"but got {len(published_messages)}: {[m.content for m in published_messages]}"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify session contains signed [HIDDEN:*] messages (cryptographic visibility markers)
|
||||||
|
session = sessions.get_or_create("test:suppress")
|
||||||
|
hidden_messages = [m for m in session.messages if m.get("content") and "[HIDDEN:" in str(m.get("content"))]
|
||||||
|
|
||||||
|
assert len(hidden_messages) >= 1, (
|
||||||
|
f"Expected signed [HIDDEN:*] messages in session, "
|
||||||
|
f"but found {len(hidden_messages)}. Total messages: {len(session.messages)}"
|
||||||
|
)
|
||||||
@@ -0,0 +1,65 @@
|
|||||||
|
# tests/test_telegram_suppress.py
|
||||||
|
import pytest
|
||||||
|
from nanobot.bus.events import OutboundMessage
|
||||||
|
from nanobot.channels.telegram import TelegramChannel
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_suppressed_message_not_sent():
|
||||||
|
"""Test that messages with suppressed=True metadata are not sent to Telegram API."""
|
||||||
|
|
||||||
|
config = MagicMock()
|
||||||
|
config.token = "test-token"
|
||||||
|
bus = MagicMock()
|
||||||
|
|
||||||
|
channel = TelegramChannel(config, bus)
|
||||||
|
|
||||||
|
# Mock the internal _app and bot directly (skip start())
|
||||||
|
mock_app = MagicMock()
|
||||||
|
mock_bot = AsyncMock()
|
||||||
|
mock_app.bot = mock_bot
|
||||||
|
channel._app = mock_app
|
||||||
|
|
||||||
|
# Send a suppressed message
|
||||||
|
msg = OutboundMessage(
|
||||||
|
channel="telegram",
|
||||||
|
chat_id="12345",
|
||||||
|
content="[HIDDEN] This should not be sent",
|
||||||
|
metadata={"suppressed": True}
|
||||||
|
)
|
||||||
|
|
||||||
|
await channel.send(msg)
|
||||||
|
|
||||||
|
# Verify bot.send_message was NOT called
|
||||||
|
mock_bot.send_message.assert_not_called()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_normal_message_sent():
|
||||||
|
"""Test that normal messages are sent to Telegram API."""
|
||||||
|
|
||||||
|
config = MagicMock()
|
||||||
|
config.token = "test-token"
|
||||||
|
bus = MagicMock()
|
||||||
|
|
||||||
|
channel = TelegramChannel(config, bus)
|
||||||
|
|
||||||
|
# Mock the internal _app and bot directly (skip start())
|
||||||
|
mock_app = MagicMock()
|
||||||
|
mock_bot = AsyncMock()
|
||||||
|
mock_app.bot = mock_bot
|
||||||
|
channel._app = mock_app
|
||||||
|
|
||||||
|
# Send a normal message
|
||||||
|
msg = OutboundMessage(
|
||||||
|
channel="telegram",
|
||||||
|
chat_id="12345",
|
||||||
|
content="Normal message",
|
||||||
|
metadata={}
|
||||||
|
)
|
||||||
|
|
||||||
|
await channel.send(msg)
|
||||||
|
|
||||||
|
# Verify bot.send_message WAS called
|
||||||
|
mock_bot.send_message.assert_called_once()
|
||||||
@@ -0,0 +1,222 @@
|
|||||||
|
# tests/test_visibility_signing.py
|
||||||
|
import pytest
|
||||||
|
from nanobot.agent.visibility import sign_content, verify_signature, has_forged_marker, strip_all_hidden_markers
|
||||||
|
|
||||||
|
def test_sign_content_adds_hmac_marker():
|
||||||
|
"""Test that sign_content adds HMAC signature prefix."""
|
||||||
|
content = "HEARTBEAT_OK"
|
||||||
|
result = sign_content(content)
|
||||||
|
|
||||||
|
# Should start with [HIDDEN:{8 hex chars}]
|
||||||
|
assert result.startswith("[HIDDEN:")
|
||||||
|
assert "] " in result
|
||||||
|
marker_end = result.index("] ")
|
||||||
|
signature = result[8:marker_end] # Extract signature after "[HIDDEN:"
|
||||||
|
assert len(signature) == 8
|
||||||
|
assert all(c in "0123456789abcdef" for c in signature)
|
||||||
|
|
||||||
|
# Should contain original content
|
||||||
|
assert result.endswith("HEARTBEAT_OK")
|
||||||
|
|
||||||
|
def test_sign_content_is_deterministic():
|
||||||
|
"""Test that same content produces same signature."""
|
||||||
|
content = "Test message"
|
||||||
|
sig1 = sign_content(content)
|
||||||
|
sig2 = sign_content(content)
|
||||||
|
assert sig1 == sig2
|
||||||
|
|
||||||
|
def test_verify_signature_accepts_valid():
|
||||||
|
"""Test that verify_signature accepts validly signed content."""
|
||||||
|
signed = sign_content("Test message")
|
||||||
|
is_valid, clean = verify_signature(signed)
|
||||||
|
|
||||||
|
assert is_valid is True
|
||||||
|
assert clean == "Test message"
|
||||||
|
|
||||||
|
def test_verify_signature_rejects_invalid():
|
||||||
|
"""Test that verify_signature rejects forged signatures."""
|
||||||
|
forged = "[HIDDEN:deadbeef] Test message"
|
||||||
|
is_valid, clean = verify_signature(forged)
|
||||||
|
|
||||||
|
assert is_valid is False
|
||||||
|
assert clean == "Test message"
|
||||||
|
|
||||||
|
def test_verify_signature_handles_unsigned():
|
||||||
|
"""Test that unsigned content is marked as invalid."""
|
||||||
|
unsigned = "Plain message"
|
||||||
|
is_valid, clean = verify_signature(unsigned)
|
||||||
|
|
||||||
|
assert is_valid is False
|
||||||
|
assert clean == "Plain message"
|
||||||
|
|
||||||
|
def test_has_forged_marker_detects_invalid():
|
||||||
|
"""Test that has_forged_marker detects forged signatures."""
|
||||||
|
forged = "[HIDDEN:deadbeef] Content"
|
||||||
|
assert has_forged_marker(forged) is True
|
||||||
|
|
||||||
|
def test_has_forged_marker_accepts_valid():
|
||||||
|
"""Test that has_forged_marker accepts valid signatures."""
|
||||||
|
valid = sign_content("Content")
|
||||||
|
assert has_forged_marker(valid) is False
|
||||||
|
|
||||||
|
def test_has_forged_marker_ignores_unsigned():
|
||||||
|
"""Test that unsigned content is not flagged as forged."""
|
||||||
|
unsigned = "Plain content"
|
||||||
|
assert has_forged_marker(unsigned) is False
|
||||||
|
|
||||||
|
def test_strip_all_hidden_markers_removes_markers():
|
||||||
|
"""Test that strip_all_hidden_markers removes all markers."""
|
||||||
|
signed = sign_content("Message")
|
||||||
|
stripped = strip_all_hidden_markers(signed)
|
||||||
|
assert stripped == "Message"
|
||||||
|
|
||||||
|
forged = "[HIDDEN:deadbeef] Message"
|
||||||
|
stripped = strip_all_hidden_markers(forged)
|
||||||
|
assert stripped == "Message"
|
||||||
|
|
||||||
|
def test_system_prompt_includes_visibility_docs():
|
||||||
|
"""Test that system prompt documents visibility markers."""
|
||||||
|
from nanobot.agent.context import ContextBuilder
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
builder = ContextBuilder(workspace=Path("/tmp/test"))
|
||||||
|
prompt = builder.build_system_prompt()
|
||||||
|
|
||||||
|
# Should document visibility markers
|
||||||
|
assert "[HIDDEN:" in prompt
|
||||||
|
assert "cryptographically signed" in prompt.lower()
|
||||||
|
assert "do not generate" in prompt.lower() or "don't generate" in prompt.lower()
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_suppress_mode_adds_signed_marker(tmp_path):
|
||||||
|
"""Test that suppress mode adds cryptographically signed markers."""
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.session.manager import SessionManager
|
||||||
|
from nanobot.providers.base import LLMResponse
|
||||||
|
from nanobot.bus.events import InboundMessage
|
||||||
|
from unittest.mock import AsyncMock, Mock
|
||||||
|
|
||||||
|
# Setup
|
||||||
|
bus = MessageBus()
|
||||||
|
sessions = SessionManager(tmp_path)
|
||||||
|
|
||||||
|
# Mock provider
|
||||||
|
mock_provider = Mock()
|
||||||
|
mock_provider.default_model = "mock-model"
|
||||||
|
mock_provider.thinking_budget = 0
|
||||||
|
|
||||||
|
# Mock successful response
|
||||||
|
mock_response = LLMResponse(
|
||||||
|
content="Test response",
|
||||||
|
tool_calls=[],
|
||||||
|
reasoning_content=None
|
||||||
|
)
|
||||||
|
mock_provider.chat = AsyncMock(return_value=mock_response)
|
||||||
|
|
||||||
|
# Create agent loop
|
||||||
|
from nanobot.agent.loop import AgentLoop
|
||||||
|
loop = AgentLoop(
|
||||||
|
provider=mock_provider,
|
||||||
|
bus=bus,
|
||||||
|
session_manager=sessions,
|
||||||
|
workspace=tmp_path
|
||||||
|
)
|
||||||
|
|
||||||
|
# Process message with suppress_output=True
|
||||||
|
msg = InboundMessage(
|
||||||
|
channel="test",
|
||||||
|
sender_id="user",
|
||||||
|
chat_id="123",
|
||||||
|
content="Test message",
|
||||||
|
metadata={"suppress_output": True}
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await loop._process_message(msg)
|
||||||
|
|
||||||
|
# Verify response has suppressed metadata
|
||||||
|
assert response.metadata.get("suppressed") is True
|
||||||
|
|
||||||
|
# Verify session contains signed marker
|
||||||
|
session = sessions.get_or_create("test:123")
|
||||||
|
assistant_messages = [m for m in session.messages if m.get("role") == "assistant"]
|
||||||
|
assert len(assistant_messages) > 0
|
||||||
|
|
||||||
|
last_msg = assistant_messages[-1]["content"]
|
||||||
|
assert last_msg.startswith("[HIDDEN:")
|
||||||
|
|
||||||
|
# Verify signature is valid
|
||||||
|
is_valid, clean = verify_signature(last_msg)
|
||||||
|
assert is_valid is True
|
||||||
|
assert clean == "Test response"
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_forged_marker_triggers_rejection(tmp_path):
|
||||||
|
"""Test that forged markers trigger rejection and retry."""
|
||||||
|
from nanobot.bus.queue import MessageBus
|
||||||
|
from nanobot.session.manager import SessionManager
|
||||||
|
from nanobot.providers.base import LLMResponse
|
||||||
|
from nanobot.bus.events import InboundMessage
|
||||||
|
from unittest.mock import AsyncMock, Mock
|
||||||
|
|
||||||
|
# Setup
|
||||||
|
bus = MessageBus()
|
||||||
|
sessions = SessionManager(tmp_path)
|
||||||
|
|
||||||
|
# Mock provider
|
||||||
|
mock_provider = Mock()
|
||||||
|
mock_provider.default_model = "mock-model"
|
||||||
|
mock_provider.thinking_budget = 0
|
||||||
|
|
||||||
|
# First response: model tries to forge marker
|
||||||
|
forged_response = LLMResponse(
|
||||||
|
content="[HIDDEN:deadbeef] Forged message",
|
||||||
|
tool_calls=[],
|
||||||
|
reasoning_content=None
|
||||||
|
)
|
||||||
|
|
||||||
|
# Second response: clean response after correction
|
||||||
|
clean_response = LLMResponse(
|
||||||
|
content="Clean message",
|
||||||
|
tool_calls=[],
|
||||||
|
reasoning_content=None
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_provider.chat = AsyncMock(side_effect=[forged_response, clean_response])
|
||||||
|
|
||||||
|
# Create agent loop
|
||||||
|
from nanobot.agent.loop import AgentLoop
|
||||||
|
loop = AgentLoop(
|
||||||
|
provider=mock_provider,
|
||||||
|
bus=bus,
|
||||||
|
session_manager=sessions,
|
||||||
|
workspace=tmp_path
|
||||||
|
)
|
||||||
|
|
||||||
|
# Process message with suppress_output=True
|
||||||
|
msg = InboundMessage(
|
||||||
|
channel="test",
|
||||||
|
sender_id="user",
|
||||||
|
chat_id="123",
|
||||||
|
content="Test message",
|
||||||
|
metadata={"suppress_output": True}
|
||||||
|
)
|
||||||
|
|
||||||
|
response = await loop._process_message(msg)
|
||||||
|
|
||||||
|
# Verify provider.chat was called twice (initial + retry)
|
||||||
|
assert mock_provider.chat.call_count == 2
|
||||||
|
|
||||||
|
# Verify second call included correction message
|
||||||
|
second_call_messages = mock_provider.chat.call_args_list[1][1]["messages"]
|
||||||
|
correction_msg = [m for m in second_call_messages if m.get("role") == "user" and "rejected" in m.get("content", "").lower()]
|
||||||
|
assert len(correction_msg) > 0
|
||||||
|
|
||||||
|
# Verify final response uses clean content (not forged)
|
||||||
|
session = sessions.get_or_create("test:123")
|
||||||
|
assistant_messages = [m for m in session.messages if m.get("role") == "assistant"]
|
||||||
|
last_msg = assistant_messages[-1]["content"]
|
||||||
|
|
||||||
|
# Should be signed version of "Clean message", not "Forged message"
|
||||||
|
is_valid, clean = verify_signature(last_msg)
|
||||||
|
assert is_valid is True
|
||||||
|
assert clean == "Clean message"
|
||||||
Reference in New Issue
Block a user