Compare commits

..
Author SHA1 Message Date
chengyongru 59152877af fix(feishu): support video (media) download by converting type to 'file'
Feishu's GetMessageResource API only accepts 'image' or 'file' as the
type parameter. Video messages have msg_type='media', which was passed
through unchanged, causing error 234001 (Invalid request param). Now
both 'audio' and 'media' are converted to 'file' for download.
2026-04-04 01:39:53 +08:00
chengyongruandchengyongru 791282fc75 test(feishu): add unit tests for reaction add/remove and auto-cleanup 2026-04-03 22:57:11 +08:00
Jiajun Xieandchengyongru a0979d17ee feat: auto-remove reaction after message processing complete
- _add_reaction now returns reaction_id on success
- Add _remove_reaction_sync and _remove_reaction methods
- Remove reaction when stream ends to clear processing indicator
- Store reaction_id in metadata for later removal
2026-04-03 22:57:11 +08:00
FloandGitHub a6aa0b7932 feat(telegram): render tool hints as expandable blockquotes (#2752) 2026-04-03 18:27:53 +08:00
chengyongru a662ace8dd fix(memory): use parent tree in dream-restore revert to properly undo commit
The revert method was using the commit's own tree instead of its parent's,
which meant /dream-restore would restore TO that commit rather than UNDO it.
Also add root commit guard to prevent crash when reverting the initial commit.
2026-04-03 14:14:34 +08:00
Bob JohnsonandGitHub 475ea06294 fix(openai): use max_completion_tokens for OpenAI provider (#2758) 2026-04-03 10:24:03 +08:00
chengyongruandGitHub b2598270bf feat(memory): add git-backed version control for dream memory files (#2753)
- Add GitStore class wrapping dulwich for memory file versioning
- Auto-commit memory changes during Dream consolidation
- Add /dream-log and /dream-restore commands for history browsing
- Pass tracked_files as constructor param, generate .gitignore dynamically
2026-04-02 23:52:13 +08:00
JiajunandGitHub 9d07093e6d feat(provider): add Qianfan provider support (#2699) 2026-04-02 22:16:25 +08:00
chengyongru 4fd4658146 fix(memory): extract successful solutions in consolidate prompt
Add "Solutions" category to consolidate prompt so trial-and-error
workflows that reach a working approach are captured in history for
Dream to persist. Remove overly broad "debug steps" skip rule that
discarded these valuable findings.
2026-04-01 17:53:40 +08:00
chengyongru c93bf114d2 Revert "fix(feishu): handle _resuming in send_delta to prevent duplicate messages (#2667)"
This reverts commit 5d2a13db6a.
2026-04-01 17:30:46 +08:00
FloandGitHub e747d32dda fix(telegram): change drop_pending_updates to False on startup (#2686) 2026-04-01 16:47:41 +08:00
LeftXandGitHub 5d2a13db6a fix(feishu): handle _resuming in send_delta to prevent duplicate messages (#2667)
When a tool call triggers a mid-turn _stream_end(resuming=True), Feishu
channel unconditionally popped the buffer and closed the streaming card.
The next segment then created a new card, causing duplicate/fragmented
replies.

Now send_delta checks _resuming: if True, it flushes current text to
the existing card but keeps the buffer alive so subsequent segments
append to the same card. Only the final _stream_end (resuming=False)
pops the buffer and closes streaming mode.


Made-with: Cursor
2026-04-01 14:54:12 +08:00
FloandGitHub 9a468ab69b fix(tools): strip <think> blocks from message tool content (#2621) 2026-04-01 14:42:18 +08:00
FloandGitHub c41ab9e52e feat(telegram): include author context in reply tags (#2605) (#2606)
* feat(telegram): include author context in reply tags (#2605)

* fix(telegram): handle missing attributes in reply_user safely
2026-04-01 14:16:51 +08:00
FloandGitHub 2780292c9d fix(telegram): remove acknowledgment reaction when response completes (#2564) 2026-04-01 14:14:42 +08:00
FloandGitHub ebf6201b42 fix(telegram): handle RetryAfter delay internally in channel (#2552) 2026-04-01 14:13:08 +08:00
daliu858andGitHub 94e6d569b3 feat(qq): add configurable instant acknowledgment message (#2561)
Add ack_message config field to QQConfig (default: Processing...). When non-empty, sends an instant text reply before agent processing begins, filling the silence gap for users. Uses existing _send_text_only method; failure is logged but never blocks normal message handling.

Made-with: Cursor
2026-04-01 14:10:54 +08:00
FloandGitHub 973fcced2f fix(telegram): support commands with bot username suffix in groups (#2553)
* fix(telegram): support commands with bot username suffix in groups

* fix(command): preserve metadata in builtin command responses
2026-04-01 14:00:52 +08:00
Jiajun Xieandchengyongru 0d6deb9197 fix(cli): prevent spinner ANSI escape codes from being printed verbatim
Fixes #2591

The "nanobot is thinking..." spinner was printing ANSI escape codes
literally in some terminals, causing garbled output like:
  ?[2K?[32m⠧?[0m ?[2mnanobot is thinking...?[0m

Root causes:
1. Console created without force_terminal=True, so Rich couldn't
   reliably detect terminal capabilities
2. Spinner continued running during user input prompt, conflicting
   with prompt_toolkit

Changes:
- Set force_terminal=True in _make_console() for proper ANSI handling
- Add stop_for_input() method to StreamRenderer
- Call stop_for_input() before reading user input in interactive mode
- Add tests for the new functionality
2026-04-01 10:30:16 +08:00
chengyongru 5da86258cc feat(agent): two-stage memory system with Dream consolidation
Replace single-stage MemoryConsolidator with a two-stage architecture:

- Consolidator: lightweight token-budget triggered summarization,
  appends to HISTORY.md with cursor-based tracking
- Dream: cron-scheduled two-phase processor that analyzes HISTORY.md
  and updates SOUL.md, USER.md, MEMORY.md via AgentRunner with
  edit_file tools for surgical, fault-tolerant updates

New files: MemoryStore (pure file I/O), Dream class, DreamConfig,
/dream and /dream-log commands. 89 tests covering all components.
2026-04-01 09:45:24 +08:00
chengyongru cda627f956 Merge branch 'main' into nightly 2026-03-31 17:09:32 +08:00
chengyongruandchengyongru 0334fa9944 feat(provider): show cache hit rate in /status (#2645) 2026-03-31 10:57:39 +08:00
chengyongru 4741026538 Merge branch 'main' into nightly 2026-03-31 09:46:09 +08:00
VanditKumarandGitHub 7c34910fd1 feat(channel): add iMessage integration (#2539)
* feat(channel): add iMessage integration

Add iMessage as a built-in channel via Photon's iMessage platform.

Local mode (macOS) reads ~/Library/Messages/chat.db via sqlite3 with
attachment support, sends via AppleScript, and handles voice
transcription.  Remote mode connects to a Photon endpoint over pure
HTTP with proxy support, implementing send, attachments, tapback
reactions, typing indicators, mark-as-read, polls, groups, contacts,
and health checks.

Paragraph splitting sends each \n\n-separated block as a separate
iMessage bubble.  Startup seeds existing message IDs to avoid
replaying old messages on restart.

Tested end-to-end on both modes.

* fix(imessage): address review feedback

- Add retry with exponential backoff for SQLite lock contention
- Escape newlines/tabs in AppleScript to prevent parse failures
- Lower _MAX_MESSAGE_LEN to 6000 (paragraph splitting handles UX)
- Add privacy note in README for Photon remote mode
- Merge poll interval constants; make local mode configurable too
2026-03-30 15:13:23 +08:00
dd73d5d8df feat(discord): configurable read receipt + subagent working indicator (#2330)
* feat(discord): channel-side read receipt and subagent indicator

- Add 👀 reaction on message receipt, removed after bot reply
- Add 🔧 reaction on first progress message, removed on final reply
- Both managed purely in discord.py channel layer, no subagent.py changes
- Config: read_receipt_emoji, subagent_emoji with sensible defaults

Addresses maintainer feedback on HKUDS/nanobot#2330

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* fix(discord): add both reactions on inbound, not on progress

_progress flag is for streaming chunks, not subagent lifecycle.
Add 👀 + 🔧 immediately on message receipt, clear both on final reply.

* fix: remove stale _subagent_active reference in _clear_reactions

* fix(discord): clean up reactions on message handling failure

Previously, if _handle_message raised an exception, pending reactions
(read receipt + subagent indicator) would remain on the user's message
indefinitely since send() — which handles normal cleanup — would never
be called.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* refactor(discord): replace subagent_emoji with delayed working indicator

- Rename subagent_emoji → working_emoji (honest naming: not tied to
  subagent lifecycle)
- Add working_emoji_delay (default 2s) — cosmetic delay so 🔧 appears
  after 👀, cancelled if bot replies before delay fires
- Clean up: cancel pending task + remove both reactions on reply/error

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-30 14:02:43 +08:00
chengyongruandchengyongru 0d10c129e8 fix(agent): preserve inbound metadata in streaming callbacks
The on_stream and on_stream_end closures in _dispatch hardcoded
their metadata dicts, dropping channel-specific fields like
message_thread_id. Copy msg.metadata first, then add internal
streaming flags, matching the pattern already used by _bus_progress.
2026-03-30 11:58:46 +08:00
Paresh MathurchengyongruPares Mathurgemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
80403d352b feat(discord): Use discord.py for stable discord channel (#2486)
Co-authored-by: Pares Mathur <paresh.2047@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
2026-03-29 21:12:22 +08:00
5118ea824a feat(matrix): streaming support (#2447)
* Added streaming message support with incremental updates for Matrix channel

* Improve Matrix message handling and add tests

* Adjust Matrix streaming edit interval to 2 seconds

---------

Co-authored-by: natan <natan@podbielski>
2026-03-29 21:12:14 +08:00
52 changed files with 3880 additions and 2104 deletions
+79 -2
View File
@@ -261,6 +261,7 @@ Connect nanobot to your favorite chat platform. Want to build your own? See the
| **Email** | IMAP/SMTP credentials |
| **QQ** | App ID + App Secret |
| **Wecom** | Bot ID + Bot Secret |
| **iMessage** | macOS (local) or Photon server credentials (remote) |
| **Mochat** | Claw token (auto-setup available) |
<details>
@@ -831,6 +832,82 @@ nanobot gateway
</details>
<details>
<summary><b>iMessage</b></summary>
Supports two modes via [Photon](https://photon.codes):
- **Local mode**: macOS only. Reads the on-device iMessage database and sends via AppleScript. No external server needed.
- **Remote mode**: Get your endpoint and API key from [Photon](https://photon.codes) and connect from any platform. Supports tapback reactions, typing indicators, mark-as-read, attachments, and inline replies.
**Local mode (macOS)**
1. Grant **Full Disk Access** to your terminal in **System Settings → Privacy & Security → Full Disk Access**
2. Ensure iMessage is signed in and working on the Mac
```json
{
"channels": {
"imessage": {
"enabled": true,
"local": true,
"allowFrom": ["+1234567890"]
}
}
}
```
```bash
nanobot gateway
```
> Local mode supports sending/receiving text, images, and files. For reactions, typing indicators, and inline replies, use remote mode.
**Remote mode**
1. Get your **endpoint URL** and **API key** from [Photon](https://photon.codes)
2. Configure:
```json
{
"channels": {
"imessage": {
"enabled": true,
"local": false,
"serverUrl": "https://xxxxx.imsgd.photon.codes",
"apiKey": "your-api-key",
"allowFrom": ["+1234567890"]
}
}
}
```
```bash
nanobot gateway
```
> `allowFrom`: Add phone numbers or email addresses. Use `["*"]` to allow all senders.
> `groupPolicy`: `"open"` (default — respond to all messages) or `"ignore"` (skip group chats entirely).
> `proxy`: Optional HTTP proxy URL (e.g. `"http://127.0.0.1:7890"`).
> `pollInterval`: Polling interval in seconds (default `2.0`).
> **Note:** Remote mode routes messages through Photon's [advanced-imessage-http-proxy](https://github.com/photon-hq/advanced-imessage-http-proxy). Your messages and attachments transit Photon's infrastructure — the same provider that hosts your iMessage Kit server. If you need full on-device privacy, use local mode instead.
**Feature comparison:**
| Feature | Local | Remote |
|---------|-------|--------|
| Send/receive messages | ✅ | ✅ |
| Images & files | ✅ | ✅ |
| Message history | ✅ | ✅ |
| Reactions (tapbacks) | ❌ | ✅ |
| Typing indicators | ❌ | ✅ |
| Mark as read | ❌ | ✅ |
| Inline replies | ❌ | ✅ (`replyToMessage: true`) |
| Runs on any platform | ❌ | ✅ |
</details>
## 🌐 Agent Social Network
🐈 nanobot is capable of linking to the agent social network (agent community). **Just send one message and your nanobot joins automatically!**
@@ -882,6 +959,8 @@ Config file: `~/.nanobot/config.json`
| `vllm` | LLM (local, any OpenAI-compatible server) | — |
| `openai_codex` | LLM (Codex, OAuth) | `nanobot provider login openai-codex` |
| `github_copilot` | LLM (GitHub Copilot, OAuth) | `nanobot provider login github-copilot` |
| `qianfan` | LLM (Baidu Qianfan) | [cloud.baidu.com](https://cloud.baidu.com/doc/qianfan/s/Hmh4suq26) |
<details>
<summary><b>OpenAI Codex (OAuth)</b></summary>
@@ -1387,8 +1466,6 @@ MCP tools are automatically discovered and registered on startup. The LLM can us
| `tools.restrictToWorkspace` | `false` | When `true`, restricts **all** agent tools (shell, file read/write/edit, list) to the workspace directory. Prevents path traversal and out-of-scope access. |
| `tools.exec.enable` | `true` | When `false`, the shell `exec` tool is not registered at all. Use this to completely disable shell command execution. |
| `tools.exec.pathAppend` | `""` | Extra directories to append to `PATH` when running shell commands (e.g. `/usr/sbin` for `ufw`). |
| `tools.exec.commandWrapper` | `""` | Sandbox wrapper command template. See [Exec Tool Sandbox](docs/COMMAND_WRAPPER.md) for details and examples. |
| `channels.*.allowFrom` | `[]` (deny all) | Whitelist of user IDs. Empty denies all; use `["*"]` to allow everyone. |
-86
View File
@@ -1,86 +0,0 @@
# Exec Tool Sandbox (`commandWrapper`)
The `tools.exec.commandWrapper` config option wraps every shell command in a user-defined template before execution. This allows you to add a sandbox layer (e.g. bubblewrap, firejail, nsjail) without any code changes to nanobot.
## Configuration
```json
{
"tools": {
"exec": {
"commandWrapper": "<template>"
}
}
}
```
Leave empty (the default) to run commands directly with no wrapper.
## Placeholders
Two placeholders are available in the template:
| Placeholder | Value |
|---|---|
| `{command}` | The original shell command generated by the LLM |
| `{cwd}` | Absolute path of the working directory |
nanobot performs plain string replacement — it does not parse, validate, or shell-escape the values. The wrapper template is trusted configuration.
## Examples
### bubblewrap
```json
{
"tools": {
"exec": {
"commandWrapper": "bwrap --ro-bind /usr /usr --ro-bind-try /bin /bin --ro-bind-try /lib /lib --ro-bind-try /lib64 /lib64 --proc /proc --dev /dev --tmpfs /tmp --bind {cwd} {cwd} --chdir {cwd} -- sh -c \"{command}\""
}
}
}
```
Requires: `apt install bubblewrap` (or equivalent for your distro).
### firejail
```json
{
"tools": {
"exec": {
"commandWrapper": "firejail --noprofile --private={cwd} -- {command}"
}
}
}
```
### nsjail
```json
{
"tools": {
"exec": {
"commandWrapper": "nsjail -Mo --chroot /sandbox --cwd {cwd} -- {command}"
}
}
}
```
## Caveats
> [!WARNING]
> **Do not wrap `{command}` in shell quotes.** If the original command contains the same quote character, the shell will break the quoting context. For example, `sh -c '{command}'` will fail on any command that contains single quotes.
This is an inherent limitation of the template approach — nanobot substitutes `{command}` as a raw string and cannot safely shell-quote it (the command may contain compound syntax like `&&`, `|`, `;` that must be preserved for the inner shell).
### Interaction with `create_subprocess_shell`
nanobot executes the wrapped command via `create_subprocess_shell`, which adds an outer shell layer. Keep this in mind when designing your template:
- **Without `sh -c`** (e.g. `firejail ... -- {command}`): The outer shell parses `{command}` directly. Compound commands with `&&` and `|` work as expected because they are parsed by the outer shell before the sandbox tool receives them.
- **With `sh -c`** (e.g. `bwrap ... -- sh -c "{command}"`): The command is passed through two shell layers. This is only needed if the sandbox tool requires a single command argument but you want to support compound syntax.
### `restrict_to_workspace` is independent
The `tools.restrictToWorkspace` setting and `commandWrapper` are orthogonal features. The workspace restriction guards against path traversal in the original command (before wrapping). The sandbox wrapper provides OS-level isolation. You can use either or both — they address different threat models.
+2 -1
View File
@@ -3,7 +3,7 @@
from nanobot.agent.context import ContextBuilder
from nanobot.agent.hook import AgentHook, AgentHookContext, CompositeHook
from nanobot.agent.loop import AgentLoop
from nanobot.agent.memory import MemoryStore
from nanobot.agent.memory import Consolidator, Dream, MemoryStore
from nanobot.agent.skills import SkillsLoader
from nanobot.agent.subagent import SubagentManager
@@ -13,6 +13,7 @@ __all__ = [
"AgentLoop",
"CompositeHook",
"ContextBuilder",
"Dream",
"MemoryStore",
"SkillsLoader",
"SubagentManager",
+2 -2
View File
@@ -82,8 +82,8 @@ You are nanobot, a helpful AI assistant.
## Workspace
Your workspace is at: {workspace_path}
- Long-term memory: {workspace_path}/memory/MEMORY.md (write important facts here)
- History log: {workspace_path}/memory/HISTORY.md (grep-searchable). Each entry starts with [YYYY-MM-DD HH:MM].
- Long-term memory: {workspace_path}/memory/MEMORY.md (automatically managed by Dream — do not edit directly)
- History log: {workspace_path}/memory/history.jsonl (append-only JSONL, not grep-searchable).
- Custom skills: {workspace_path}/skills/{{skill-name}}/SKILL.md
{platform_policy}
+307
View File
@@ -0,0 +1,307 @@
"""Git-backed version control for memory files, using dulwich."""
from __future__ import annotations
import io
import time
from dataclasses import dataclass
from pathlib import Path
from loguru import logger
@dataclass
class CommitInfo:
sha: str # Short SHA (8 chars)
message: str
timestamp: str # Formatted datetime
def format(self, diff: str = "") -> str:
"""Format this commit for display, optionally with a diff."""
header = f"## {self.message.splitlines()[0]}\n`{self.sha}` — {self.timestamp}\n"
if diff:
return f"{header}\n```diff\n{diff}\n```"
return f"{header}\n(no file changes)"
class GitStore:
"""Git-backed version control for memory files."""
def __init__(self, workspace: Path, tracked_files: list[str]):
self._workspace = workspace
self._tracked_files = tracked_files
def is_initialized(self) -> bool:
"""Check if the git repo has been initialized."""
return (self._workspace / ".git").is_dir()
# -- init ------------------------------------------------------------------
def init(self) -> bool:
"""Initialize a git repo if not already initialized.
Creates .gitignore and makes an initial commit.
Returns True if a new repo was created, False if already exists.
"""
if self.is_initialized():
return False
try:
from dulwich import porcelain
porcelain.init(str(self._workspace))
# Write .gitignore
gitignore = self._workspace / ".gitignore"
gitignore.write_text(self._build_gitignore(), encoding="utf-8")
# Ensure tracked files exist (touch them if missing) so the initial
# commit has something to track.
for rel in self._tracked_files:
p = self._workspace / rel
p.parent.mkdir(parents=True, exist_ok=True)
if not p.exists():
p.write_text("", encoding="utf-8")
# Initial commit
porcelain.add(str(self._workspace), paths=[".gitignore"] + self._tracked_files)
porcelain.commit(
str(self._workspace),
message=b"init: nanobot memory store",
author=b"nanobot <nanobot@dream>",
committer=b"nanobot <nanobot@dream>",
)
logger.info("Git store initialized at {}", self._workspace)
return True
except Exception:
logger.warning("Git store init failed for {}", self._workspace)
return False
# -- daily operations ------------------------------------------------------
def auto_commit(self, message: str) -> str | None:
"""Stage tracked memory files and commit if there are changes.
Returns the short commit SHA, or None if nothing to commit.
"""
if not self.is_initialized():
return None
try:
from dulwich import porcelain
# .gitignore excludes everything except tracked files,
# so any staged/unstaged change must be in our files.
st = porcelain.status(str(self._workspace))
if not st.unstaged and not any(st.staged.values()):
return None
msg_bytes = message.encode("utf-8") if isinstance(message, str) else message
porcelain.add(str(self._workspace), paths=self._tracked_files)
sha_bytes = porcelain.commit(
str(self._workspace),
message=msg_bytes,
author=b"nanobot <nanobot@dream>",
committer=b"nanobot <nanobot@dream>",
)
if sha_bytes is None:
return None
sha = sha_bytes.hex()[:8]
logger.debug("Git auto-commit: {} ({})", sha, message)
return sha
except Exception:
logger.warning("Git auto-commit failed: {}", message)
return None
# -- internal helpers ------------------------------------------------------
def _resolve_sha(self, short_sha: str) -> bytes | None:
"""Resolve a short SHA prefix to the full SHA bytes."""
try:
from dulwich.repo import Repo
with Repo(str(self._workspace)) as repo:
try:
sha = repo.refs[b"HEAD"]
except KeyError:
return None
while sha:
if sha.hex().startswith(short_sha):
return sha
commit = repo[sha]
if commit.type_name != b"commit":
break
sha = commit.parents[0] if commit.parents else None
return None
except Exception:
return None
def _build_gitignore(self) -> str:
"""Generate .gitignore content from tracked files."""
dirs: set[str] = set()
for f in self._tracked_files:
parent = str(Path(f).parent)
if parent != ".":
dirs.add(parent)
lines = ["/*"]
for d in sorted(dirs):
lines.append(f"!{d}/")
for f in self._tracked_files:
lines.append(f"!{f}")
lines.append("!.gitignore")
return "\n".join(lines) + "\n"
# -- query -----------------------------------------------------------------
def log(self, max_entries: int = 20) -> list[CommitInfo]:
"""Return simplified commit log."""
if not self.is_initialized():
return []
try:
from dulwich.repo import Repo
entries: list[CommitInfo] = []
with Repo(str(self._workspace)) as repo:
try:
head = repo.refs[b"HEAD"]
except KeyError:
return []
sha = head
while sha and len(entries) < max_entries:
commit = repo[sha]
if commit.type_name != b"commit":
break
ts = time.strftime(
"%Y-%m-%d %H:%M",
time.localtime(commit.commit_time),
)
msg = commit.message.decode("utf-8", errors="replace").strip()
entries.append(CommitInfo(
sha=sha.hex()[:8],
message=msg,
timestamp=ts,
))
sha = commit.parents[0] if commit.parents else None
return entries
except Exception:
logger.warning("Git log failed")
return []
def diff_commits(self, sha1: str, sha2: str) -> str:
"""Show diff between two commits."""
if not self.is_initialized():
return ""
try:
from dulwich import porcelain
full1 = self._resolve_sha(sha1)
full2 = self._resolve_sha(sha2)
if not full1 or not full2:
return ""
out = io.BytesIO()
porcelain.diff(
str(self._workspace),
commit=full1,
commit2=full2,
outstream=out,
)
return out.getvalue().decode("utf-8", errors="replace")
except Exception:
logger.warning("Git diff_commits failed")
return ""
def find_commit(self, short_sha: str, max_entries: int = 20) -> CommitInfo | None:
"""Find a commit by short SHA prefix match."""
for c in self.log(max_entries=max_entries):
if c.sha.startswith(short_sha):
return c
return None
def show_commit_diff(self, short_sha: str, max_entries: int = 20) -> tuple[CommitInfo, str] | None:
"""Find a commit and return it with its diff vs the parent."""
commits = self.log(max_entries=max_entries)
for i, c in enumerate(commits):
if c.sha.startswith(short_sha):
if i + 1 < len(commits):
diff = self.diff_commits(commits[i + 1].sha, c.sha)
else:
diff = ""
return c, diff
return None
# -- restore ---------------------------------------------------------------
def revert(self, commit: str) -> str | None:
"""Revert (undo) the changes introduced by the given commit.
Restores all tracked memory files to the state at the commit's parent,
then creates a new commit recording the revert.
Returns the new commit SHA, or None on failure.
"""
if not self.is_initialized():
return None
try:
from dulwich.repo import Repo
full_sha = self._resolve_sha(commit)
if not full_sha:
logger.warning("Git revert: SHA not found: {}", commit)
return None
with Repo(str(self._workspace)) as repo:
commit_obj = repo[full_sha]
if commit_obj.type_name != b"commit":
return None
if not commit_obj.parents:
logger.warning("Git revert: cannot revert root commit {}", commit)
return None
# Use the parent's tree — this undoes the commit's changes
parent_obj = repo[commit_obj.parents[0]]
tree = repo[parent_obj.tree]
restored: list[str] = []
for filepath in self._tracked_files:
content = self._read_blob_from_tree(repo, tree, filepath)
if content is not None:
dest = self._workspace / filepath
dest.write_text(content, encoding="utf-8")
restored.append(filepath)
if not restored:
return None
# Commit the restored state
msg = f"revert: undo {commit}"
return self.auto_commit(msg)
except Exception:
logger.warning("Git revert failed for {}", commit)
return None
@staticmethod
def _read_blob_from_tree(repo, tree, filepath: str) -> str | None:
"""Read a blob's content from a tree object by walking path parts."""
parts = Path(filepath).parts
current = tree
for part in parts:
try:
entry = current[part.encode()]
except KeyError:
return None
obj = repo[entry[1]]
if obj.type_name == b"blob":
return obj.data.decode("utf-8", errors="replace")
if obj.type_name == b"tree":
current = obj
else:
return None
return None
+23 -12
View File
@@ -15,7 +15,7 @@ from loguru import logger
from nanobot.agent.context import ContextBuilder
from nanobot.agent.hook import AgentHook, AgentHookContext, CompositeHook
from nanobot.agent.memory import MemoryConsolidator
from nanobot.agent.memory import Consolidator, Dream
from nanobot.agent.runner import AgentRunSpec, AgentRunner
from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.cron import CronTool
@@ -97,6 +97,15 @@ class _LoopHook(AgentHook):
logger.info("Tool call: {}({})", tc.name, args_str[:200])
self._loop._set_tool_context(self._channel, self._chat_id, self._message_id)
async def after_iteration(self, context: AgentHookContext) -> None:
u = context.usage or {}
logger.debug(
"LLM usage: prompt={} completion={} cached={}",
u.get("prompt_tokens", 0),
u.get("completion_tokens", 0),
u.get("cached_tokens", 0),
)
def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None:
return self._loop._strip_think(content)
@@ -221,8 +230,8 @@ class AgentLoop:
self._concurrency_gate: asyncio.Semaphore | None = (
asyncio.Semaphore(_max) if _max > 0 else None
)
self.memory_consolidator = MemoryConsolidator(
workspace=workspace,
self.consolidator = Consolidator(
store=self.context.memory,
provider=provider,
model=self.model,
sessions=self.sessions,
@@ -231,6 +240,11 @@ class AgentLoop:
get_tool_definitions=self.tools.get_definitions,
max_completion_tokens=provider.generation.max_tokens,
)
self.dream = Dream(
store=self.context.memory,
provider=provider,
model=self.model,
)
self._register_default_tools()
self.commands = CommandRouter()
register_builtin_commands(self.commands)
@@ -248,7 +262,6 @@ class AgentLoop:
timeout=self.exec_config.timeout,
restrict_to_workspace=self.restrict_to_workspace,
path_append=self.exec_config.path_append,
command_wrapper=self.exec_config.command_wrapper,
))
self.tools.register(WebSearchTool(config=self.web_search_config, proxy=self.web_proxy))
self.tools.register(WebFetchTool(proxy=self.web_proxy))
@@ -409,8 +422,7 @@ class AgentLoop:
meta["_stream_id"] = _current_stream_id()
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content=delta,
metadata=meta,
content=delta, metadata=meta,
))
async def on_stream_end(*, resuming: bool = False) -> None:
@@ -421,8 +433,7 @@ class AgentLoop:
meta["_stream_id"] = _current_stream_id()
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content="",
metadata=meta,
content="", metadata=meta,
))
stream_segment += 1
@@ -485,7 +496,7 @@ class AgentLoop:
logger.info("Processing system message from {}", msg.sender_id)
key = f"{channel}:{chat_id}"
session = self.sessions.get_or_create(key)
await self.memory_consolidator.maybe_consolidate_by_tokens(session)
await self.consolidator.maybe_consolidate_by_tokens(session)
self._set_tool_context(channel, chat_id, msg.metadata.get("message_id"))
history = session.get_history(max_messages=0)
current_role = "assistant" if msg.sender_id == "subagent" else "user"
@@ -500,7 +511,7 @@ class AgentLoop:
)
self._save_turn(session, all_msgs, 1 + len(history))
self.sessions.save(session)
self._schedule_background(self.memory_consolidator.maybe_consolidate_by_tokens(session))
self._schedule_background(self.consolidator.maybe_consolidate_by_tokens(session))
return OutboundMessage(channel=channel, chat_id=chat_id,
content=final_content or "Background task completed.")
@@ -516,7 +527,7 @@ class AgentLoop:
if result := await self.commands.dispatch(ctx):
return result
await self.memory_consolidator.maybe_consolidate_by_tokens(session)
await self.consolidator.maybe_consolidate_by_tokens(session)
self._set_tool_context(msg.channel, msg.chat_id, msg.metadata.get("message_id"))
if message_tool := self.tools.get("message"):
@@ -553,7 +564,7 @@ class AgentLoop:
self._save_turn(session, all_msgs, 1 + len(history))
self.sessions.save(session)
self._schedule_background(self.memory_consolidator.maybe_consolidate_by_tokens(session))
self._schedule_background(self.consolidator.maybe_consolidate_by_tokens(session))
if (mt := self.tools.get("message")) and isinstance(mt, MessageTool) and mt._sent_in_turn:
return None
+411 -184
View File
@@ -1,4 +1,4 @@
"""Memory system for persistent agent memory."""
"""Memory system: pure file I/O store, lightweight Consolidator, and Dream processor."""
from __future__ import annotations
@@ -11,94 +11,189 @@ from typing import TYPE_CHECKING, Any, Callable
from loguru import logger
from nanobot.utils.helpers import ensure_dir, estimate_message_tokens, estimate_prompt_tokens_chain
from nanobot.utils.helpers import ensure_dir, estimate_message_tokens, estimate_prompt_tokens_chain, strip_think
from nanobot.agent.runner import AgentRunSpec, AgentRunner
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.agent.git_store import GitStore
if TYPE_CHECKING:
from nanobot.providers.base import LLMProvider
from nanobot.session.manager import Session, SessionManager
_SAVE_MEMORY_TOOL = [
{
"type": "function",
"function": {
"name": "save_memory",
"description": "Save the memory consolidation result to persistent storage.",
"parameters": {
"type": "object",
"properties": {
"history_entry": {
"type": "string",
"description": "A paragraph summarizing key events/decisions/topics. "
"Start with [YYYY-MM-DD HH:MM]. Include detail useful for grep search.",
},
"memory_update": {
"type": "string",
"description": "Full updated long-term memory as markdown. Include all existing "
"facts plus new ones. Return unchanged if nothing new.",
},
},
"required": ["history_entry", "memory_update"],
},
},
}
]
def _ensure_text(value: Any) -> str:
"""Normalize tool-call payload values to text for file storage."""
return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)
def _normalize_save_memory_args(args: Any) -> dict[str, Any] | None:
"""Normalize provider tool-call arguments to the expected dict shape."""
if isinstance(args, str):
args = json.loads(args)
if isinstance(args, list):
return args[0] if args and isinstance(args[0], dict) else None
return args if isinstance(args, dict) else None
_TOOL_CHOICE_ERROR_MARKERS = (
"tool_choice",
"toolchoice",
"does not support",
'should be ["none", "auto"]',
)
def _is_tool_choice_unsupported(content: str | None) -> bool:
"""Detect provider errors caused by forced tool_choice being unsupported."""
text = (content or "").lower()
return any(m in text for m in _TOOL_CHOICE_ERROR_MARKERS)
# ---------------------------------------------------------------------------
# MemoryStore — pure file I/O layer
# ---------------------------------------------------------------------------
class MemoryStore:
"""Two-layer memory: MEMORY.md (long-term facts) + HISTORY.md (grep-searchable log)."""
"""Pure file I/O for memory files: MEMORY.md, history.jsonl, SOUL.md, USER.md."""
_MAX_FAILURES_BEFORE_RAW_ARCHIVE = 3
_DEFAULT_MAX_HISTORY = 1000
def __init__(self, workspace: Path):
def __init__(self, workspace: Path, max_history_entries: int = _DEFAULT_MAX_HISTORY):
self.workspace = workspace
self.max_history_entries = max_history_entries
self.memory_dir = ensure_dir(workspace / "memory")
self.memory_file = self.memory_dir / "MEMORY.md"
self.history_file = self.memory_dir / "HISTORY.md"
self._consecutive_failures = 0
self.history_file = self.memory_dir / "history.jsonl"
self.soul_file = workspace / "SOUL.md"
self.user_file = workspace / "USER.md"
self._dream_log_file = self.memory_dir / ".dream-log.md"
self._cursor_file = self.memory_dir / ".cursor"
self._dream_cursor_file = self.memory_dir / ".dream_cursor"
self._git = GitStore(workspace, tracked_files=[
"SOUL.md", "USER.md", "memory/MEMORY.md",
])
def read_long_term(self) -> str:
if self.memory_file.exists():
return self.memory_file.read_text(encoding="utf-8")
return ""
@property
def git(self) -> GitStore:
return self._git
def write_long_term(self, content: str) -> None:
# -- generic helpers -----------------------------------------------------
@staticmethod
def read_file(path: Path) -> str:
try:
return path.read_text(encoding="utf-8")
except FileNotFoundError:
return ""
# -- MEMORY.md (long-term facts) -----------------------------------------
def read_memory(self) -> str:
return self.read_file(self.memory_file)
def write_memory(self, content: str) -> None:
self.memory_file.write_text(content, encoding="utf-8")
def append_history(self, entry: str) -> None:
with open(self.history_file, "a", encoding="utf-8") as f:
f.write(entry.rstrip() + "\n\n")
# -- SOUL.md -------------------------------------------------------------
def read_soul(self) -> str:
return self.read_file(self.soul_file)
def write_soul(self, content: str) -> None:
self.soul_file.write_text(content, encoding="utf-8")
# -- USER.md -------------------------------------------------------------
def read_user(self) -> str:
return self.read_file(self.user_file)
def write_user(self, content: str) -> None:
self.user_file.write_text(content, encoding="utf-8")
# -- context injection (used by context.py) ------------------------------
def get_memory_context(self) -> str:
long_term = self.read_long_term()
long_term = self.read_memory()
return f"## Long-term Memory\n{long_term}" if long_term else ""
# -- history.jsonl — append-only, JSONL format ---------------------------
def append_history(self, entry: str) -> int:
"""Append *entry* to history.jsonl and return its auto-incrementing cursor."""
cursor = self._next_cursor()
ts = datetime.now().strftime("%Y-%m-%d %H:%M")
record = {"cursor": cursor, "timestamp": ts, "content": strip_think(entry.rstrip()) or entry.rstrip()}
with open(self.history_file, "a", encoding="utf-8") as f:
f.write(json.dumps(record, ensure_ascii=False) + "\n")
self._cursor_file.write_text(str(cursor), encoding="utf-8")
return cursor
def _next_cursor(self) -> int:
"""Read the current cursor counter and return next value."""
if self._cursor_file.exists():
try:
return int(self._cursor_file.read_text(encoding="utf-8").strip()) + 1
except (ValueError, OSError):
pass
# Fallback: read last line's cursor from the JSONL file.
last = self._read_last_entry()
if last:
return last["cursor"] + 1
return 1
def read_unprocessed_history(self, since_cursor: int) -> list[dict[str, Any]]:
"""Return history entries with cursor > *since_cursor*."""
return [e for e in self._read_entries() if e["cursor"] > since_cursor]
def compact_history(self) -> None:
"""Drop oldest entries if the file exceeds *max_history_entries*."""
if self.max_history_entries <= 0:
return
entries = self._read_entries()
if len(entries) <= self.max_history_entries:
return
kept = entries[-self.max_history_entries:]
self._write_entries(kept)
# -- JSONL helpers -------------------------------------------------------
def _read_entries(self) -> list[dict[str, Any]]:
"""Read all entries from history.jsonl."""
entries: list[dict[str, Any]] = []
try:
with open(self.history_file, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
try:
entries.append(json.loads(line))
except json.JSONDecodeError:
continue
except FileNotFoundError:
pass
return entries
def _read_last_entry(self) -> dict[str, Any] | None:
"""Read the last entry from the JSONL file efficiently."""
try:
with open(self.history_file, "rb") as f:
f.seek(0, 2)
size = f.tell()
if size == 0:
return None
read_size = min(size, 4096)
f.seek(size - read_size)
data = f.read().decode("utf-8")
lines = [l for l in data.split("\n") if l.strip()]
if not lines:
return None
return json.loads(lines[-1])
except (FileNotFoundError, json.JSONDecodeError):
return None
def _write_entries(self, entries: list[dict[str, Any]]) -> None:
"""Overwrite history.jsonl with the given entries."""
with open(self.history_file, "w", encoding="utf-8") as f:
for entry in entries:
f.write(json.dumps(entry, ensure_ascii=False) + "\n")
# -- dream cursor --------------------------------------------------------
def get_last_dream_cursor(self) -> int:
if self._dream_cursor_file.exists():
try:
return int(self._dream_cursor_file.read_text(encoding="utf-8").strip())
except (ValueError, OSError):
pass
return 0
def set_last_dream_cursor(self, cursor: int) -> None:
self._dream_cursor_file.write_text(str(cursor), encoding="utf-8")
# -- dream log -----------------------------------------------------------
def read_dream_log(self) -> str:
return self.read_file(self._dream_log_file)
def append_dream_log(self, entry: str) -> None:
with open(self._dream_log_file, "a", encoding="utf-8") as f:
f.write(f"{entry.rstrip()}\n\n")
# -- message formatting utility ------------------------------------------
@staticmethod
def _format_messages(messages: list[dict]) -> str:
lines = []
@@ -111,107 +206,10 @@ class MemoryStore:
)
return "\n".join(lines)
async def consolidate(
self,
messages: list[dict],
provider: LLMProvider,
model: str,
) -> bool:
"""Consolidate the provided message chunk into MEMORY.md + HISTORY.md."""
if not messages:
return True
current_memory = self.read_long_term()
prompt = f"""Process this conversation and call the save_memory tool with your consolidation.
## Current Long-term Memory
{current_memory or "(empty)"}
## Conversation to Process
{self._format_messages(messages)}"""
chat_messages = [
{"role": "system", "content": "You are a memory consolidation agent. Call the save_memory tool with your consolidation of the conversation."},
{"role": "user", "content": prompt},
]
try:
forced = {"type": "function", "function": {"name": "save_memory"}}
response = await provider.chat_with_retry(
messages=chat_messages,
tools=_SAVE_MEMORY_TOOL,
model=model,
tool_choice=forced,
)
if response.finish_reason == "error" and _is_tool_choice_unsupported(
response.content
):
logger.warning("Forced tool_choice unsupported, retrying with auto")
response = await provider.chat_with_retry(
messages=chat_messages,
tools=_SAVE_MEMORY_TOOL,
model=model,
tool_choice="auto",
)
if not response.has_tool_calls:
logger.warning(
"Memory consolidation: LLM did not call save_memory "
"(finish_reason={}, content_len={}, content_preview={})",
response.finish_reason,
len(response.content or ""),
(response.content or "")[:200],
)
return self._fail_or_raw_archive(messages)
args = _normalize_save_memory_args(response.tool_calls[0].arguments)
if args is None:
logger.warning("Memory consolidation: unexpected save_memory arguments")
return self._fail_or_raw_archive(messages)
if "history_entry" not in args or "memory_update" not in args:
logger.warning("Memory consolidation: save_memory payload missing required fields")
return self._fail_or_raw_archive(messages)
entry = args["history_entry"]
update = args["memory_update"]
if entry is None or update is None:
logger.warning("Memory consolidation: save_memory payload contains null required fields")
return self._fail_or_raw_archive(messages)
entry = _ensure_text(entry).strip()
if not entry:
logger.warning("Memory consolidation: history_entry is empty after normalization")
return self._fail_or_raw_archive(messages)
self.append_history(entry)
update = _ensure_text(update)
if update != current_memory:
self.write_long_term(update)
self._consecutive_failures = 0
logger.info("Memory consolidation done for {} messages", len(messages))
return True
except Exception:
logger.exception("Memory consolidation failed")
return self._fail_or_raw_archive(messages)
def _fail_or_raw_archive(self, messages: list[dict]) -> bool:
"""Increment failure count; after threshold, raw-archive messages and return True."""
self._consecutive_failures += 1
if self._consecutive_failures < self._MAX_FAILURES_BEFORE_RAW_ARCHIVE:
return False
self._raw_archive(messages)
self._consecutive_failures = 0
return True
def _raw_archive(self, messages: list[dict]) -> None:
def raw_archive(self, messages: list[dict]) -> None:
"""Fallback: dump raw messages to HISTORY.md without LLM summarization."""
ts = datetime.now().strftime("%Y-%m-%d %H:%M")
self.append_history(
f"[{ts}] [RAW] {len(messages)} messages\n"
f"[RAW] {len(messages)} messages\n"
f"{self._format_messages(messages)}"
)
logger.warning(
@@ -219,8 +217,14 @@ class MemoryStore:
)
class MemoryConsolidator:
"""Owns consolidation policy, locking, and session offset updates."""
# ---------------------------------------------------------------------------
# Consolidator — lightweight token-budget triggered consolidation
# ---------------------------------------------------------------------------
class Consolidator:
"""Lightweight consolidation: summarizes evicted messages, appends to HISTORY.md."""
_MAX_CONSOLIDATION_ROUNDS = 5
@@ -228,7 +232,7 @@ class MemoryConsolidator:
def __init__(
self,
workspace: Path,
store: MemoryStore,
provider: LLMProvider,
model: str,
sessions: SessionManager,
@@ -237,7 +241,7 @@ class MemoryConsolidator:
get_tool_definitions: Callable[[], list[dict[str, Any]]],
max_completion_tokens: int = 4096,
):
self.store = MemoryStore(workspace)
self.store = store
self.provider = provider
self.model = model
self.sessions = sessions
@@ -245,16 +249,14 @@ class MemoryConsolidator:
self.max_completion_tokens = max_completion_tokens
self._build_messages = build_messages
self._get_tool_definitions = get_tool_definitions
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
weakref.WeakValueDictionary()
)
def get_lock(self, session_key: str) -> asyncio.Lock:
"""Return the shared consolidation lock for one session."""
return self._locks.setdefault(session_key, asyncio.Lock())
async def consolidate_messages(self, messages: list[dict[str, object]]) -> bool:
"""Archive a selected message chunk into persistent memory."""
return await self.store.consolidate(messages, self.provider, self.model)
def pick_consolidation_boundary(
self,
session: Session,
@@ -294,14 +296,50 @@ class MemoryConsolidator:
self._get_tool_definitions(),
)
async def archive_messages(self, messages: list[dict[str, object]]) -> bool:
"""Archive messages with guaranteed persistence (retries until raw-dump fallback)."""
async def archive(self, messages: list[dict]) -> bool:
"""Summarize messages via LLM and append to HISTORY.md.
Returns True on success (or degraded success), False if nothing to do.
"""
if not messages:
return False
try:
formatted = MemoryStore._format_messages(messages)
response = await self.provider.chat_with_retry(
model=self.model,
messages=[
{
"role": "system",
"content": (
"Extract key facts from this conversation. "
"Only output items matching these categories, skip everything else:\n"
"- User facts: personal info, preferences, stated opinions, habits\n"
"- Decisions: choices made, conclusions reached\n"
"- Solutions: working approaches discovered through trial and error, "
"especially non-obvious methods that succeeded after failed attempts\n"
"- Events: plans, deadlines, notable occurrences\n"
"- Preferences: communication style, tool preferences\n\n"
"Priority: user corrections and preferences > solutions > decisions > events > environment facts. "
"The most valuable memory prevents the user from having to repeat themselves.\n\n"
"Skip: code patterns derivable from source, git history, "
"or anything already captured in existing memory.\n\n"
"Output as concise bullet points, one fact per line. "
"No preamble, no commentary.\n"
"If nothing noteworthy happened, output: (nothing)"
),
},
{"role": "user", "content": formatted},
],
tools=None,
tool_choice=None,
)
summary = response.content or "[no summary]"
self.store.append_history(summary)
return True
except Exception:
logger.warning("Consolidation LLM call failed, raw-dumping to history")
self.store.raw_archive(messages)
return True
for _ in range(self.store._MAX_FAILURES_BEFORE_RAW_ARCHIVE):
if await self.consolidate_messages(messages):
return True
return True
async def maybe_consolidate_by_tokens(self, session: Session) -> None:
"""Loop: archive old messages until prompt fits within safe budget.
@@ -356,7 +394,7 @@ class MemoryConsolidator:
source,
len(chunk),
)
if not await self.consolidate_messages(chunk):
if not await self.archive(chunk):
return
session.last_consolidated = end_idx
self.sessions.save(session)
@@ -364,3 +402,192 @@ class MemoryConsolidator:
estimated, source = self.estimate_session_prompt_tokens(session)
if estimated <= 0:
return
# ---------------------------------------------------------------------------
# Dream — heavyweight cron-scheduled memory consolidation
# ---------------------------------------------------------------------------
class Dream:
"""Two-phase memory processor: analyze HISTORY.md, then edit files via AgentRunner.
Phase 1 produces an analysis summary (plain LLM call).
Phase 2 delegates to AgentRunner with read_file / edit_file tools so the
LLM can make targeted, incremental edits instead of replacing entire files.
"""
_PHASE1_SYSTEM = (
"Compare conversation history against current memory files. "
"Output one line per finding:\n"
"[FILE] atomic fact or change description\n\n"
"Files: USER (identity, preferences, habits), "
"SOUL (bot behavior, tone), "
"MEMORY (knowledge, project context, tool patterns)\n\n"
"Rules:\n"
"- Only new or conflicting information — skip duplicates and ephemera\n"
"- Prefer atomic facts: \"has a cat named Luna\" not \"discussed pet care\"\n"
"- Corrections: [USER] location is Tokyo, not Osaka\n"
"- Also capture confirmed approaches: if the user validated a non-obvious choice, note it\n\n"
"If nothing needs updating: [SKIP] no new information"
)
_PHASE2_SYSTEM = (
"Update memory files based on the analysis below.\n\n"
"## Quality standards\n"
"- Every line must carry standalone value — no filler\n"
"- Concise bullet points under clear headers\n"
"- Remove outdated or contradicted information\n\n"
"## Editing\n"
"- File contents provided below — edit directly, no read_file needed\n"
"- Batch changes to the same file into one edit_file call\n"
"- Surgical edits only — never rewrite entire files\n"
"- Do NOT overwrite correct entries — only add, update, or remove\n"
"- If nothing to update, stop without calling tools"
)
def __init__(
self,
store: MemoryStore,
provider: LLMProvider,
model: str,
max_batch_size: int = 20,
max_iterations: int = 10,
):
self.store = store
self.provider = provider
self.model = model
self.max_batch_size = max_batch_size
self.max_iterations = max_iterations
self._runner = AgentRunner(provider)
self._tools = self._build_tools()
# -- tool registry -------------------------------------------------------
def _build_tools(self) -> ToolRegistry:
"""Build a minimal tool registry for the Dream agent."""
from nanobot.agent.tools.filesystem import EditFileTool, ReadFileTool
tools = ToolRegistry()
workspace = self.store.workspace
tools.register(ReadFileTool(workspace=workspace, allowed_dir=workspace))
tools.register(EditFileTool(workspace=workspace, allowed_dir=workspace))
return tools
# -- main entry ----------------------------------------------------------
async def run(self) -> bool:
"""Process unprocessed history entries. Returns True if work was done."""
last_cursor = self.store.get_last_dream_cursor()
entries = self.store.read_unprocessed_history(since_cursor=last_cursor)
if not entries:
return False
batch = entries[: self.max_batch_size]
logger.info(
"Dream: processing {} entries (cursor {}{}), batch={}",
len(entries), last_cursor, batch[-1]["cursor"], len(batch),
)
# Build history text for LLM
history_text = "\n".join(
f"[{e['timestamp']}] {e['content']}" for e in batch
)
# Current file contents
current_memory = self.store.read_memory() or "(empty)"
current_soul = self.store.read_soul() or "(empty)"
current_user = self.store.read_user() or "(empty)"
file_context = (
f"## Current MEMORY.md\n{current_memory}\n\n"
f"## Current SOUL.md\n{current_soul}\n\n"
f"## Current USER.md\n{current_user}"
)
# Phase 1: Analyze
phase1_prompt = (
f"## Conversation History\n{history_text}\n\n{file_context}"
)
try:
phase1_response = await self.provider.chat_with_retry(
model=self.model,
messages=[
{"role": "system", "content": self._PHASE1_SYSTEM},
{"role": "user", "content": phase1_prompt},
],
tools=None,
tool_choice=None,
)
analysis = phase1_response.content or ""
logger.debug("Dream Phase 1 complete ({} chars)", len(analysis))
except Exception:
logger.exception("Dream Phase 1 failed")
return False
# Phase 2: Delegate to AgentRunner with read_file / edit_file
phase2_prompt = f"## Analysis Result\n{analysis}\n\n{file_context}"
tools = self._tools
messages: list[dict[str, Any]] = [
{"role": "system", "content": self._PHASE2_SYSTEM},
{"role": "user", "content": phase2_prompt},
]
try:
result = await self._runner.run(AgentRunSpec(
initial_messages=messages,
tools=tools,
model=self.model,
max_iterations=self.max_iterations,
fail_on_tool_error=True,
))
logger.debug(
"Dream Phase 2 complete: stop_reason={}, tool_events={}",
result.stop_reason, len(result.tool_events),
)
except Exception:
logger.exception("Dream Phase 2 failed")
result = None
# Build changelog from tool events
changelog: list[str] = []
if result and result.tool_events:
for event in result.tool_events:
if event["status"] == "ok":
changelog.append(f"{event['name']}: {event['detail']}")
# Advance cursor — always, to avoid re-processing Phase 1
new_cursor = batch[-1]["cursor"]
self.store.set_last_dream_cursor(new_cursor)
self.store.compact_history()
if result and result.stop_reason == "completed":
logger.info(
"Dream done: {} change(s), cursor advanced to {}",
len(changelog), new_cursor,
)
else:
reason = result.stop_reason if result else "exception"
logger.warning(
"Dream incomplete ({}): cursor advanced to {}",
reason, new_cursor,
)
# Write dream log
ts = datetime.now().strftime("%Y-%m-%d %H:%M")
if changelog:
log_entry = f"## {ts}\n"
for change in changelog:
log_entry += f"- {change}\n"
self.store.append_dream_log(log_entry)
else:
self.store.append_dream_log(f"## {ts}\nNo changes.\n")
# Git auto-commit (only when there are actual changes)
if changelog and self.store.git.is_initialized():
sha = self.store.git.auto_commit(f"dream: {ts}, {len(changelog)} change(s)")
if sha:
logger.info("Dream commit: {}", sha)
return True
+8 -6
View File
@@ -60,7 +60,7 @@ class AgentRunner:
messages = list(spec.initial_messages)
final_content: str | None = None
tools_used: list[str] = []
usage = {"prompt_tokens": 0, "completion_tokens": 0}
usage: dict[str, int] = {}
error: str | None = None
stop_reason = "completed"
tool_events: list[dict[str, str]] = []
@@ -92,13 +92,15 @@ class AgentRunner:
response = await self.provider.chat_with_retry(**kwargs)
raw_usage = response.usage or {}
usage = {
"prompt_tokens": int(raw_usage.get("prompt_tokens", 0) or 0),
"completion_tokens": int(raw_usage.get("completion_tokens", 0) or 0),
}
context.response = response
context.usage = usage
context.usage = raw_usage
context.tool_calls = list(response.tool_calls)
# Accumulate standard fields into result usage.
usage["prompt_tokens"] = usage.get("prompt_tokens", 0) + int(raw_usage.get("prompt_tokens", 0) or 0)
usage["completion_tokens"] = usage.get("completion_tokens", 0) + int(raw_usage.get("completion_tokens", 0) or 0)
cached = raw_usage.get("cached_tokens")
if cached:
usage["cached_tokens"] = usage.get("cached_tokens", 0) + int(cached)
if response.has_tool_calls:
if hook.wants_streaming():
-1
View File
@@ -121,7 +121,6 @@ class SubagentManager:
timeout=self.exec_config.timeout,
restrict_to_workspace=self.restrict_to_workspace,
path_append=self.exec_config.path_append,
command_wrapper=self.exec_config.command_wrapper,
))
tools.register(WebSearchTool(config=self.web_search_config, proxy=self.web_proxy))
tools.register(WebFetchTool(proxy=self.web_proxy))
+3
View File
@@ -84,6 +84,9 @@ class MessageTool(Tool):
media: list[str] | None = None,
**kwargs: Any
) -> str:
from nanobot.utils.helpers import strip_think
content = strip_think(content)
channel = channel or self._default_channel
chat_id = chat_id or self._default_chat_id
message_id = message_id or self._default_message_id
+1 -8
View File
@@ -23,11 +23,9 @@ class ExecTool(Tool):
allow_patterns: list[str] | None = None,
restrict_to_workspace: bool = False,
path_append: str = "",
command_wrapper: str = "",
):
self.timeout = timeout
self.working_dir = working_dir
self.command_wrapper = command_wrapper
self.deny_patterns = deny_patterns or [
r"\brm\s+-[rf]{1,2}\b", # rm -r, rm -rf, rm -fr
r"\bdel\s+/[fq]\b", # del /f, del /q
@@ -84,16 +82,11 @@ class ExecTool(Tool):
self, command: str, working_dir: str | None = None,
timeout: int | None = None, **kwargs: Any,
) -> str:
cwd = os.path.abspath(working_dir or self.working_dir or os.getcwd())
cwd = working_dir or self.working_dir or os.getcwd()
guard_error = self._guard_command(command, cwd)
if guard_error:
return guard_error
if self.command_wrapper:
original_command = command
command = self.command_wrapper.replace("{cwd}", cwd).replace("{command}", command)
logger.debug("command_wrapper applied: {} -> {}", original_command, command)
effective_timeout = min(timeout or self.timeout, self._MAX_TIMEOUT)
env = os.environ.copy()
+43 -7
View File
@@ -417,7 +417,7 @@ class FeishuChannel(BaseChannel):
return True
return self._is_bot_mentioned(message)
def _add_reaction_sync(self, message_id: str, emoji_type: str) -> None:
def _add_reaction_sync(self, message_id: str, emoji_type: str) -> str | None:
"""Sync helper for adding reaction (runs in thread pool)."""
from lark_oapi.api.im.v1 import CreateMessageReactionRequest, CreateMessageReactionRequestBody, Emoji
try:
@@ -433,22 +433,54 @@ class FeishuChannel(BaseChannel):
if not response.success():
logger.warning("Failed to add reaction: code={}, msg={}", response.code, response.msg)
return None
else:
logger.debug("Added {} reaction to message {}", emoji_type, message_id)
return response.data.reaction_id if response.data else None
except Exception as e:
logger.warning("Error adding reaction: {}", e)
return None
async def _add_reaction(self, message_id: str, emoji_type: str = "THUMBSUP") -> None:
async def _add_reaction(self, message_id: str, emoji_type: str = "THUMBSUP") -> str | None:
"""
Add a reaction emoji to a message (non-blocking).
Common emoji types: THUMBSUP, OK, EYES, DONE, OnIt, HEART
"""
if not self._client:
return None
loop = asyncio.get_running_loop()
return await loop.run_in_executor(None, self._add_reaction_sync, message_id, emoji_type)
def _remove_reaction_sync(self, message_id: str, reaction_id: str) -> None:
"""Sync helper for removing reaction (runs in thread pool)."""
from lark_oapi.api.im.v1 import DeleteMessageReactionRequest
try:
request = DeleteMessageReactionRequest.builder() \
.message_id(message_id) \
.reaction_id(reaction_id) \
.build()
response = self._client.im.v1.message_reaction.delete(request)
if response.success():
logger.debug("Removed reaction {} from message {}", reaction_id, message_id)
else:
logger.debug("Failed to remove reaction: code={}, msg={}", response.code, response.msg)
except Exception as e:
logger.debug("Error removing reaction: {}", e)
async def _remove_reaction(self, message_id: str, reaction_id: str) -> None:
"""
Remove a reaction emoji from a message (non-blocking).
Used to clear the "processing" indicator after bot replies.
"""
if not self._client or not reaction_id:
return
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self._add_reaction_sync, message_id, emoji_type)
await loop.run_in_executor(None, self._remove_reaction_sync, message_id, reaction_id)
# Regex to match markdown tables (header + separator + data rows)
_TABLE_RE = re.compile(
@@ -783,9 +815,9 @@ class FeishuChannel(BaseChannel):
"""Download a file/audio/media from a Feishu message by message_id and file_key."""
from lark_oapi.api.im.v1 import GetMessageResourceRequest
# Feishu API only accepts 'image' or 'file' as type parameter
# Convert 'audio' to 'file' for API compatibility
if resource_type == "audio":
# Feishu resource download API only accepts 'image' or 'file' as type.
# Both 'audio' and 'media' (video) messages use type='file' for download.
if resource_type in ("audio", "media"):
resource_type = "file"
try:
@@ -1046,6 +1078,9 @@ class FeishuChannel(BaseChannel):
# --- stream end: final update or fallback ---
if meta.get("_stream_end"):
if (message_id := meta.get("message_id")) and (reaction_id := meta.get("reaction_id")):
await self._remove_reaction(message_id, reaction_id)
buf = self._stream_bufs.pop(chat_id, None)
if not buf or not buf.text:
return
@@ -1227,7 +1262,7 @@ class FeishuChannel(BaseChannel):
return
# Add reaction
await self._add_reaction(message_id, self.config.react_emoji)
reaction_id = await self._add_reaction(message_id, self.config.react_emoji)
# Parse content
content_parts = []
@@ -1305,6 +1340,7 @@ class FeishuChannel(BaseChannel):
media=media_paths,
metadata={
"message_id": message_id,
"reaction_id": reaction_id,
"chat_type": chat_type,
"msg_type": msg_type,
"parent_id": parent_id,
+941
View File
@@ -0,0 +1,941 @@
"""iMessage channel using local macOS database or Photon advanced-imessage-http-proxy."""
from __future__ import annotations
import asyncio
import base64
import mimetypes
import platform
import sqlite3
import subprocess
from collections import OrderedDict
from pathlib import Path
from typing import Any, Literal
import httpx
from loguru import logger
from pydantic import Field
from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel
from nanobot.config.paths import get_media_dir
from nanobot.config.schema import Base
from nanobot.utils.helpers import split_message
_DEFAULT_DB_PATH = str(Path.home() / "Library" / "Messages" / "chat.db")
_DEFAULT_POLL_INTERVAL = 2.0
_AUDIO_EXTENSIONS = frozenset({".m4a", ".mp3", ".wav", ".aac", ".ogg", ".caf", ".opus"})
_MAX_MESSAGE_LEN = 6000
def _split_paragraphs(text: str) -> list[str]:
"""Split text on ``\\n\\n`` boundaries, then apply length limits to each chunk."""
parts: list[str] = []
for para in text.split("\n\n"):
stripped = para.strip()
if stripped:
parts.extend(split_message(stripped, _MAX_MESSAGE_LEN))
return parts or [text]
# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
class IMessageConfig(Base):
"""iMessage channel configuration."""
enabled: bool = False
local: bool = True
server_url: str = ""
api_key: str = ""
proxy: str | None = None
poll_interval: float = _DEFAULT_POLL_INTERVAL
database_path: str = _DEFAULT_DB_PATH
allow_from: list[str] = Field(default_factory=list)
group_policy: Literal["open", "ignore"] = "open"
reply_to_message: bool = False
react_tapback: str = "love"
done_tapback: str = ""
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
_PHOTON_PROXY_URL = "https://imessage-swagger.photon.codes"
_PHOTON_KIT_PATTERN = ".imsgd.photon.codes"
def _is_photon_kit_url(url: str) -> bool:
"""Return True when *url* points to a Photon iMessage Kit server (upstream)."""
from urllib.parse import urlparse
return _PHOTON_KIT_PATTERN in (urlparse(url).hostname or "")
def _make_bearer_token(server_url: str, api_key: str) -> str:
"""Build the Bearer token expected by advanced-imessage-http-proxy.
If the key already decodes to a ``url|key`` pair it is used as-is.
"""
try:
decoded = base64.b64decode(api_key, validate=True).decode()
if "|" in decoded:
return api_key
except Exception as e:
logger.debug("API key not pre-encoded, will encode: {}", type(e).__name__)
raw = f"{server_url}|{api_key}"
return base64.b64encode(raw.encode()).decode()
def _resolve_proxy_url(server_url: str) -> str:
"""Return the actual HTTP proxy base URL to use for API calls.
When the user provides a Photon iMessage Kit server URL (e.g.
``https://xxxxx.imsgd.photon.codes``), requests must go through
the shared ``advanced-imessage-http-proxy`` at a fixed endpoint.
The original Kit URL is only used inside the Bearer token.
"""
if _is_photon_kit_url(server_url):
logger.info(
"Photon Kit URL detected — routing through proxy at {} "
"(hosted by Photon, the same provider as your iMessage Kit server).",
_PHOTON_PROXY_URL,
)
return _PHOTON_PROXY_URL
return server_url
def _extract_address(chat_id: str) -> str:
"""Convert a chatGuid to the proxy's address format.
``iMessage;-;+1234567890`` → ``+1234567890``
``iMessage;+;chat123`` → ``group:chat123``
``+1234567890`` → ``+1234567890`` (passthrough)
"""
if ";-;" in chat_id:
return chat_id.split(";-;", 1)[1]
if ";+;" in chat_id:
return "group:" + chat_id.split(";+;", 1)[1]
return chat_id
# ---------------------------------------------------------------------------
# Channel
# ---------------------------------------------------------------------------
class IMessageChannel(BaseChannel):
"""iMessage channel with local (macOS) and remote (Photon) modes.
Local mode reads from the native iMessage SQLite database and sends
via AppleScript — pure Python, no external dependencies.
Remote mode talks to a Photon ``advanced-imessage-http-proxy`` server.
See https://github.com/photon-hq/advanced-imessage-http-proxy for the
full API reference.
"""
name = "imessage"
display_name = "iMessage"
@classmethod
def default_config(cls) -> dict[str, Any]:
return IMessageConfig().model_dump(by_alias=True)
def __init__(self, config: Any, bus: MessageBus):
if isinstance(config, dict):
config = IMessageConfig.model_validate(config)
super().__init__(config, bus)
self.config: IMessageConfig = config
self._processed_ids: OrderedDict[str, None] = OrderedDict()
self._http: httpx.AsyncClient | None = None
self._last_rowid: int = 0
# ---- lifecycle ---------------------------------------------------------
async def start(self) -> None:
if self.config.local:
await self._start_local()
else:
await self._start_remote()
async def stop(self) -> None:
self._running = False
if self._http:
await self._http.aclose()
self._http = None
async def send(self, msg: OutboundMessage) -> None:
if self.config.local:
await self._send_local(msg)
else:
await self._send_remote(msg)
# ======================================================================
# LOCAL MODE (macOS — sqlite3 + AppleScript)
# ======================================================================
async def _start_local(self) -> None:
if platform.system() != "Darwin":
logger.error("iMessage local mode requires macOS")
return
db_path = self.config.database_path
if not Path(db_path).exists():
logger.error(
"iMessage database not found at {}. "
"Ensure Full Disk Access is granted to your terminal.",
db_path,
)
return
self._running = True
self._last_rowid = self._get_max_rowid(db_path)
logger.info("iMessage local watcher started (polling {})", db_path)
while self._running:
try:
await self._poll_local_db(db_path)
except Exception as e:
logger.warning("iMessage local poll error: {}", e)
await asyncio.sleep(max(0.5, self.config.poll_interval))
def _get_max_rowid(self, db_path: str) -> int:
try:
with sqlite3.connect(db_path, uri=True) as conn:
cur = conn.execute("SELECT MAX(ROWID) FROM message")
row = cur.fetchone()
return row[0] or 0
except Exception:
return 0
async def _poll_local_db(self, db_path: str) -> None:
loop = asyncio.get_running_loop()
rows = await loop.run_in_executor(None, self._fetch_new_messages, db_path)
for row in rows:
await self._handle_local_message(row)
self._last_rowid = max(self._last_rowid, int(row["ROWID"]))
def _fetch_new_messages(self, db_path: str) -> list[dict[str, Any]]:
for attempt in range(3):
try:
return self._query_new_messages(db_path)
except sqlite3.OperationalError as e:
if "locked" in str(e) and attempt < 2:
import time
time.sleep(0.5 * (attempt + 1))
continue
raise
return []
def _query_new_messages(self, db_path: str) -> list[dict[str, Any]]:
with sqlite3.connect(db_path, uri=True, timeout=10) as conn:
conn.row_factory = sqlite3.Row
cur = conn.execute(
"""
SELECT
m.ROWID,
m.guid,
m.text,
m.is_from_me,
m.date AS msg_date,
m.service,
h.id AS sender,
c.chat_identifier,
c.style AS chat_style,
a.ROWID AS att_rowid,
a.filename AS att_filename,
a.mime_type AS att_mime,
a.transfer_name AS att_transfer_name
FROM message m
LEFT JOIN handle h ON m.handle_id = h.ROWID
LEFT JOIN chat_message_join cmj ON m.ROWID = cmj.message_id
LEFT JOIN chat c ON cmj.chat_id = c.ROWID
LEFT JOIN message_attachment_join maj ON m.ROWID = maj.message_id
LEFT JOIN attachment a ON maj.attachment_id = a.ROWID
WHERE m.ROWID > ?
ORDER BY m.ROWID ASC
""",
(self._last_rowid,),
)
msg_map: dict[int, dict[str, Any]] = {}
for row in cur:
d = dict(row)
rowid = d["ROWID"]
if rowid not in msg_map:
msg_map[rowid] = {**d, "attachments": []}
if d.get("att_rowid"):
raw_path = d.get("att_filename") or ""
resolved = (
raw_path.replace("~", str(Path.home()), 1)
if raw_path.startswith("~")
else raw_path
)
msg_map[rowid]["attachments"].append(
{
"filename": resolved,
"mime_type": d.get("att_mime") or "",
"transfer_name": d.get("att_transfer_name") or "",
}
)
return list(msg_map.values())
async def _handle_local_message(self, row: dict[str, Any]) -> None:
if row.get("is_from_me"):
return
message_id = row.get("guid", "")
if self._is_seen(message_id):
return
sender = row.get("sender") or ""
chat_id = row.get("chat_identifier") or sender
content = row.get("text") or ""
is_group = (row.get("chat_style") or 0) == 43
if is_group and self.config.group_policy == "ignore":
return
media_paths: list[str] = []
for att in row.get("attachments") or []:
file_path = att.get("filename", "")
if not file_path or not Path(file_path).exists():
continue
mime = att.get("mime_type") or ""
ext = Path(file_path).suffix.lower()
if ext in _AUDIO_EXTENSIONS or mime.startswith("audio/"):
transcription = await self.transcribe_audio(file_path)
if transcription:
voice_tag = f"[Voice Message: {transcription}]"
content = f"{content}\n{voice_tag}" if content else voice_tag
continue
media_paths.append(file_path)
tag = "image" if mime.startswith("image/") else "file"
media_tag = f"[{tag}: {file_path}]"
content = f"{content}\n{media_tag}" if content else media_tag
await self._handle_message(
sender_id=sender,
chat_id=chat_id,
content=content,
media=media_paths,
metadata={
"message_id": message_id,
"service": row.get("service", "iMessage"),
"is_group": is_group,
"source": "local",
},
)
self._mark_seen(message_id)
async def _send_local(self, msg: OutboundMessage) -> None:
if (msg.metadata or {}).get("_progress"):
return
recipient = msg.chat_id
if msg.content:
for chunk in _split_paragraphs(msg.content):
await self._applescript_send_text(recipient, chunk)
for media_path in msg.media or []:
await self._applescript_send_file(recipient, media_path)
@staticmethod
def _escape_applescript(s: str) -> str:
return (
s.replace("\\", "\\\\")
.replace('"', '\\"')
.replace("\n", "\\n")
.replace("\r", "\\r")
.replace("\t", "\\t")
)
async def _applescript_send_text(self, recipient: str, text: str) -> None:
escaped_recipient = self._escape_applescript(recipient)
escaped_text = self._escape_applescript(text)
script = (
f'tell application "Messages"\n'
f" set targetService to 1st account whose service type = iMessage\n"
f' set targetBuddy to participant "{escaped_recipient}" of targetService\n'
f' send "{escaped_text}" to targetBuddy\n'
f"end tell"
)
await self._run_osascript(script)
async def _applescript_send_file(self, recipient: str, file_path: str) -> None:
escaped_recipient = self._escape_applescript(recipient)
escaped_path = self._escape_applescript(file_path)
script = (
f'tell application "Messages"\n'
f" set targetService to 1st account whose service type = iMessage\n"
f' set targetBuddy to participant "{escaped_recipient}" of targetService\n'
f' send POSIX file "{escaped_path}" to targetBuddy\n'
f"end tell"
)
await self._run_osascript(script)
async def _run_osascript(self, script: str) -> None:
loop = asyncio.get_running_loop()
try:
await loop.run_in_executor(
None,
lambda: subprocess.run(
["osascript", "-e", script],
check=True,
capture_output=True,
timeout=15,
),
)
except subprocess.CalledProcessError as e:
logger.error(
"AppleScript send failed: {}", e.stderr.decode()[:200] if e.stderr else str(e)
)
raise
except subprocess.TimeoutExpired:
logger.error("AppleScript send timed out")
raise
# ======================================================================
# REMOTE MODE (advanced-imessage-http-proxy)
# https://github.com/photon-hq/advanced-imessage-http-proxy
# ======================================================================
def _build_http_client(self) -> httpx.AsyncClient:
token = _make_bearer_token(self.config.server_url, self.config.api_key)
base_url = _resolve_proxy_url(self.config.server_url)
return httpx.AsyncClient(
base_url=base_url.rstrip("/"),
headers={"Authorization": f"Bearer {token}"},
proxy=self.config.proxy or None,
timeout=30.0,
)
async def _start_remote(self) -> None:
if not self.config.server_url:
logger.error("iMessage remote mode requires serverUrl")
return
if not self.config.api_key:
logger.error("iMessage remote mode requires apiKey")
return
self._running = True
self._http = self._build_http_client()
if not await self._api_health():
logger.error("iMessage server health check failed — will retry in poll loop")
await self._seed_existing_message_ids()
poll_interval = max(0.5, self.config.poll_interval)
logger.info(
"iMessage remote polling started ({}s interval, proxy={})",
poll_interval,
self.config.proxy or "none",
)
while self._running:
try:
await self._poll_remote()
except Exception as e:
logger.warning("iMessage remote poll error: {}", e)
await asyncio.sleep(poll_interval)
async def _seed_existing_message_ids(self) -> None:
"""Mark all existing messages as seen so we only process new ones after startup."""
try:
resp = await self._api_get_messages(limit=100)
if resp and isinstance(resp, list):
for msg in resp:
msg_id = msg.get("id") or msg.get("guid", "")
if msg_id:
self._mark_seen(msg_id)
logger.info("Seeded {} existing message IDs", len(self._processed_ids))
except Exception as e:
logger.debug("Could not seed existing message IDs: {}", e)
# ---- inbound -----------------------------------------------------------
async def _poll_remote(self) -> None:
if not self._http:
return
messages = await self._api_get_messages(limit=50)
if not messages:
return
for msg in reversed(messages):
await self._handle_remote_message(msg)
async def _handle_remote_message(self, data: dict[str, Any]) -> None:
sender_raw = data.get("from") or ""
if sender_raw == "me" or data.get("isFromMe"):
return
message_id = data.get("id") or data.get("guid", "")
if self._is_seen(message_id):
return
self._mark_seen(message_id)
sender = sender_raw
if not sender:
handle = data.get("handle")
if isinstance(handle, dict):
sender = handle.get("address", "")
address = data.get("chat") or sender
if not address:
chats = data.get("chats") or []
chat_guid = chats[0].get("guid", "") if chats else ""
address = _extract_address(chat_guid) if chat_guid else sender
content = data.get("text") or ""
is_group = address.startswith("group:") or (";+;" in address)
if is_group and self.config.group_policy == "ignore":
return
await self._api_react(address, message_id, self.config.react_tapback)
await self._api_mark_read(address)
media_paths: list[str] = []
for att in data.get("attachments") or []:
att_guid = att.get("guid", "")
name = att.get("transferName") or att.get("filename") or ""
if not att_guid or not self._http:
continue
local_path = await self._api_download_attachment(att_guid, name)
if not local_path:
continue
mime, _ = mimetypes.guess_type(local_path)
ext = Path(local_path).suffix.lower()
if ext in _AUDIO_EXTENSIONS or (mime and mime.startswith("audio/")):
transcription = await self.transcribe_audio(local_path)
if transcription:
voice_tag = f"[Voice Message: {transcription}]"
content = f"{content}\n{voice_tag}" if content else voice_tag
continue
media_paths.append(local_path)
tag = "image" if mime and mime.startswith("image/") else "file"
media_tag = f"[{tag}: {local_path}]"
content = f"{content}\n{media_tag}" if content else media_tag
await self._handle_message(
sender_id=sender,
chat_id=address,
content=content,
media=media_paths,
metadata={
"message_id": message_id,
"is_group": is_group,
"source": "remote",
"timestamp": data.get("sentAt") or data.get("dateCreated"),
},
)
# ---- outbound ----------------------------------------------------------
async def _send_remote(self, msg: OutboundMessage) -> None:
if not self._http:
raise RuntimeError("iMessage remote HTTP client not initialised")
meta = msg.metadata or {}
if meta.get("_progress"):
return
to = msg.chat_id
await self._api_start_typing(to)
try:
if msg.content:
chunks = _split_paragraphs(msg.content)
for i, chunk in enumerate(chunks):
body: dict[str, Any] = {"to": to, "text": chunk, "service": "iMessage"}
if i == 0 and self.config.reply_to_message and msg.reply_to:
body["replyTo"] = msg.reply_to
if await self._api_send(body) is None:
raise RuntimeError(f"iMessage text delivery failed for {to}")
for media_path in msg.media or []:
if await self._api_send_file(to, media_path) is None:
raise RuntimeError(f"iMessage media delivery failed: {media_path}")
finally:
await self._api_stop_typing(to)
message_id = meta.get("message_id")
if message_id and self.config.react_tapback:
await self._api_remove_react(to, message_id, self.config.react_tapback)
if self.config.done_tapback:
await self._api_react(to, message_id, self.config.done_tapback)
# ======================================================================
# PROXY API METHODS
# https://github.com/photon-hq/advanced-imessage-http-proxy
# ======================================================================
# ---- messaging ---------------------------------------------------------
async def _api_send(self, body: dict[str, Any]) -> dict[str, Any] | None:
"""``POST /send`` — send text message with optional effect / reply."""
return await self._post("/send", body)
async def _api_send_file(
self, to: str, file_path: str, audio: bool | None = None
) -> dict[str, Any] | None:
"""``POST /send/file`` — send attachment (image, file, audio message)."""
if not self._http:
return None
if not Path(file_path).exists():
logger.warning("iMessage attachment not found: {}", file_path)
return None
mime, _ = mimetypes.guess_type(file_path)
ext = Path(file_path).suffix.lower()
is_audio = (
audio
if audio is not None
else (ext in _AUDIO_EXTENSIONS or (mime or "").startswith("audio/"))
)
data: dict[str, str] = {"to": to}
if is_audio:
data["audio"] = "true"
with open(file_path, "rb") as f:
resp = await self._http.post(
"/send/file",
data=data,
files={"file": (Path(file_path).name, f, mime or "application/octet-stream")},
)
return self._unwrap(resp)
async def _api_send_sticker(
self,
to: str,
file_path: str,
reply_to: str | None = None,
**kwargs: Any,
) -> dict[str, Any] | None:
"""``POST /send/sticker`` — send standalone or reply sticker."""
if not self._http:
return None
data: dict[str, str] = {"to": to}
if reply_to:
data["replyTo"] = reply_to
for k in ("stickerX", "stickerY", "stickerScale", "stickerRotation", "stickerWidth"):
if k in kwargs:
data[k] = str(kwargs[k])
with open(file_path, "rb") as f:
resp = await self._http.post(
"/send/sticker",
data=data,
files={"file": (Path(file_path).name, f, "image/png")},
)
return self._unwrap(resp)
async def _api_unsend(self, message_id: str) -> dict[str, Any] | None:
"""``DELETE /messages/:id`` — retract a sent message."""
return await self._delete(f"/messages/{message_id}")
# ---- reactions ---------------------------------------------------------
async def _api_react(self, chat: str, message_id: str, tapback: str) -> None:
"""``POST /messages/:id/react`` — add tapback (best-effort)."""
if not self._http or not tapback:
return
try:
await self._http.post(
f"/messages/{message_id}/react",
json={"chat": chat, "type": tapback},
)
except Exception as e:
logger.debug("iMessage tapback failed: {}", e)
async def _api_remove_react(self, chat: str, message_id: str, tapback: str) -> None:
"""``DELETE /messages/:id/react`` — remove tapback (best-effort)."""
if not self._http or not tapback:
return
try:
await self._http.request(
"DELETE",
f"/messages/{message_id}/react",
json={"chat": chat, "type": tapback},
)
except Exception as e:
logger.debug("iMessage remove tapback failed: {}", e)
# ---- messages ----------------------------------------------------------
async def _api_get_messages(
self,
limit: int = 50,
chat: str | None = None,
) -> list[dict[str, Any]]:
"""``GET /messages`` — query messages."""
params: dict[str, Any] = {"limit": limit}
if chat:
params["chat"] = chat
data = await self._get("/messages", params=params)
return data if isinstance(data, list) else []
async def _api_search_messages(
self, query: str, chat: str | None = None
) -> list[dict[str, Any]]:
"""``GET /messages/search`` — search messages by text."""
params: dict[str, Any] = {"q": query}
if chat:
params["chat"] = chat
data = await self._get("/messages/search", params=params)
return data if isinstance(data, list) else []
async def _api_get_message(self, message_id: str) -> dict[str, Any] | None:
"""``GET /messages/:id`` — get single message details."""
data = await self._get(f"/messages/{message_id}")
return data if isinstance(data, dict) else None
# ---- chats -------------------------------------------------------------
async def _api_get_chats(self) -> list[dict[str, Any]]:
"""``GET /chats`` — list all conversations."""
data = await self._get("/chats")
return data if isinstance(data, list) else []
async def _api_get_chat(self, address: str) -> dict[str, Any] | None:
"""``GET /chats/:id`` — get chat details."""
data = await self._get(f"/chats/{address}")
return data if isinstance(data, dict) else None
async def _api_get_chat_messages(
self,
address: str,
limit: int = 50,
) -> list[dict[str, Any]]:
"""``GET /chats/:id/messages`` — get message history for a chat."""
data = await self._get(f"/chats/{address}/messages", params={"limit": limit})
return data if isinstance(data, list) else []
async def _api_get_chat_participants(self, address: str) -> list[dict[str, Any]]:
"""``GET /chats/:id/participants`` — get group participants."""
data = await self._get(f"/chats/{address}/participants")
return data if isinstance(data, list) else []
async def _api_mark_read(self, address: str) -> None:
"""``POST /chats/:id/read`` — clear unread badge."""
if not self._http:
return
try:
await self._http.post(f"/chats/{address}/read")
except Exception as e:
logger.debug("iMessage mark-read failed: {}", e)
async def _api_start_typing(self, address: str) -> None:
"""``POST /chats/:id/typing`` — show typing indicator."""
if not self._http:
return
try:
await self._http.post(f"/chats/{address}/typing")
except Exception as e:
logger.debug("iMessage typing start failed: {}", e)
async def _api_stop_typing(self, address: str) -> None:
"""``DELETE /chats/:id/typing`` — stop typing indicator."""
if not self._http:
return
try:
await self._http.request("DELETE", f"/chats/{address}/typing")
except Exception as e:
logger.debug("iMessage typing stop failed: {}", e)
# ---- groups ------------------------------------------------------------
async def _api_create_group(
self, members: list[str], name: str | None = None
) -> dict[str, Any] | None:
"""``POST /groups`` — create a group chat."""
body: dict[str, Any] = {"members": members}
if name:
body["name"] = name
return await self._post("/groups", body)
async def _api_update_group(self, group_id: str, name: str) -> dict[str, Any] | None:
"""``PATCH /groups/:id`` — rename a group."""
if not self._http:
return None
resp = await self._http.patch(f"/groups/{group_id}", json={"name": name})
return self._unwrap(resp)
# ---- polls -------------------------------------------------------------
async def _api_create_poll(
self,
to: str,
question: str,
options: list[str],
) -> dict[str, Any] | None:
"""``POST /polls`` — create an interactive poll."""
return await self._post("/polls", {"to": to, "question": question, "options": options})
async def _api_get_poll(self, poll_id: str) -> dict[str, Any] | None:
"""``GET /polls/:id`` — get poll details."""
data = await self._get(f"/polls/{poll_id}")
return data if isinstance(data, dict) else None
async def _api_vote_poll(
self, poll_id: str, chat: str, option_id: str
) -> dict[str, Any] | None:
"""``POST /polls/:id/vote`` — vote on a poll option."""
return await self._post(f"/polls/{poll_id}/vote", {"chat": chat, "optionId": option_id})
async def _api_unvote_poll(
self, poll_id: str, chat: str, option_id: str
) -> dict[str, Any] | None:
"""``POST /polls/:id/unvote`` — remove vote from poll."""
return await self._post(f"/polls/{poll_id}/unvote", {"chat": chat, "optionId": option_id})
async def _api_add_poll_option(
self, poll_id: str, chat: str, text: str
) -> dict[str, Any] | None:
"""``POST /polls/:id/options`` — add option to existing poll."""
return await self._post(f"/polls/{poll_id}/options", {"chat": chat, "text": text})
# ---- attachments -------------------------------------------------------
async def _api_download_attachment(self, att_guid: str, filename: str) -> str | None:
"""``GET /attachments/:id`` — download to local media dir."""
if not self._http:
return None
try:
resp = await self._http.get(f"/attachments/{att_guid}")
if not resp.is_success:
return None
media_dir = get_media_dir("imessage")
sanitized_guid = att_guid.replace("/", "_").replace("\\", "_").replace("\x00", "")
safe_name = Path(filename).name.replace("\x00", "") if filename else ""
if not safe_name:
safe_name = f"{sanitized_guid}.bin"
dest = (media_dir / safe_name).resolve()
if not dest.is_relative_to(media_dir.resolve()):
dest = (media_dir / f"{sanitized_guid}.bin").resolve()
dest.write_bytes(resp.content)
return str(dest)
except Exception as e:
logger.warning("Failed to download iMessage attachment {}: {}", att_guid, e)
return None
async def _api_attachment_info(self, att_guid: str) -> dict[str, Any] | None:
"""``GET /attachments/:id/info`` — get attachment metadata."""
data = await self._get(f"/attachments/{att_guid}/info")
return data if isinstance(data, dict) else None
# ---- contacts & handles ------------------------------------------------
async def _api_check_imessage(self, address: str) -> bool:
"""``GET /check/:address`` — check if address uses iMessage."""
data = await self._get(f"/check/{address}")
if isinstance(data, dict):
return bool(data.get("available") or data.get("imessage"))
return False
async def _api_get_contacts(self) -> list[dict[str, Any]]:
"""``GET /contacts`` — list device contacts."""
data = await self._get("/contacts")
return data if isinstance(data, list) else []
async def _api_get_handles(self) -> list[dict[str, Any]]:
"""``GET /handles`` — list known handles."""
data = await self._get("/handles")
return data if isinstance(data, list) else []
# ---- server ------------------------------------------------------------
async def _api_server_info(self) -> dict[str, Any] | None:
"""``GET /server`` — get server info."""
data = await self._get("/server")
return data if isinstance(data, dict) else None
async def _api_health(self) -> bool:
"""``GET /health`` — basic health check."""
if not self._http:
return False
try:
resp = await self._http.get("/health")
if resp.is_success:
logger.info("iMessage server health check passed")
return True
logger.warning("iMessage health check HTTP {}", resp.status_code)
except Exception as e:
logger.warning("iMessage health check failed: {}", e)
return False
# ---- HTTP helpers ------------------------------------------------------
async def _get(self, path: str, params: dict[str, Any] | None = None) -> Any:
if not self._http:
return None
try:
resp = await self._http.get(path, params=params)
return self._unwrap(resp)
except Exception as e:
logger.warning("iMessage GET {} failed: {}", path, e)
return None
async def _post(self, path: str, body: dict[str, Any]) -> dict[str, Any] | None:
if not self._http:
return None
try:
resp = await self._http.post(path, json=body)
if not resp.is_success:
logger.warning(
"iMessage POST {} HTTP {}: {}", path, resp.status_code, resp.text[:200]
)
return self._unwrap(resp)
except Exception as e:
logger.warning("iMessage POST {} failed: {}", path, e)
raise
async def _delete(self, path: str, body: dict[str, Any] | None = None) -> dict[str, Any] | None:
if not self._http:
return None
try:
resp = await self._http.request("DELETE", path, json=body)
return self._unwrap(resp)
except Exception as e:
logger.warning("iMessage DELETE {} failed: {}", path, e)
return None
@staticmethod
def _unwrap(resp: httpx.Response) -> Any:
"""Unwrap the proxy's ``{"ok": true, "data": ...}`` envelope."""
if not resp.is_success:
return None
try:
body = resp.json()
except Exception:
logger.debug("iMessage server returned non-JSON response")
return None
if isinstance(body, dict) and "data" in body:
return body["data"]
return body
# ---- dedup helper ------------------------------------------------------
def _is_seen(self, message_id: str) -> bool:
if not message_id:
return False
return message_id in self._processed_ids
def _mark_seen(self, message_id: str) -> None:
if not message_id:
return
self._processed_ids[message_id] = None
while len(self._processed_ids) > 1000:
self._processed_ids.popitem(last=False)
+9 -26
View File
@@ -132,11 +132,7 @@ def _render_markdown_html(text: str) -> str | None:
return formatted
def _build_matrix_text_content(
text: str,
event_id: str | None = None,
thread_relates_to: dict[str, object] | None = None,
) -> dict[str, object]:
def _build_matrix_text_content(text: str, event_id: str | None = None) -> dict[str, object]:
"""
Constructs and returns a dictionary representing the matrix text content with optional
HTML formatting and reference to an existing event for replacement. This function is
@@ -148,9 +144,6 @@ def _build_matrix_text_content(
include information indicating that the message is a replacement of the specified
event.
:type event_id: str | None
:param thread_relates_to: Optional Matrix thread relation metadata. For edits this is
stored in ``m.new_content`` so the replacement remains in the same thread.
:type thread_relates_to: dict[str, object] | None
:return: A dictionary containing the matrix text content, potentially enriched with
HTML formatting and replacement metadata if applicable.
:rtype: dict[str, object]
@@ -160,18 +153,14 @@ def _build_matrix_text_content(
content["format"] = MATRIX_HTML_FORMAT
content["formatted_body"] = html
if event_id:
content["m.new_content"] = {
content["m.new_content"] = {
"body": text,
"msgtype": "m.text",
"msgtype": "m.text"
}
content["m.relates_to"] = {
"rel_type": "m.replace",
"event_id": event_id,
"event_id": event_id
}
if thread_relates_to:
content["m.new_content"]["m.relates_to"] = thread_relates_to
elif thread_relates_to:
content["m.relates_to"] = thread_relates_to
return content
@@ -486,11 +475,9 @@ class MatrixChannel(BaseChannel):
await self._stop_typing_keepalive(chat_id, clear_typing=True)
content = _build_matrix_text_content(
buf.text,
buf.event_id,
thread_relates_to=relates_to,
)
content = _build_matrix_text_content(buf.text, buf.event_id)
if relates_to:
content["m.relates_to"] = relates_to
await self._send_room_content(chat_id, content)
return
@@ -507,18 +494,14 @@ class MatrixChannel(BaseChannel):
if not buf.last_edit or (now - buf.last_edit) >= self._STREAM_EDIT_INTERVAL:
try:
content = _build_matrix_text_content(
buf.text,
buf.event_id,
thread_relates_to=relates_to,
)
content = _build_matrix_text_content(buf.text, buf.event_id)
response = await self._send_room_content(chat_id, content)
buf.last_edit = now
if not buf.event_id:
# we are editing the same message all the time, so only the first time the event id needs to be set
buf.event_id = response.event_id
except Exception:
await self._stop_typing_keepalive(chat_id, clear_typing=True)
await self._stop_typing_keepalive(metadata["room_id"], clear_typing=True)
pass
+12
View File
@@ -134,6 +134,7 @@ class QQConfig(Base):
secret: str = ""
allow_from: list[str] = Field(default_factory=list)
msg_format: Literal["plain", "markdown"] = "plain"
ack_message: str = "⏳ Processing..."
# Optional: directory to save inbound attachments. If empty, use nanobot get_media_dir("qq").
media_dir: str = ""
@@ -484,6 +485,17 @@ class QQChannel(BaseChannel):
if not content and not media_paths:
return
if self.config.ack_message:
try:
await self._send_text_only(
chat_id=chat_id,
is_group=is_group,
msg_id=data.id,
content=self.config.ack_message,
)
except Exception:
logger.debug("QQ ack message failed for chat_id={}", chat_id)
await self._handle_message(
sender_id=user_id,
chat_id=chat_id,
+87 -20
View File
@@ -28,6 +28,16 @@ TELEGRAM_MAX_MESSAGE_LEN = 4000 # Telegram message character limit
TELEGRAM_REPLY_CONTEXT_MAX_LEN = TELEGRAM_MAX_MESSAGE_LEN # Max length for reply context in user message
def _escape_telegram_html(text: str) -> str:
"""Escape text for Telegram HTML parse mode."""
return text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
def _tool_hint_to_telegram_blockquote(text: str) -> str:
"""Render tool hints as an expandable blockquote (collapsed by default)."""
return f"<blockquote expandable>{_escape_telegram_html(text)}</blockquote>" if text else ""
def _strip_md(s: str) -> str:
"""Strip markdown inline formatting from text."""
s = re.sub(r'\*\*(.+?)\*\*', r'\1', s)
@@ -120,7 +130,7 @@ def _markdown_to_telegram_html(text: str) -> str:
text = re.sub(r'^>\s*(.*)$', r'\1', text, flags=re.MULTILINE)
# 5. Escape HTML special characters
text = text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
text = _escape_telegram_html(text)
# 6. Links [text](url) - must be before bold/italic to handle nested cases
text = re.sub(r'\[([^\]]+)\]\(([^)]+)\)', r'<a href="\2">\1</a>', text)
@@ -141,13 +151,13 @@ def _markdown_to_telegram_html(text: str) -> str:
# 11. Restore inline code with HTML tags
for i, code in enumerate(inline_codes):
# Escape HTML in code content
escaped = code.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
escaped = _escape_telegram_html(code)
text = text.replace(f"\x00IC{i}\x00", f"<code>{escaped}</code>")
# 12. Restore code blocks with HTML tags
for i, code in enumerate(code_blocks):
# Escape HTML in code content
escaped = code.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
escaped = _escape_telegram_html(code)
text = text.replace(f"\x00CB{i}\x00", f"<pre><code>{escaped}</code></pre>")
return text
@@ -275,13 +285,10 @@ class TelegramChannel(BaseChannel):
self._app = builder.build()
self._app.add_error_handler(self._on_error)
# Add command handlers
self._app.add_handler(CommandHandler("start", self._on_start))
self._app.add_handler(CommandHandler("new", self._forward_command))
self._app.add_handler(CommandHandler("stop", self._forward_command))
self._app.add_handler(CommandHandler("restart", self._forward_command))
self._app.add_handler(CommandHandler("status", self._forward_command))
self._app.add_handler(CommandHandler("help", self._on_help))
# Add command handlers (using Regex to support @username suffixes before bot initialization)
self._app.add_handler(MessageHandler(filters.Regex(r"^/start(?:@\w+)?$"), self._on_start))
self._app.add_handler(MessageHandler(filters.Regex(r"^/(new|stop|restart|status)(?:@\w+)?$"), self._forward_command))
self._app.add_handler(MessageHandler(filters.Regex(r"^/help(?:@\w+)?$"), self._on_help))
# Add message handler for text, photos, voice, documents
self._app.add_handler(
@@ -313,7 +320,7 @@ class TelegramChannel(BaseChannel):
# Start polling (this runs until stopped)
await self._app.updater.start_polling(
allowed_updates=["message"],
drop_pending_updates=True # Ignore old messages on startup
drop_pending_updates=False # Process pending messages on startup
)
# Keep running until stopped
@@ -362,9 +369,14 @@ class TelegramChannel(BaseChannel):
logger.warning("Telegram bot not running")
return
# Only stop typing indicator for final responses
# Only stop typing indicator and remove reaction for final responses
if not msg.metadata.get("_progress", False):
self._stop_typing(msg.chat_id)
if reply_to_message_id := msg.metadata.get("message_id"):
try:
await self._remove_reaction(msg.chat_id, int(reply_to_message_id))
except ValueError:
pass
try:
chat_id = int(msg.chat_id)
@@ -431,11 +443,17 @@ class TelegramChannel(BaseChannel):
# Send text content
if msg.content and msg.content != "[empty message]":
render_as_blockquote = bool(msg.metadata.get("_tool_hint"))
for chunk in split_message(msg.content, TELEGRAM_MAX_MESSAGE_LEN):
await self._send_text(chat_id, chunk, reply_params, thread_kwargs)
await self._send_text(
chat_id, chunk, reply_params, thread_kwargs,
render_as_blockquote=render_as_blockquote,
)
async def _call_with_retry(self, fn, *args, **kwargs):
"""Call an async Telegram API function with retry on pool/network timeout."""
"""Call an async Telegram API function with retry on pool/network timeout and RetryAfter."""
from telegram.error import RetryAfter
for attempt in range(1, _SEND_MAX_RETRIES + 1):
try:
return await fn(*args, **kwargs)
@@ -448,6 +466,15 @@ class TelegramChannel(BaseChannel):
attempt, _SEND_MAX_RETRIES, delay,
)
await asyncio.sleep(delay)
except RetryAfter as e:
if attempt == _SEND_MAX_RETRIES:
raise
delay = float(e.retry_after)
logger.warning(
"Telegram Flood Control (attempt {}/{}), retrying in {:.1f}s",
attempt, _SEND_MAX_RETRIES, delay,
)
await asyncio.sleep(delay)
async def _send_text(
self,
@@ -455,10 +482,11 @@ class TelegramChannel(BaseChannel):
text: str,
reply_params=None,
thread_kwargs: dict | None = None,
render_as_blockquote: bool = False,
) -> None:
"""Send a plain text message with HTML fallback."""
try:
html = _markdown_to_telegram_html(text)
html = _tool_hint_to_telegram_blockquote(text) if render_as_blockquote else _markdown_to_telegram_html(text)
await self._call_with_retry(
self._app.bot.send_message,
chat_id=chat_id, text=html, parse_mode="HTML",
@@ -498,6 +526,11 @@ class TelegramChannel(BaseChannel):
if stream_id is not None and buf.stream_id is not None and buf.stream_id != stream_id:
return
self._stop_typing(chat_id)
if reply_to_message_id := meta.get("message_id"):
try:
await self._remove_reaction(chat_id, int(reply_to_message_id))
except ValueError:
pass
try:
html = _markdown_to_telegram_html(buf.text)
await self._call_with_retry(
@@ -619,8 +652,7 @@ class TelegramChannel(BaseChannel):
"reply_to_message_id": getattr(reply_to, "message_id", None) if reply_to else None,
}
@staticmethod
def _extract_reply_context(message) -> str | None:
async def _extract_reply_context(self, message) -> str | None:
"""Extract text from the message being replied to, if any."""
reply = getattr(message, "reply_to_message", None)
if not reply:
@@ -628,7 +660,21 @@ class TelegramChannel(BaseChannel):
text = getattr(reply, "text", None) or getattr(reply, "caption", None) or ""
if len(text) > TELEGRAM_REPLY_CONTEXT_MAX_LEN:
text = text[:TELEGRAM_REPLY_CONTEXT_MAX_LEN] + "..."
return f"[Reply to: {text}]" if text else None
if not text:
return None
bot_id, _ = await self._ensure_bot_identity()
reply_user = getattr(reply, "from_user", None)
if bot_id and reply_user and getattr(reply_user, "id", None) == bot_id:
return f"[Reply to bot: {text}]"
elif reply_user and getattr(reply_user, "username", None):
return f"[Reply to @{reply_user.username}: {text}]"
elif reply_user and getattr(reply_user, "first_name", None):
return f"[Reply to {reply_user.first_name}: {text}]"
else:
return f"[Reply to: {text}]"
async def _download_message_media(
self, msg, *, add_failure_content: bool = False
@@ -765,10 +811,18 @@ class TelegramChannel(BaseChannel):
message = update.message
user = update.effective_user
self._remember_thread_context(message)
# Strip @bot_username suffix if present
content = message.text or ""
if content.startswith("/") and "@" in content:
cmd_part, *rest = content.split(" ", 1)
cmd_part = cmd_part.split("@")[0]
content = f"{cmd_part} {rest[0]}" if rest else cmd_part
await self._handle_message(
sender_id=self._sender_id(user),
chat_id=str(message.chat_id),
content=message.text or "",
content=content,
metadata=self._build_message_metadata(message, user),
session_key=self._derive_topic_session_key(message),
)
@@ -812,7 +866,7 @@ class TelegramChannel(BaseChannel):
# Reply context: text and/or media from the replied-to message
reply = getattr(message, "reply_to_message", None)
if reply is not None:
reply_ctx = self._extract_reply_context(message)
reply_ctx = await self._extract_reply_context(message)
reply_media, reply_media_parts = await self._download_message_media(reply)
if reply_media:
media_paths = reply_media + media_paths
@@ -903,6 +957,19 @@ class TelegramChannel(BaseChannel):
except Exception as e:
logger.debug("Telegram reaction failed: {}", e)
async def _remove_reaction(self, chat_id: str, message_id: int) -> None:
"""Remove emoji reaction from a message (best-effort, non-blocking)."""
if not self._app:
return
try:
await self._app.bot.set_message_reaction(
chat_id=int(chat_id),
message_id=message_id,
reaction=[],
)
except Exception as e:
logger.debug("Telegram reaction removal failed: {}", e)
async def _typing_loop(self, chat_id: str) -> None:
"""Repeatedly send 'typing' action until cancelled."""
try:
+94 -360
View File
@@ -15,7 +15,6 @@ import hashlib
import json
import mimetypes
import os
import random
import re
import time
import uuid
@@ -54,26 +53,7 @@ MESSAGE_TYPE_BOT = 2
MESSAGE_STATE_FINISH = 2
WEIXIN_MAX_MESSAGE_LEN = 4000
WEIXIN_CHANNEL_VERSION = "2.1.1"
ILINK_APP_ID = "bot"
def _build_client_version(version: str) -> int:
"""Encode semantic version as 0x00MMNNPP (major/minor/patch in one uint32)."""
parts = version.split(".")
def _as_int(idx: int) -> int:
try:
return int(parts[idx])
except Exception:
return 0
major = _as_int(0)
minor = _as_int(1)
patch = _as_int(2)
return ((major & 0xFF) << 16) | ((minor & 0xFF) << 8) | (patch & 0xFF)
ILINK_APP_CLIENT_VERSION = _build_client_version(WEIXIN_CHANNEL_VERSION)
WEIXIN_CHANNEL_VERSION = "1.0.3"
BASE_INFO: dict[str, str] = {"channel_version": WEIXIN_CHANNEL_VERSION}
# Session-expired error code
@@ -85,32 +65,18 @@ MAX_CONSECUTIVE_FAILURES = 3
BACKOFF_DELAY_S = 30
RETRY_DELAY_S = 2
MAX_QR_REFRESH_COUNT = 3
TYPING_STATUS_TYPING = 1
TYPING_STATUS_CANCEL = 2
TYPING_TICKET_TTL_S = 24 * 60 * 60
TYPING_KEEPALIVE_INTERVAL_S = 5
CONFIG_CACHE_INITIAL_RETRY_S = 2
CONFIG_CACHE_MAX_RETRY_S = 60 * 60
# Default long-poll timeout; overridden by server via longpolling_timeout_ms.
DEFAULT_LONG_POLL_TIMEOUT_S = 35
# Media-type codes for getuploadurl (1=image, 2=video, 3=file, 4=voice)
# Media-type codes for getuploadurl (1=image, 2=video, 3=file)
UPLOAD_MEDIA_IMAGE = 1
UPLOAD_MEDIA_VIDEO = 2
UPLOAD_MEDIA_FILE = 3
UPLOAD_MEDIA_VOICE = 4
# File extensions considered as images / videos for outbound media
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp", ".tiff", ".ico", ".svg"}
_VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv"}
_VOICE_EXTS = {".mp3", ".wav", ".amr", ".silk", ".ogg", ".m4a", ".aac", ".flac"}
def _has_downloadable_media_locator(media: dict[str, Any] | None) -> bool:
if not isinstance(media, dict):
return False
return bool(str(media.get("encrypt_query_param", "") or "") or str(media.get("full_url", "") or "").strip())
class WeixinConfig(Base):
@@ -158,7 +124,6 @@ class WeixinChannel(BaseChannel):
self._poll_task: asyncio.Task | None = None
self._next_poll_timeout_s: int = DEFAULT_LONG_POLL_TIMEOUT_S
self._session_pause_until: float = 0.0
self._typing_tickets: dict[str, dict[str, Any]] = {}
# ------------------------------------------------------------------
# State persistence
@@ -197,7 +162,8 @@ class WeixinChannel(BaseChannel):
if base_url:
self.config.base_url = base_url
return bool(self._token)
except Exception:
except Exception as e:
logger.warning("Failed to load WeChat state: {}", e)
return False
def _save_state(self) -> None:
@@ -210,8 +176,8 @@ class WeixinChannel(BaseChannel):
"base_url": self.config.base_url,
}
state_file.write_text(json.dumps(data, ensure_ascii=False))
except Exception:
pass
except Exception as e:
logger.warning("Failed to save WeChat state: {}", e)
# ------------------------------------------------------------------
# HTTP helpers (matches api.ts buildHeaders / apiFetch)
@@ -233,8 +199,6 @@ class WeixinChannel(BaseChannel):
"X-WECHAT-UIN": self._random_wechat_uin(),
"Content-Type": "application/json",
"AuthorizationType": "ilink_bot_token",
"iLink-App-Id": ILINK_APP_ID,
"iLink-App-ClientVersion": str(ILINK_APP_CLIENT_VERSION),
}
if auth and self._token:
headers["Authorization"] = f"Bearer {self._token}"
@@ -242,15 +206,6 @@ class WeixinChannel(BaseChannel):
headers["SKRouteTag"] = str(self.config.route_tag).strip()
return headers
@staticmethod
def _is_retryable_media_download_error(err: Exception) -> bool:
if isinstance(err, httpx.TimeoutException | httpx.TransportError):
return True
if isinstance(err, httpx.HTTPStatusError):
status_code = err.response.status_code if err.response is not None else 0
return status_code >= 500
return False
async def _api_get(
self,
endpoint: str,
@@ -268,25 +223,6 @@ class WeixinChannel(BaseChannel):
resp.raise_for_status()
return resp.json()
async def _api_get_with_base(
self,
*,
base_url: str,
endpoint: str,
params: dict | None = None,
auth: bool = True,
extra_headers: dict[str, str] | None = None,
) -> dict:
"""GET helper that allows overriding base_url for QR redirect polling."""
assert self._client is not None
url = f"{base_url.rstrip('/')}/{endpoint}"
hdrs = self._make_headers(auth=auth)
if extra_headers:
hdrs.update(extra_headers)
resp = await self._client.get(url, params=params, headers=hdrs)
resp.raise_for_status()
return resp.json()
async def _api_post(
self,
endpoint: str,
@@ -323,27 +259,23 @@ class WeixinChannel(BaseChannel):
async def _qr_login(self) -> bool:
"""Perform QR code login flow. Returns True on success."""
try:
logger.info("Starting WeChat QR code login...")
refresh_count = 0
qrcode_id, scan_url = await self._fetch_qr_code()
self._print_qr_code(scan_url)
current_poll_base_url = self.config.base_url
logger.info("Waiting for QR code scan...")
while self._running:
try:
status_data = await self._api_get_with_base(
base_url=current_poll_base_url,
endpoint="ilink/bot/get_qrcode_status",
# Reference plugin sends iLink-App-ClientVersion header for
# QR status polling (login-qr.ts:81).
status_data = await self._api_get(
"ilink/bot/get_qrcode_status",
params={"qrcode": qrcode_id},
auth=False,
extra_headers={"iLink-App-ClientVersion": "1"},
)
except Exception as e:
if self._is_retryable_qr_poll_error(e):
await asyncio.sleep(1)
continue
raise
if not isinstance(status_data, dict):
await asyncio.sleep(1)
except httpx.TimeoutException:
continue
status = status_data.get("status", "")
@@ -366,15 +298,8 @@ class WeixinChannel(BaseChannel):
else:
logger.error("Login confirmed but no bot_token in response")
return False
elif status == "scaned_but_redirect":
redirect_host = str(status_data.get("redirect_host", "") or "").strip()
if redirect_host:
if redirect_host.startswith("http://") or redirect_host.startswith("https://"):
redirected_base = redirect_host
else:
redirected_base = f"https://{redirect_host}"
if redirected_base != current_poll_base_url:
current_poll_base_url = redirected_base
elif status == "scaned":
logger.info("QR code scanned, waiting for confirmation...")
elif status == "expired":
refresh_count += 1
if refresh_count > MAX_QR_REFRESH_COUNT:
@@ -384,9 +309,14 @@ class WeixinChannel(BaseChannel):
MAX_QR_REFRESH_COUNT,
)
return False
logger.warning(
"QR code expired, refreshing... ({}/{})",
refresh_count,
MAX_QR_REFRESH_COUNT,
)
qrcode_id, scan_url = await self._fetch_qr_code()
current_poll_base_url = self.config.base_url
self._print_qr_code(scan_url)
logger.info("New QR code generated, waiting for scan...")
continue
# status == "wait" — keep polling
@@ -397,16 +327,6 @@ class WeixinChannel(BaseChannel):
return False
@staticmethod
def _is_retryable_qr_poll_error(err: Exception) -> bool:
if isinstance(err, httpx.TimeoutException | httpx.TransportError):
return True
if isinstance(err, httpx.HTTPStatusError):
status_code = err.response.status_code if err.response is not None else 0
if status_code >= 500:
return True
return False
@staticmethod
def _print_qr_code(url: str) -> None:
try:
@@ -417,6 +337,7 @@ class WeixinChannel(BaseChannel):
qr.make(fit=True)
qr.print_ascii(invert=True)
except ImportError:
logger.info("QR code URL (install 'qrcode' for terminal display): {}", url)
print(f"\nLogin URL: {url}\n")
# ------------------------------------------------------------------
@@ -478,6 +399,12 @@ class WeixinChannel(BaseChannel):
if not self._running:
break
consecutive_failures += 1
logger.error(
"WeChat poll error ({}/{}): {}",
consecutive_failures,
MAX_CONSECUTIVE_FAILURES,
e,
)
if consecutive_failures >= MAX_CONSECUTIVE_FAILURES:
consecutive_failures = 0
await asyncio.sleep(BACKOFF_DELAY_S)
@@ -492,6 +419,8 @@ class WeixinChannel(BaseChannel):
await self._client.aclose()
self._client = None
self._save_state()
logger.info("WeChat channel stopped")
# ------------------------------------------------------------------
# Polling (matches monitor.ts monitorWeixinProvider)
# ------------------------------------------------------------------
@@ -517,6 +446,10 @@ class WeixinChannel(BaseChannel):
async def _poll_once(self) -> None:
remaining = self._session_pause_remaining_s()
if remaining > 0:
logger.warning(
"WeChat session paused, waiting {} min before next poll.",
max((remaining + 59) // 60, 1),
)
await asyncio.sleep(remaining)
return
@@ -566,8 +499,8 @@ class WeixinChannel(BaseChannel):
for msg in msgs:
try:
await self._process_message(msg)
except Exception:
pass
except Exception as e:
logger.error("Error processing WeChat message: {}", e)
# ------------------------------------------------------------------
# Inbound message processing (matches inbound.ts + process-message.ts)
@@ -603,7 +536,6 @@ class WeixinChannel(BaseChannel):
item_list: list[dict] = msg.get("item_list") or []
content_parts: list[str] = []
media_paths: list[str] = []
has_top_level_downloadable_media = False
for item in item_list:
item_type = item.get("type", 0)
@@ -640,8 +572,6 @@ class WeixinChannel(BaseChannel):
elif item_type == ITEM_IMAGE:
image_item = item.get("image_item") or {}
if _has_downloadable_media_locator(image_item.get("media")):
has_top_level_downloadable_media = True
file_path = await self._download_media_item(image_item, "image")
if file_path:
content_parts.append(f"[image]\n[Image: source: {file_path}]")
@@ -656,8 +586,6 @@ class WeixinChannel(BaseChannel):
if voice_text:
content_parts.append(f"[voice] {voice_text}")
else:
if _has_downloadable_media_locator(voice_item.get("media")):
has_top_level_downloadable_media = True
file_path = await self._download_media_item(voice_item, "voice")
if file_path:
transcription = await self.transcribe_audio(file_path)
@@ -671,8 +599,6 @@ class WeixinChannel(BaseChannel):
elif item_type == ITEM_FILE:
file_item = item.get("file_item") or {}
if _has_downloadable_media_locator(file_item.get("media")):
has_top_level_downloadable_media = True
file_name = file_item.get("file_name", "unknown")
file_path = await self._download_media_item(
file_item,
@@ -687,8 +613,6 @@ class WeixinChannel(BaseChannel):
elif item_type == ITEM_VIDEO:
video_item = item.get("video_item") or {}
if _has_downloadable_media_locator(video_item.get("media")):
has_top_level_downloadable_media = True
file_path = await self._download_media_item(video_item, "video")
if file_path:
content_parts.append(f"[video]\n[Video: source: {file_path}]")
@@ -696,56 +620,17 @@ class WeixinChannel(BaseChannel):
else:
content_parts.append("[video]")
# Fallback: when no top-level media was downloaded, try quoted/referenced media.
# This aligns with the reference plugin behavior that checks ref_msg.message_item
# when main item_list has no downloadable media.
if not media_paths and not has_top_level_downloadable_media:
ref_media_item: dict[str, Any] | None = None
for item in item_list:
if item.get("type", 0) != ITEM_TEXT:
continue
ref = item.get("ref_msg") or {}
candidate = ref.get("message_item") or {}
if candidate.get("type", 0) in (ITEM_IMAGE, ITEM_VOICE, ITEM_FILE, ITEM_VIDEO):
ref_media_item = candidate
break
if ref_media_item:
ref_type = ref_media_item.get("type", 0)
if ref_type == ITEM_IMAGE:
image_item = ref_media_item.get("image_item") or {}
file_path = await self._download_media_item(image_item, "image")
if file_path:
content_parts.append(f"[image]\n[Image: source: {file_path}]")
media_paths.append(file_path)
elif ref_type == ITEM_VOICE:
voice_item = ref_media_item.get("voice_item") or {}
file_path = await self._download_media_item(voice_item, "voice")
if file_path:
transcription = await self.transcribe_audio(file_path)
if transcription:
content_parts.append(f"[voice] {transcription}")
else:
content_parts.append(f"[voice]\n[Audio: source: {file_path}]")
media_paths.append(file_path)
elif ref_type == ITEM_FILE:
file_item = ref_media_item.get("file_item") or {}
file_name = file_item.get("file_name", "unknown")
file_path = await self._download_media_item(file_item, "file", file_name)
if file_path:
content_parts.append(f"[file: {file_name}]\n[File: source: {file_path}]")
media_paths.append(file_path)
elif ref_type == ITEM_VIDEO:
video_item = ref_media_item.get("video_item") or {}
file_path = await self._download_media_item(video_item, "video")
if file_path:
content_parts.append(f"[video]\n[Video: source: {file_path}]")
media_paths.append(file_path)
content = "\n".join(content_parts)
if not content:
return
logger.info(
"WeChat inbound: from={} items={} bodyLen={}",
from_user_id,
",".join(str(i.get("type", 0)) for i in item_list),
len(content),
)
await self._handle_message(
sender_id=from_user_id,
chat_id=from_user_id,
@@ -767,10 +652,9 @@ class WeixinChannel(BaseChannel):
"""Download + AES-decrypt a media item. Returns local path or None."""
try:
media = typed_item.get("media") or {}
encrypt_query_param = str(media.get("encrypt_query_param", "") or "")
full_url = str(media.get("full_url", "") or "").strip()
encrypt_query_param = media.get("encrypt_query_param", "")
if not encrypt_query_param and not full_url:
if not encrypt_query_param:
return None
# Resolve AES key (media-download.ts:43-45, pic-decrypt.ts:40-52)
@@ -787,50 +671,21 @@ class WeixinChannel(BaseChannel):
elif media_aes_key_b64:
aes_key_b64 = media_aes_key_b64
# Reference protocol behavior: VOICE/FILE/VIDEO require aes_key;
# only IMAGE may be downloaded as plain bytes when key is missing.
if media_type != "image" and not aes_key_b64:
return None
# Build CDN download URL with proper URL-encoding (cdn-url.ts:7)
cdn_url = (
f"{self.config.cdn_base_url}/download"
f"?encrypted_query_param={quote(encrypt_query_param)}"
)
assert self._client is not None
fallback_url = ""
if encrypt_query_param:
fallback_url = (
f"{self.config.cdn_base_url}/download"
f"?encrypted_query_param={quote(encrypt_query_param)}"
)
download_candidates: list[tuple[str, str]] = []
if full_url:
download_candidates.append(("full_url", full_url))
if fallback_url and (not full_url or fallback_url != full_url):
download_candidates.append(("encrypt_query_param", fallback_url))
data = b""
for idx, (download_source, cdn_url) in enumerate(download_candidates):
try:
resp = await self._client.get(cdn_url)
resp.raise_for_status()
data = resp.content
break
except Exception as e:
has_more_candidates = idx + 1 < len(download_candidates)
should_fallback = (
download_source == "full_url"
and has_more_candidates
and self._is_retryable_media_download_error(e)
)
if should_fallback:
logger.warning(
"WeChat media download failed via full_url, falling back to encrypt_query_param: type={} err={}",
media_type,
e,
)
continue
raise
resp = await self._client.get(cdn_url)
resp.raise_for_status()
data = resp.content
if aes_key_b64 and data:
data = _decrypt_aes_ecb(data, aes_key_b64)
elif not aes_key_b64:
logger.debug("No AES key for {} item, using raw bytes", media_type)
if not data:
return None
@@ -839,12 +694,12 @@ class WeixinChannel(BaseChannel):
ext = _ext_for_type(media_type)
if not filename:
ts = int(time.time())
hash_seed = encrypt_query_param or full_url
h = abs(hash(hash_seed)) % 100000
h = abs(hash(encrypt_query_param)) % 100000
filename = f"{media_type}_{ts}_{h}{ext}"
safe_name = os.path.basename(filename)
file_path = media_dir / safe_name
file_path.write_bytes(data)
logger.debug("Downloaded WeChat {} to {}", media_type, file_path)
return str(file_path)
except Exception as e:
@@ -855,76 +710,14 @@ class WeixinChannel(BaseChannel):
# Outbound (matches send.ts buildTextMessageReq + sendMessageWeixin)
# ------------------------------------------------------------------
async def _get_typing_ticket(self, user_id: str, context_token: str = "") -> str:
"""Get typing ticket with per-user refresh + failure backoff cache."""
now = time.time()
entry = self._typing_tickets.get(user_id)
if entry and now < float(entry.get("next_fetch_at", 0)):
return str(entry.get("ticket", "") or "")
body: dict[str, Any] = {
"ilink_user_id": user_id,
"context_token": context_token or None,
"base_info": BASE_INFO,
}
data = await self._api_post("ilink/bot/getconfig", body)
if data.get("ret", 0) == 0:
ticket = str(data.get("typing_ticket", "") or "")
self._typing_tickets[user_id] = {
"ticket": ticket,
"ever_succeeded": True,
"next_fetch_at": now + (random.random() * TYPING_TICKET_TTL_S),
"retry_delay_s": CONFIG_CACHE_INITIAL_RETRY_S,
}
return ticket
prev_delay = float(entry.get("retry_delay_s", CONFIG_CACHE_INITIAL_RETRY_S)) if entry else CONFIG_CACHE_INITIAL_RETRY_S
next_delay = min(prev_delay * 2, CONFIG_CACHE_MAX_RETRY_S)
if entry:
entry["next_fetch_at"] = now + next_delay
entry["retry_delay_s"] = next_delay
return str(entry.get("ticket", "") or "")
self._typing_tickets[user_id] = {
"ticket": "",
"ever_succeeded": False,
"next_fetch_at": now + CONFIG_CACHE_INITIAL_RETRY_S,
"retry_delay_s": CONFIG_CACHE_INITIAL_RETRY_S,
}
return ""
async def _send_typing(self, user_id: str, typing_ticket: str, status: int) -> None:
"""Best-effort sendtyping wrapper."""
if not typing_ticket:
return
body: dict[str, Any] = {
"ilink_user_id": user_id,
"typing_ticket": typing_ticket,
"status": status,
"base_info": BASE_INFO,
}
await self._api_post("ilink/bot/sendtyping", body)
async def _typing_keepalive_loop(self, user_id: str, typing_ticket: str, stop_event: asyncio.Event) -> None:
try:
while not stop_event.is_set():
await asyncio.sleep(TYPING_KEEPALIVE_INTERVAL_S)
if stop_event.is_set():
break
try:
await self._send_typing(user_id, typing_ticket, TYPING_STATUS_TYPING)
except Exception:
pass
finally:
pass
async def send(self, msg: OutboundMessage) -> None:
if not self._client or not self._token:
logger.warning("WeChat client not initialized or not authenticated")
return
try:
self._assert_session_active()
except RuntimeError:
except RuntimeError as e:
logger.warning("WeChat send blocked: {}", e)
return
content = msg.content.strip()
@@ -936,62 +729,29 @@ class WeixinChannel(BaseChannel):
)
return
typing_ticket = ""
try:
typing_ticket = await self._get_typing_ticket(msg.chat_id, ctx_token)
except Exception:
typing_ticket = ""
if typing_ticket:
# --- Send media files first (following Telegram channel pattern) ---
for media_path in (msg.media or []):
try:
await self._send_typing(msg.chat_id, typing_ticket, TYPING_STATUS_TYPING)
except Exception:
pass
await self._send_media_file(msg.chat_id, media_path, ctx_token)
except Exception as e:
filename = Path(media_path).name
logger.error("Failed to send WeChat media {}: {}", media_path, e)
# Notify user about failure via text
await self._send_text(
msg.chat_id, f"[Failed to send: {filename}]", ctx_token,
)
typing_keepalive_stop = asyncio.Event()
typing_keepalive_task: asyncio.Task | None = None
if typing_ticket:
typing_keepalive_task = asyncio.create_task(
self._typing_keepalive_loop(msg.chat_id, typing_ticket, typing_keepalive_stop)
)
# --- Send text content ---
if not content:
return
try:
# --- Send media files first (following Telegram channel pattern) ---
for media_path in (msg.media or []):
try:
await self._send_media_file(msg.chat_id, media_path, ctx_token)
except Exception as e:
filename = Path(media_path).name
logger.error("Failed to send WeChat media {}: {}", media_path, e)
# Notify user about failure via text
await self._send_text(
msg.chat_id, f"[Failed to send: {filename}]", ctx_token,
)
# --- Send text content ---
if not content:
return
chunks = split_message(content, WEIXIN_MAX_MESSAGE_LEN)
for chunk in chunks:
await self._send_text(msg.chat_id, chunk, ctx_token)
except Exception as e:
logger.error("Error sending WeChat message: {}", e)
raise
finally:
if typing_keepalive_task:
typing_keepalive_stop.set()
typing_keepalive_task.cancel()
try:
await typing_keepalive_task
except asyncio.CancelledError:
pass
if typing_ticket:
try:
await self._send_typing(msg.chat_id, typing_ticket, TYPING_STATUS_CANCEL)
except Exception:
pass
async def _send_text(
self,
@@ -1065,10 +825,6 @@ class WeixinChannel(BaseChannel):
upload_type = UPLOAD_MEDIA_VIDEO
item_type = ITEM_VIDEO
item_key = "video_item"
elif ext in _VOICE_EXTS:
upload_type = UPLOAD_MEDIA_VOICE
item_type = ITEM_VOICE
item_key = "voice_item"
else:
upload_type = UPLOAD_MEDIA_FILE
item_type = ITEM_FILE
@@ -1082,7 +838,7 @@ class WeixinChannel(BaseChannel):
# Matches aesEcbPaddedSize: Math.ceil((size + 1) / 16) * 16
padded_size = ((raw_size + 1 + 15) // 16) * 16
# Step 1: Get upload URL from server (prefer upload_full_url, fallback to upload_param)
# Step 1: Get upload URL (upload_param) from server
file_key = os.urandom(16).hex()
upload_body: dict[str, Any] = {
"filekey": file_key,
@@ -1097,27 +853,22 @@ class WeixinChannel(BaseChannel):
assert self._client is not None
upload_resp = await self._api_post("ilink/bot/getuploadurl", upload_body)
logger.debug("WeChat getuploadurl response: {}", upload_resp)
upload_full_url = str(upload_resp.get("upload_full_url", "") or "").strip()
upload_param = str(upload_resp.get("upload_param", "") or "")
if not upload_full_url and not upload_param:
raise RuntimeError(
"getuploadurl returned no upload URL "
f"(need upload_full_url or upload_param): {upload_resp}"
)
upload_param = upload_resp.get("upload_param", "")
if not upload_param:
raise RuntimeError(f"getuploadurl returned no upload_param: {upload_resp}")
# Step 2: AES-128-ECB encrypt and POST to CDN
aes_key_b64 = base64.b64encode(aes_key_raw).decode()
encrypted_data = _encrypt_aes_ecb(raw_data, aes_key_b64)
if upload_full_url:
cdn_upload_url = upload_full_url
else:
cdn_upload_url = (
f"{self.config.cdn_base_url}/upload"
f"?encrypted_query_param={quote(upload_param)}"
f"&filekey={quote(file_key)}"
)
cdn_upload_url = (
f"{self.config.cdn_base_url}/upload"
f"?encrypted_query_param={quote(upload_param)}"
f"&filekey={quote(file_key)}"
)
logger.debug("WeChat CDN POST url={} ciphertextSize={}", cdn_upload_url[:80], len(encrypted_data))
cdn_resp = await self._client.post(
cdn_upload_url,
@@ -1133,6 +884,7 @@ class WeixinChannel(BaseChannel):
"CDN upload response missing x-encrypted-param header; "
f"status={cdn_resp.status_code} headers={dict(cdn_resp.headers)}"
)
logger.debug("WeChat CDN upload success for {}, got download_param", p.name)
# Step 3: Send message with the media item
# aes_key for CDNMedia is the hex key encoded as base64
@@ -1181,6 +933,7 @@ class WeixinChannel(BaseChannel):
raise RuntimeError(
f"WeChat send media error (code {errcode}): {data.get('errmsg', '')}"
)
logger.info("WeChat media sent: {} (type={})", p.name, item_key)
# ---------------------------------------------------------------------------
@@ -1252,42 +1005,23 @@ def _decrypt_aes_ecb(data: bytes, aes_key_b64: str) -> bytes:
logger.warning("Failed to parse AES key, returning raw data: {}", e)
return data
decrypted: bytes | None = None
try:
from Crypto.Cipher import AES
cipher = AES.new(key, AES.MODE_ECB)
decrypted = cipher.decrypt(data)
return cipher.decrypt(data) # pycryptodome auto-strips PKCS7 with unpad
except ImportError:
pass
if decrypted is None:
try:
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
try:
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
cipher_obj = Cipher(algorithms.AES(key), modes.ECB())
decryptor = cipher_obj.decryptor()
decrypted = decryptor.update(data) + decryptor.finalize()
except ImportError:
logger.warning("Cannot decrypt media: install 'pycryptodome' or 'cryptography'")
return data
return _pkcs7_unpad_safe(decrypted)
def _pkcs7_unpad_safe(data: bytes, block_size: int = 16) -> bytes:
"""Safely remove PKCS7 padding when valid; otherwise return original bytes."""
if not data:
cipher_obj = Cipher(algorithms.AES(key), modes.ECB())
decryptor = cipher_obj.decryptor()
return decryptor.update(data) + decryptor.finalize()
except ImportError:
logger.warning("Cannot decrypt media: install 'pycryptodome' or 'cryptography'")
return data
if len(data) % block_size != 0:
return data
pad_len = data[-1]
if pad_len < 1 or pad_len > block_size:
return data
if data[-pad_len:] != bytes([pad_len]) * pad_len:
return data
return data[:-pad_len]
def _ext_for_type(media_type: str) -> str:
+28
View File
@@ -22,6 +22,7 @@ if sys.platform == "win32":
pass
import typer
from loguru import logger
from prompt_toolkit import PromptSession, print_formatted_text
from prompt_toolkit.application import run_in_terminal
from prompt_toolkit.formatted_text import ANSI, HTML
@@ -640,6 +641,15 @@ def gateway(
# Set cron callback (needs agent)
async def on_cron_job(job: CronJob) -> str | None:
"""Execute a cron job through the agent."""
# Dream is an internal job — run directly, not through the agent loop.
if job.name == "dream":
try:
await agent.dream.run()
logger.info("Dream cron job completed")
except Exception:
logger.exception("Dream cron job failed")
return None
from nanobot.agent.tools.cron import CronTool
from nanobot.agent.tools.message import MessageTool
from nanobot.utils.evaluator import evaluate_response
@@ -759,6 +769,21 @@ def gateway(
console.print(f"[green]✓[/green] Heartbeat: every {hb_cfg.interval_s}s")
# Register Dream cron job (always-on, idempotent on restart)
dream_cfg = config.agents.defaults.dream
if dream_cfg.model:
agent.dream.model = dream_cfg.model
agent.dream.max_batch_size = dream_cfg.max_batch_size
agent.dream.max_iterations = dream_cfg.max_iterations
from nanobot.cron.types import CronJob, CronPayload, CronSchedule
cron.register_system_job(CronJob(
id="dream",
name="dream",
schedule=CronSchedule(kind="cron", expr=dream_cfg.cron, tz=config.agents.defaults.timezone),
payload=CronPayload(kind="system_event"),
))
console.print(f"[green]✓[/green] Dream: cron {dream_cfg.cron}")
async def run():
try:
await cron.start()
@@ -959,6 +984,9 @@ def agent(
while True:
try:
_flush_pending_tty_input()
# Stop spinner before user input to avoid prompt_toolkit conflicts
if renderer:
renderer.stop_for_input()
user_input = await _read_interactive_input_async()
command = user_input.strip()
if not command:
+5 -1
View File
@@ -18,7 +18,7 @@ from nanobot import __logo__
def _make_console() -> Console:
return Console(file=sys.stdout)
return Console(file=sys.stdout, force_terminal=True)
class ThinkingSpinner:
@@ -120,6 +120,10 @@ class StreamRenderer:
else:
_make_console().print()
def stop_for_input(self) -> None:
"""Stop spinner before user input to avoid prompt_toolkit conflicts."""
self._stop_spinner()
async def close(self) -> None:
"""Stop spinner/live without rendering a final streamed round."""
if self._live:
+118 -6
View File
@@ -26,7 +26,10 @@ async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
sub_cancelled = await loop.subagents.cancel_by_session(msg.session_key)
total = cancelled + sub_cancelled
content = f"Stopped {total} task(s)." if total else "No active task to stop."
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id, content=content)
return OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id, content=content,
metadata=dict(msg.metadata or {})
)
async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
@@ -38,7 +41,10 @@ async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
os.execv(sys.executable, [sys.executable, "-m", "nanobot"] + sys.argv[1:])
asyncio.create_task(_do_restart())
return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id, content="Restarting...")
return OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id, content="Restarting...",
metadata=dict(msg.metadata or {})
)
async def cmd_status(ctx: CommandContext) -> OutboundMessage:
@@ -47,7 +53,7 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
session = ctx.session or loop.sessions.get_or_create(ctx.key)
ctx_est = 0
try:
ctx_est, _ = loop.memory_consolidator.estimate_session_prompt_tokens(session)
ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(session)
except Exception:
pass
if ctx_est <= 0:
@@ -62,7 +68,7 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
session_msg_count=len(session.get_history(max_messages=0)),
context_tokens_estimate=ctx_est,
),
metadata={"render_as": "text"},
metadata={**dict(ctx.msg.metadata or {}), "render_as": "text"},
)
@@ -75,10 +81,108 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
loop.sessions.save(session)
loop.sessions.invalidate(session.key)
if snapshot:
loop._schedule_background(loop.memory_consolidator.archive_messages(snapshot))
loop._schedule_background(loop.consolidator.archive(snapshot))
return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content="New session started.",
metadata=dict(ctx.msg.metadata or {})
)
async def cmd_dream(ctx: CommandContext) -> OutboundMessage:
"""Manually trigger a Dream consolidation run."""
loop = ctx.loop
try:
did_work = await loop.dream.run()
content = "Dream completed." if did_work else "Dream: nothing to process."
except Exception as e:
content = f"Dream failed: {e}"
return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id, content=content,
)
async def cmd_dream_log(ctx: CommandContext) -> OutboundMessage:
"""Show what the last Dream changed.
Default: diff of the latest commit (HEAD~1 vs HEAD).
With /dream-log <sha>: diff of that specific commit.
"""
store = ctx.loop.consolidator.store
git = store.git
if not git.is_initialized():
if store.get_last_dream_cursor() == 0:
msg = "Dream has not run yet."
else:
msg = "Git not initialized for memory files."
return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content=msg, metadata={"render_as": "text"},
)
args = ctx.args.strip()
if args:
# Show diff of a specific commit
sha = args.split()[0]
result = git.show_commit_diff(sha)
if not result:
content = f"Commit `{sha}` not found."
else:
commit, diff = result
content = commit.format(diff)
else:
# Default: show the latest commit's diff
result = git.show_commit_diff(git.log(max_entries=1)[0].sha) if git.log(max_entries=1) else None
if result:
commit, diff = result
content = commit.format(diff)
else:
content = "No commits yet."
return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content=content, metadata={"render_as": "text"},
)
async def cmd_dream_restore(ctx: CommandContext) -> OutboundMessage:
"""Restore memory files from a previous dream commit.
Usage:
/dream-restore — list recent commits
/dream-restore <sha> — revert a specific commit
"""
store = ctx.loop.consolidator.store
git = store.git
if not git.is_initialized():
return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content="Git not initialized for memory files.",
)
args = ctx.args.strip()
if not args:
# Show recent commits for the user to pick
commits = git.log(max_entries=10)
if not commits:
content = "No commits found."
else:
lines = ["## Recent Dream Commits\n", "Use `/dream-restore <sha>` to revert a commit.\n"]
for c in commits:
lines.append(f"- `{c.sha}` {c.message.splitlines()[0]} ({c.timestamp})")
content = "\n".join(lines)
else:
sha = args.split()[0]
new_sha = git.revert(sha)
if new_sha:
content = f"Reverted commit `{sha}` → new commit `{new_sha}`."
else:
content = f"Failed to revert commit `{sha}`. Check if the SHA is correct."
return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content=content, metadata={"render_as": "text"},
)
@@ -88,7 +192,7 @@ async def cmd_help(ctx: CommandContext) -> OutboundMessage:
channel=ctx.msg.channel,
chat_id=ctx.msg.chat_id,
content=build_help_text(),
metadata={"render_as": "text"},
metadata={**dict(ctx.msg.metadata or {}), "render_as": "text"},
)
@@ -100,6 +204,9 @@ def build_help_text() -> str:
"/stop — Stop the current task",
"/restart — Restart the bot",
"/status — Show bot status",
"/dream — Manually trigger Dream consolidation",
"/dream-log — Show what the last Dream changed",
"/dream-restore — Revert memory to a previous state",
"/help — Show available commands",
]
return "\n".join(lines)
@@ -112,4 +219,9 @@ def register_builtin_commands(router: CommandRouter) -> None:
router.priority("/status", cmd_status)
router.exact("/new", cmd_new)
router.exact("/status", cmd_status)
router.exact("/dream", cmd_dream)
router.exact("/dream-log", cmd_dream_log)
router.prefix("/dream-log ", cmd_dream_log)
router.exact("/dream-restore", cmd_dream_restore)
router.prefix("/dream-restore ", cmd_dream_restore)
router.exact("/help", cmd_help)
+11 -1
View File
@@ -28,6 +28,15 @@ class ChannelsConfig(Base):
send_max_retries: int = Field(default=3, ge=0, le=10) # Max delivery attempts (initial send included)
class DreamConfig(Base):
"""Dream memory consolidation configuration."""
cron: str = "0 */2 * * *" # Every 2 hours
model: str | None = None # Override model for Dream
max_batch_size: int = Field(default=20, ge=1) # Max history entries per run
max_iterations: int = Field(default=10, ge=1) # Max tool calls per Phase 2
class AgentDefaults(Base):
"""Default agent configuration."""
@@ -42,6 +51,7 @@ class AgentDefaults(Base):
max_tool_iterations: int = 40
reasoning_effort: str | None = None # low / medium / high - enables LLM thinking mode
timezone: str = "UTC" # IANA timezone, e.g. "Asia/Shanghai", "America/New_York"
dream: DreamConfig = Field(default_factory=DreamConfig)
class AgentsConfig(Base):
@@ -86,6 +96,7 @@ class ProvidersConfig(Base):
byteplus_coding_plan: ProviderConfig = Field(default_factory=ProviderConfig) # BytePlus Coding Plan
openai_codex: ProviderConfig = Field(default_factory=ProviderConfig, exclude=True) # OpenAI Codex (OAuth)
github_copilot: ProviderConfig = Field(default_factory=ProviderConfig, exclude=True) # Github Copilot (OAuth)
qianfan: ProviderConfig = Field(default_factory=ProviderConfig) # Qianfan (百度千帆)
class HeartbeatConfig(Base):
@@ -136,7 +147,6 @@ class ExecToolConfig(Base):
enable: bool = True
timeout: int = 60
path_append: str = ""
command_wrapper: str = "" # sandbox wrapper command template; supports {command} and {cwd}
class MCPServerConfig(Base):
"""MCP server connection configuration (stdio or HTTP)."""
+14
View File
@@ -351,6 +351,20 @@ class CronService:
logger.info("Cron: added job '{}' ({})", name, job.id)
return job
def register_system_job(self, job: CronJob) -> CronJob:
"""Register an internal system job (idempotent on restart)."""
store = self._load_store()
now = _now_ms()
job.state = CronJobState(next_run_at_ms=_compute_next_run(job.schedule, now))
job.created_at_ms = now
job.updated_at_ms = now
store.jobs = [j for j in store.jobs if j.id != job.id]
store.jobs.append(job)
self._save_store()
self._arm_timer()
logger.info("Cron: registered system job '{}' ({})", job.name, job.id)
return job
def remove_job(self, job_id: str) -> bool:
"""Remove a job by ID."""
store = self._load_store()
+4
View File
@@ -379,6 +379,10 @@ class AnthropicProvider(LLMProvider):
val = getattr(response.usage, attr, 0)
if val:
usage[attr] = val
# Normalize to cached_tokens for downstream consistency.
cache_read = usage.get("cache_read_input_tokens", 0)
if cache_read:
usage["cached_tokens"] = cache_read
return LLMResponse(
content="".join(content_parts) or None,
+46 -5
View File
@@ -308,6 +308,13 @@ class OpenAICompatProvider(LLMProvider):
@classmethod
def _extract_usage(cls, response: Any) -> dict[str, int]:
"""Extract token usage from an OpenAI-compatible response.
Handles both dict-based (raw JSON) and object-based (SDK Pydantic)
responses. Provider-specific ``cached_tokens`` fields are normalised
under a single key; see the priority chain inside for details.
"""
# --- resolve usage object ---
usage_obj = None
response_map = cls._maybe_mapping(response)
if response_map is not None:
@@ -317,19 +324,53 @@ class OpenAICompatProvider(LLMProvider):
usage_map = cls._maybe_mapping(usage_obj)
if usage_map is not None:
return {
result = {
"prompt_tokens": int(usage_map.get("prompt_tokens") or 0),
"completion_tokens": int(usage_map.get("completion_tokens") or 0),
"total_tokens": int(usage_map.get("total_tokens") or 0),
}
if usage_obj:
return {
elif usage_obj:
result = {
"prompt_tokens": getattr(usage_obj, "prompt_tokens", 0) or 0,
"completion_tokens": getattr(usage_obj, "completion_tokens", 0) or 0,
"total_tokens": getattr(usage_obj, "total_tokens", 0) or 0,
}
return {}
else:
return {}
# --- cached_tokens (normalised across providers) ---
# Try nested paths first (dict), fall back to attribute (SDK object).
# Priority order ensures the most specific field wins.
for path in (
("prompt_tokens_details", "cached_tokens"), # OpenAI/Zhipu/MiniMax/Qwen/Mistral/xAI
("cached_tokens",), # StepFun/Moonshot (top-level)
("prompt_cache_hit_tokens",), # DeepSeek/SiliconFlow
):
cached = cls._get_nested_int(usage_map, path)
if not cached and usage_obj:
cached = cls._get_nested_int(usage_obj, path)
if cached:
result["cached_tokens"] = cached
break
return result
@staticmethod
def _get_nested_int(obj: Any, path: tuple[str, ...]) -> int:
"""Drill into *obj* by *path* segments and return an ``int`` value.
Supports both dict-key access and attribute access so it works
uniformly with raw JSON dicts **and** SDK Pydantic models.
"""
current = obj
for segment in path:
if current is None:
return 0
if isinstance(current, dict):
current = current.get(segment)
else:
current = getattr(current, segment, None)
return int(current or 0) if current is not None else 0
def _parse(self, response: Any) -> LLMResponse:
if isinstance(response, str):
+10
View File
@@ -200,6 +200,7 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
env_key="OPENAI_API_KEY",
display_name="OpenAI",
backend="openai_compat",
supports_max_completion_tokens=True,
),
# OpenAI Codex: OAuth-based, dedicated provider
ProviderSpec(
@@ -338,6 +339,15 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
backend="openai_compat",
default_api_base="https://api.groq.com/openai/v1",
),
# Qianfan (百度千帆): OpenAI-compatible API
ProviderSpec(
name="qianfan",
keywords=("qianfan", "ernie"),
env_key="QIANFAN_API_KEY",
display_name="Qianfan",
backend="openai_compat",
default_api_base="https://qianfan.baidubce.com/v2"
),
)
+15 -22
View File
@@ -1,6 +1,6 @@
---
name: memory
description: Two-layer memory system with grep-based recall.
description: Two-layer memory system with Dream-managed knowledge files.
always: true
---
@@ -8,30 +8,23 @@ always: true
## Structure
- `memory/MEMORY.md` — Long-term facts (preferences, project context, relationships). Always loaded into your context.
- `memory/HISTORY.md` — Append-only event log. NOT loaded into context. Search it with grep-style tools or in-memory filters. Each entry starts with [YYYY-MM-DD HH:MM].
- `SOUL.md` — Bot personality and communication style. **Managed by Dream.** Do NOT edit.
- `USER.md` — User profile and preferences. **Managed by Dream.** Do NOT edit.
- `memory/MEMORY.md` — Long-term facts (project context, important events). **Managed by Dream.** Do NOT edit.
- `memory/history.jsonl` — append-only JSONL, not loaded into context. search with `jq`-style tools.
- `memory/.dream-log.md` — Changelog of what Dream changed. View with `/dream-log`.
## Search Past Events
Choose the search method based on file size:
`memory/history.jsonl` is JSONL format — each line is a JSON object with `cursor`, `timestamp`, `content`.
- Small `memory/HISTORY.md`: use `read_file`, then search in-memory
- Large or long-lived `memory/HISTORY.md`: use the `exec` tool for targeted search
Examples (replace `keyword`):
- **Python (cross-platform):** `python -c "import json; [print(json.loads(l).get('content','')) for l in open('memory/history.jsonl','r',encoding='utf-8') if l.strip() and 'keyword' in l.lower()][-20:]"`
- **jq:** `cat memory/history.jsonl | jq -r 'select(.content | test("keyword"; "i")) | .content' | tail -20`
- **grep:** `grep -i "keyword" memory/history.jsonl`
Examples:
- **Linux/macOS:** `grep -i "keyword" memory/HISTORY.md`
- **Windows:** `findstr /i "keyword" memory\HISTORY.md`
- **Cross-platform Python:** `python -c "from pathlib import Path; text = Path('memory/HISTORY.md').read_text(encoding='utf-8'); print('\n'.join([l for l in text.splitlines() if 'keyword' in l.lower()][-20:]))"`
## Important
Prefer targeted command-line search for large history files.
## When to Update MEMORY.md
Write important facts immediately using `edit_file` or `write_file`:
- User preferences ("I prefer dark mode")
- Project context ("The API uses OAuth2")
- Relationships ("Alice is the project lead")
## Auto-consolidation
Old conversations are automatically summarized and appended to HISTORY.md when the session grows large. Long-term facts are extracted to MEMORY.md. You don't need to manage this.
- **Do NOT edit SOUL.md, USER.md, or MEMORY.md.** They are automatically managed by Dream.
- If you notice outdated information, it will be corrected when Dream runs next.
- Users can view Dream's activity with the `/dream-log` command.
+17 -2
View File
@@ -255,14 +255,18 @@ def build_status_content(
)
last_in = last_usage.get("prompt_tokens", 0)
last_out = last_usage.get("completion_tokens", 0)
cached = last_usage.get("cached_tokens", 0)
ctx_total = max(context_window_tokens, 0)
ctx_pct = int((context_tokens_estimate / ctx_total) * 100) if ctx_total > 0 else 0
ctx_used_str = f"{context_tokens_estimate // 1000}k" if context_tokens_estimate >= 1000 else str(context_tokens_estimate)
ctx_total_str = f"{ctx_total // 1024}k" if ctx_total > 0 else "n/a"
token_line = f"\U0001f4ca Tokens: {last_in} in / {last_out} out"
if cached and last_in:
token_line += f" ({cached * 100 // last_in}% cached)"
return "\n".join([
f"\U0001f408 nanobot v{version}",
f"\U0001f9e0 Model: {model}",
f"\U0001f4ca Tokens: {last_in} in / {last_out} out",
token_line,
f"\U0001f4da Context: {ctx_used_str}/{ctx_total_str} ({ctx_pct}%)",
f"\U0001f4ac Session: {session_msg_count} messages",
f"\u23f1 Uptime: {uptime}",
@@ -292,11 +296,22 @@ def sync_workspace_templates(workspace: Path, silent: bool = False) -> list[str]
if item.name.endswith(".md") and not item.name.startswith("."):
_write(item, workspace / item.name)
_write(tpl / "memory" / "MEMORY.md", workspace / "memory" / "MEMORY.md")
_write(None, workspace / "memory" / "HISTORY.md")
_write(None, workspace / "memory" / "history.jsonl")
(workspace / "skills").mkdir(exist_ok=True)
if added and not silent:
from rich.console import Console
for name in added:
Console().print(f" [dim]Created {name}[/dim]")
# Initialize git for memory version control
try:
from nanobot.agent.git_store import GitStore
gs = GitStore(workspace, tracked_files=[
"SOUL.md", "USER.md", "memory/MEMORY.md",
])
gs.init()
except Exception:
pass
return added
+1
View File
@@ -48,6 +48,7 @@ dependencies = [
"chardet>=3.0.2,<6.0.0",
"openai>=2.8.0",
"tiktoken>=0.12.0,<1.0.0",
"dulwich>=0.22.0,<1.0.0",
]
[project.optional-dependencies]
+10 -10
View File
@@ -506,7 +506,7 @@ class TestNewCommandArchival:
@pytest.mark.asyncio
async def test_new_clears_session_immediately_even_if_archive_fails(self, tmp_path: Path) -> None:
"""/new clears session immediately; archive_messages retries until raw dump."""
"""/new clears session immediately; archive is fire-and-forget."""
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
@@ -518,12 +518,12 @@ class TestNewCommandArchival:
call_count = 0
async def _failing_consolidate(_messages) -> bool:
async def _failing_summarize(_messages) -> bool:
nonlocal call_count
call_count += 1
return False
loop.memory_consolidator.consolidate_messages = _failing_consolidate # type: ignore[method-assign]
loop.consolidator.archive = _failing_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg)
@@ -535,7 +535,7 @@ class TestNewCommandArchival:
assert len(session_after.messages) == 0
await loop.close_mcp()
assert call_count == 3 # retried up to raw-archive threshold
assert call_count == 1
@pytest.mark.asyncio
async def test_new_archives_only_unconsolidated_messages(self, tmp_path: Path) -> None:
@@ -551,12 +551,12 @@ class TestNewCommandArchival:
archived_count = -1
async def _fake_consolidate(messages) -> bool:
async def _fake_summarize(messages) -> bool:
nonlocal archived_count
archived_count = len(messages)
return True
loop.memory_consolidator.consolidate_messages = _fake_consolidate # type: ignore[method-assign]
loop.consolidator.archive = _fake_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg)
@@ -578,10 +578,10 @@ class TestNewCommandArchival:
session.add_message("assistant", f"resp{i}")
loop.sessions.save(session)
async def _ok_consolidate(_messages) -> bool:
async def _ok_summarize(_messages) -> bool:
return True
loop.memory_consolidator.consolidate_messages = _ok_consolidate # type: ignore[method-assign]
loop.consolidator.archive = _ok_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg)
@@ -604,12 +604,12 @@ class TestNewCommandArchival:
archived = asyncio.Event()
async def _slow_consolidate(_messages) -> bool:
async def _slow_summarize(_messages) -> bool:
await asyncio.sleep(0.1)
archived.set()
return True
loop.memory_consolidator.consolidate_messages = _slow_consolidate # type: ignore[method-assign]
loop.consolidator.archive = _slow_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
await loop._process_message(new_msg)
+78
View File
@@ -0,0 +1,78 @@
"""Tests for the lightweight Consolidator — append-only to HISTORY.md."""
import pytest
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
from nanobot.agent.memory import Consolidator, MemoryStore
@pytest.fixture
def store(tmp_path):
return MemoryStore(tmp_path)
@pytest.fixture
def mock_provider():
p = MagicMock()
p.chat_with_retry = AsyncMock()
return p
@pytest.fixture
def consolidator(store, mock_provider):
sessions = MagicMock()
sessions.save = MagicMock()
return Consolidator(
store=store,
provider=mock_provider,
model="test-model",
sessions=sessions,
context_window_tokens=1000,
build_messages=MagicMock(return_value=[]),
get_tool_definitions=MagicMock(return_value=[]),
max_completion_tokens=100,
)
class TestConsolidatorSummarize:
async def test_summarize_appends_to_history(self, consolidator, mock_provider, store):
"""Consolidator should call LLM to summarize, then append to HISTORY.md."""
mock_provider.chat_with_retry.return_value = MagicMock(
content="User fixed a bug in the auth module."
)
messages = [
{"role": "user", "content": "fix the auth bug"},
{"role": "assistant", "content": "Done, fixed the race condition."},
]
result = await consolidator.archive(messages)
assert result is True
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1
async def test_summarize_raw_dumps_on_llm_failure(self, consolidator, mock_provider, store):
"""On LLM failure, raw-dump messages to HISTORY.md."""
mock_provider.chat_with_retry.side_effect = Exception("API error")
messages = [{"role": "user", "content": "hello"}]
result = await consolidator.archive(messages)
assert result is True # always succeeds
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1
assert "[RAW]" in entries[0]["content"]
async def test_summarize_skips_empty_messages(self, consolidator):
result = await consolidator.archive([])
assert result is False
class TestConsolidatorTokenBudget:
async def test_prompt_below_threshold_does_not_consolidate(self, consolidator):
"""No consolidation when tokens are within budget."""
session = MagicMock()
session.last_consolidated = 0
session.messages = [{"role": "user", "content": "hi"}]
session.key = "test:key"
consolidator.estimate_session_prompt_tokens = MagicMock(return_value=(100, "tiktoken"))
consolidator.archive = AsyncMock(return_value=True)
await consolidator.maybe_consolidate_by_tokens(session)
consolidator.archive.assert_not_called()
+97
View File
@@ -0,0 +1,97 @@
"""Tests for the Dream class — two-phase memory consolidation via AgentRunner."""
import pytest
from unittest.mock import AsyncMock, MagicMock
from nanobot.agent.memory import Dream, MemoryStore
from nanobot.agent.runner import AgentRunResult
@pytest.fixture
def store(tmp_path):
s = MemoryStore(tmp_path)
s.write_soul("# Soul\n- Helpful")
s.write_user("# User\n- Developer")
s.write_memory("# Memory\n- Project X active")
return s
@pytest.fixture
def mock_provider():
p = MagicMock()
p.chat_with_retry = AsyncMock()
return p
@pytest.fixture
def mock_runner():
return MagicMock()
@pytest.fixture
def dream(store, mock_provider, mock_runner):
d = Dream(store=store, provider=mock_provider, model="test-model", max_batch_size=5)
d._runner = mock_runner
return d
def _make_run_result(
stop_reason="completed",
final_content=None,
tool_events=None,
usage=None,
):
return AgentRunResult(
final_content=final_content or stop_reason,
stop_reason=stop_reason,
messages=[],
tools_used=[],
usage={},
tool_events=tool_events or [],
)
class TestDreamRun:
async def test_noop_when_no_unprocessed_history(self, dream, mock_provider, mock_runner, store):
"""Dream should not call LLM when there's nothing to process."""
result = await dream.run()
assert result is False
mock_provider.chat_with_retry.assert_not_called()
mock_runner.run.assert_not_called()
async def test_calls_runner_for_unprocessed_entries(self, dream, mock_provider, mock_runner, store):
"""Dream should call AgentRunner when there are unprocessed history entries."""
store.append_history("User prefers dark mode")
mock_provider.chat_with_retry.return_value = MagicMock(content="New fact")
mock_runner.run = AsyncMock(return_value=_make_run_result(
tool_events=[{"name": "edit_file", "status": "ok", "detail": "memory/MEMORY.md"}],
))
result = await dream.run()
assert result is True
mock_runner.run.assert_called_once()
spec = mock_runner.run.call_args[0][0]
assert spec.max_iterations == 10
assert spec.fail_on_tool_error is True
async def test_advances_dream_cursor(self, dream, mock_provider, mock_runner, store):
"""Dream should advance the cursor after processing."""
store.append_history("event 1")
store.append_history("event 2")
mock_provider.chat_with_retry.return_value = MagicMock(content="Nothing new")
mock_runner.run = AsyncMock(return_value=_make_run_result())
await dream.run()
assert store.get_last_dream_cursor() == 2
async def test_compacts_processed_history(self, dream, mock_provider, mock_runner, store):
"""Dream should compact history after processing."""
store.append_history("event 1")
store.append_history("event 2")
store.append_history("event 3")
mock_provider.chat_with_retry.return_value = MagicMock(content="Nothing new")
mock_runner.run = AsyncMock(return_value=_make_run_result())
await dream.run()
# After Dream, cursor is advanced and 3, compact keeps last max_history_entries
entries = store.read_unprocessed_history(since_cursor=0)
assert all(e["cursor"] > 0 for e in entries)
+229
View File
@@ -0,0 +1,229 @@
"""Tests for GitStore — git-backed version control for memory files."""
import pytest
from pathlib import Path
from nanobot.agent.git_store import GitStore, CommitInfo
TRACKED = ["SOUL.md", "USER.md", "memory/MEMORY.md"]
@pytest.fixture
def git(tmp_path):
"""Uninitialized GitStore."""
return GitStore(tmp_path, tracked_files=TRACKED)
@pytest.fixture
def git_ready(git):
"""Initialized GitStore with one initial commit."""
git.init()
return git
class TestInit:
def test_not_initialized_by_default(self, git, tmp_path):
assert not git.is_initialized()
assert not (tmp_path / ".git").is_dir()
def test_init_creates_git_dir(self, git, tmp_path):
assert git.init()
assert (tmp_path / ".git").is_dir()
def test_init_idempotent(self, git_ready):
assert not git_ready.init()
def test_init_creates_gitignore(self, git_ready):
gi = git_ready._workspace / ".gitignore"
assert gi.exists()
content = gi.read_text(encoding="utf-8")
for f in TRACKED:
assert f"!{f}" in content
def test_init_touches_tracked_files(self, git_ready):
for f in TRACKED:
assert (git_ready._workspace / f).exists()
def test_init_makes_initial_commit(self, git_ready):
commits = git_ready.log()
assert len(commits) == 1
assert "init" in commits[0].message
class TestBuildGitignore:
def test_subdirectory_dirs(self, git):
content = git._build_gitignore()
assert "!memory/\n" in content
for f in TRACKED:
assert f"!{f}\n" in content
assert content.startswith("/*\n")
def test_root_level_files_no_dir_entries(self, tmp_path):
gs = GitStore(tmp_path, tracked_files=["a.md", "b.md"])
content = gs._build_gitignore()
assert "!a.md\n" in content
assert "!b.md\n" in content
dir_lines = [l for l in content.split("\n") if l.startswith("!") and l.endswith("/")]
assert dir_lines == []
class TestAutoCommit:
def test_returns_none_when_not_initialized(self, git):
assert git.auto_commit("test") is None
def test_commits_file_change(self, git_ready):
(git_ready._workspace / "SOUL.md").write_text("updated", encoding="utf-8")
sha = git_ready.auto_commit("update soul")
assert sha is not None
assert len(sha) == 8
def test_returns_none_when_no_changes(self, git_ready):
assert git_ready.auto_commit("no change") is None
def test_commit_appears_in_log(self, git_ready):
ws = git_ready._workspace
(ws / "SOUL.md").write_text("v2", encoding="utf-8")
sha = git_ready.auto_commit("update soul")
commits = git_ready.log()
assert len(commits) == 2
assert commits[0].sha == sha
def test_does_not_create_empty_commits(self, git_ready):
git_ready.auto_commit("nothing 1")
git_ready.auto_commit("nothing 2")
assert len(git_ready.log()) == 1 # only init commit
class TestLog:
def test_empty_when_not_initialized(self, git):
assert git.log() == []
def test_newest_first(self, git_ready):
ws = git_ready._workspace
for i in range(3):
(ws / "SOUL.md").write_text(f"v{i}", encoding="utf-8")
git_ready.auto_commit(f"commit {i}")
commits = git_ready.log()
assert len(commits) == 4 # init + 3
assert "commit 2" in commits[0].message
assert "init" in commits[-1].message
def test_max_entries(self, git_ready):
ws = git_ready._workspace
for i in range(10):
(ws / "SOUL.md").write_text(f"v{i}", encoding="utf-8")
git_ready.auto_commit(f"c{i}")
assert len(git_ready.log(max_entries=3)) == 3
def test_commit_info_fields(self, git_ready):
c = git_ready.log()[0]
assert isinstance(c, CommitInfo)
assert len(c.sha) == 8
assert c.timestamp
assert c.message
class TestDiffCommits:
def test_empty_when_not_initialized(self, git):
assert git.diff_commits("a", "b") == ""
def test_diff_between_two_commits(self, git_ready):
ws = git_ready._workspace
(ws / "SOUL.md").write_text("original", encoding="utf-8")
git_ready.auto_commit("v1")
(ws / "SOUL.md").write_text("modified", encoding="utf-8")
git_ready.auto_commit("v2")
commits = git_ready.log()
diff = git_ready.diff_commits(commits[1].sha, commits[0].sha)
assert "modified" in diff
def test_invalid_sha_returns_empty(self, git_ready):
assert git_ready.diff_commits("deadbeef", "cafebabe") == ""
class TestFindCommit:
def test_finds_by_prefix(self, git_ready):
ws = git_ready._workspace
(ws / "SOUL.md").write_text("v2", encoding="utf-8")
sha = git_ready.auto_commit("v2")
found = git_ready.find_commit(sha[:4])
assert found is not None
assert found.sha == sha
def test_returns_none_for_unknown(self, git_ready):
assert git_ready.find_commit("deadbeef") is None
class TestShowCommitDiff:
def test_returns_commit_with_diff(self, git_ready):
ws = git_ready._workspace
(ws / "SOUL.md").write_text("content", encoding="utf-8")
sha = git_ready.auto_commit("add content")
result = git_ready.show_commit_diff(sha)
assert result is not None
commit, diff = result
assert commit.sha == sha
assert "content" in diff
def test_first_commit_has_empty_diff(self, git_ready):
init_sha = git_ready.log()[-1].sha
result = git_ready.show_commit_diff(init_sha)
assert result is not None
_, diff = result
assert diff == ""
def test_returns_none_for_unknown(self, git_ready):
assert git_ready.show_commit_diff("deadbeef") is None
class TestCommitInfoFormat:
def test_format_with_diff(self):
from nanobot.agent.git_store import CommitInfo
c = CommitInfo(sha="abcd1234", message="test commit\nsecond line", timestamp="2026-04-02 12:00")
result = c.format(diff="some diff")
assert "test commit" in result
assert "`abcd1234`" in result
assert "some diff" in result
def test_format_without_diff(self):
from nanobot.agent.git_store import CommitInfo
c = CommitInfo(sha="abcd1234", message="test", timestamp="2026-04-02 12:00")
result = c.format()
assert "(no file changes)" in result
class TestRevert:
def test_returns_none_when_not_initialized(self, git):
assert git.revert("abc") is None
def test_reverts_file_content(self, git_ready):
ws = git_ready._workspace
(ws / "SOUL.md").write_text("v2 content", encoding="utf-8")
git_ready.auto_commit("v2")
commits = git_ready.log()
new_sha = git_ready.revert(commits[0].sha) # undo v2 → back to init
assert new_sha is not None
assert (ws / "SOUL.md").read_text(encoding="utf-8") == ""
def test_cannot_revert_root_commit(self, git_ready):
commits = git_ready.log()
assert git_ready.revert(commits[-1].sha) is None
def test_invalid_sha_returns_none(self, git_ready):
assert git_ready.revert("deadbeef") is None
class TestMemoryStoreGitProperty:
def test_git_property_exposes_gitstore(self, tmp_path):
from nanobot.agent.memory import MemoryStore
store = MemoryStore(tmp_path)
assert isinstance(store.git, GitStore)
def test_git_property_is_same_object(self, tmp_path):
from nanobot.agent.memory import MemoryStore
store = MemoryStore(tmp_path)
assert store.git is store._git
+2 -1
View File
@@ -249,7 +249,8 @@ def _make_loop(tmp_path, hooks=None):
with patch("nanobot.agent.loop.ContextBuilder"), \
patch("nanobot.agent.loop.SessionManager"), \
patch("nanobot.agent.loop.SubagentManager") as mock_sub_mgr, \
patch("nanobot.agent.loop.MemoryConsolidator"):
patch("nanobot.agent.loop.Consolidator"), \
patch("nanobot.agent.loop.Dream"):
mock_sub_mgr.return_value.cancel_by_session = AsyncMock(return_value=0)
loop = AgentLoop(
bus=bus, provider=provider, workspace=tmp_path, hooks=hooks,
+138
View File
@@ -0,0 +1,138 @@
"""Tests for AgentLoop._dispatch streaming metadata passthrough."""
from unittest.mock import AsyncMock, patch
import pytest
from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.bus.queue import MessageBus
def _make_inbound(**meta) -> InboundMessage:
meta.setdefault("_wants_stream", True)
return InboundMessage(
channel="telegram",
sender_id="user1",
chat_id="chat1",
content="hello",
metadata=meta,
)
@pytest.mark.asyncio
async def test_on_stream_forwards_message_metadata() -> None:
"""on_stream should include original message metadata (e.g. message_thread_id)."""
from nanobot.agent.loop import AgentLoop
bus = MessageBus()
msg = _make_inbound(message_thread_id="42")
loop = AgentLoop.__new__(AgentLoop)
loop.bus = bus
loop._session_locks = {}
loop._concurrency_gate = None
async def fake_process_message(msg_in, **kwargs):
on_stream = kwargs.get("on_stream")
on_stream_end = kwargs.get("on_stream_end")
if on_stream:
await on_stream("hello")
if on_stream_end:
await on_stream_end()
return OutboundMessage(
channel=msg_in.channel, chat_id=msg_in.chat_id,
content="done", metadata=msg_in.metadata,
)
with patch.object(loop, "_process_message", side_effect=fake_process_message):
await loop._dispatch(msg)
# Collect all outbound messages (stream delta, stream end, final response)
outbound: list[OutboundMessage] = []
while not bus.outbound.empty():
outbound.append(await bus.outbound.get())
stream_msg = next(m for m in outbound if m.metadata.get("_stream_delta"))
assert stream_msg.metadata["message_thread_id"] == "42"
assert stream_msg.metadata["_stream_delta"] is True
assert "_stream_id" in stream_msg.metadata
@pytest.mark.asyncio
async def test_on_stream_end_forwards_message_metadata() -> None:
"""on_stream_end should include original message metadata."""
from nanobot.agent.loop import AgentLoop
bus = MessageBus()
msg = _make_inbound(message_thread_id="42")
loop = AgentLoop.__new__(AgentLoop)
loop.bus = bus
loop._session_locks = {}
loop._concurrency_gate = None
async def fake_process_message(msg_in, **kwargs):
on_stream = kwargs.get("on_stream")
on_stream_end = kwargs.get("on_stream_end")
if on_stream:
await on_stream("hello")
if on_stream_end:
await on_stream_end()
return OutboundMessage(
channel=msg_in.channel, chat_id=msg_in.chat_id,
content="done", metadata=msg_in.metadata,
)
with patch.object(loop, "_process_message", side_effect=fake_process_message):
await loop._dispatch(msg)
outbound: list[OutboundMessage] = []
while not bus.outbound.empty():
outbound.append(await bus.outbound.get())
end_msg = next(m for m in outbound if m.metadata.get("_stream_end"))
assert end_msg.metadata["message_thread_id"] == "42"
assert end_msg.metadata["_stream_end"] is True
assert end_msg.metadata["_resuming"] is False
assert "_stream_id" in end_msg.metadata
@pytest.mark.asyncio
async def test_streaming_preserves_arbitrary_metadata_keys() -> None:
"""Both streaming callbacks should forward all original metadata keys untouched."""
from nanobot.agent.loop import AgentLoop
bus = MessageBus()
msg = _make_inbound(message_thread_id="99", custom_flag="abc", reply_to_id="msg77")
loop = AgentLoop.__new__(AgentLoop)
loop.bus = bus
loop._session_locks = {}
loop._concurrency_gate = None
async def fake_process_message(msg_in, **kwargs):
on_stream = kwargs.get("on_stream")
on_stream_end = kwargs.get("on_stream_end")
if on_stream:
await on_stream("hi")
if on_stream_end:
await on_stream_end()
return OutboundMessage(
channel=msg_in.channel, chat_id=msg_in.chat_id,
content="done", metadata=msg_in.metadata,
)
with patch.object(loop, "_process_message", side_effect=fake_process_message):
await loop._dispatch(msg)
outbound: list[OutboundMessage] = []
while not bus.outbound.empty():
outbound.append(await bus.outbound.get())
stream_msg = next(m for m in outbound if m.metadata.get("_stream_delta"))
for key in ("message_thread_id", "custom_flag", "reply_to_id"):
assert stream_msg.metadata[key] == msg.metadata[key]
end_msg = next(m for m in outbound if m.metadata.get("_stream_end"))
for key in ("message_thread_id", "custom_flag", "reply_to_id"):
assert end_msg.metadata[key] == msg.metadata[key]
+18 -18
View File
@@ -26,24 +26,24 @@ def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -
context_window_tokens=context_window_tokens,
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.memory_consolidator._SAFETY_BUFFER = 0
loop.consolidator._SAFETY_BUFFER = 0
return loop
@pytest.mark.asyncio
async def test_prompt_below_threshold_does_not_consolidate(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=100, context_window_tokens=200)
loop.memory_consolidator.consolidate_messages = AsyncMock(return_value=True) # type: ignore[method-assign]
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign]
await loop.process_direct("hello", session_key="cli:test")
loop.memory_consolidator.consolidate_messages.assert_not_awaited()
loop.consolidator.archive.assert_not_awaited()
@pytest.mark.asyncio
async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypatch) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.memory_consolidator.consolidate_messages = AsyncMock(return_value=True) # type: ignore[method-assign]
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
@@ -55,13 +55,13 @@ async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypat
await loop.process_direct("hello", session_key="cli:test")
assert loop.memory_consolidator.consolidate_messages.await_count >= 1
assert loop.consolidator.archive.await_count >= 1
@pytest.mark.asyncio
async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path, monkeypatch) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.memory_consolidator.consolidate_messages = AsyncMock(return_value=True) # type: ignore[method-assign]
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
@@ -76,9 +76,9 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
token_map = {"u1": 120, "a1": 120, "u2": 120, "a2": 120, "u3": 120}
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda message: token_map[message["content"]])
await loop.memory_consolidator.maybe_consolidate_by_tokens(session)
await loop.consolidator.maybe_consolidate_by_tokens(session)
archived_chunk = loop.memory_consolidator.consolidate_messages.await_args.args[0]
archived_chunk = loop.consolidator.archive.await_args.args[0]
assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"]
assert session.last_consolidated == 4
@@ -87,7 +87,7 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> None:
"""Verify maybe_consolidate_by_tokens keeps looping until under threshold."""
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.memory_consolidator.consolidate_messages = AsyncMock(return_value=True) # type: ignore[method-assign]
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
@@ -110,12 +110,12 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No
return (300, "test")
return (80, "test")
loop.memory_consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
await loop.memory_consolidator.maybe_consolidate_by_tokens(session)
await loop.consolidator.maybe_consolidate_by_tokens(session)
assert loop.memory_consolidator.consolidate_messages.await_count == 2
assert loop.consolidator.archive.await_count == 2
assert session.last_consolidated == 6
@@ -123,7 +123,7 @@ async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> No
async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, monkeypatch) -> None:
"""Once triggered, consolidation should continue until it drops below half threshold."""
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.memory_consolidator.consolidate_messages = AsyncMock(return_value=True) # type: ignore[method-assign]
loop.consolidator.archive = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
@@ -147,12 +147,12 @@ async def test_consolidation_continues_below_trigger_until_half_target(tmp_path,
return (150, "test")
return (80, "test")
loop.memory_consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
await loop.memory_consolidator.maybe_consolidate_by_tokens(session)
await loop.consolidator.maybe_consolidate_by_tokens(session)
assert loop.memory_consolidator.consolidate_messages.await_count == 2
assert loop.consolidator.archive.await_count == 2
assert session.last_consolidated == 6
@@ -166,7 +166,7 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
async def track_consolidate(messages):
order.append("consolidate")
return True
loop.memory_consolidator.consolidate_messages = track_consolidate # type: ignore[method-assign]
loop.consolidator.archive = track_consolidate # type: ignore[method-assign]
async def track_llm(*args, **kwargs):
order.append("llm")
@@ -187,7 +187,7 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
def mock_estimate(_session):
call_count[0] += 1
return (1000 if call_count[0] <= 1 else 80, "test")
loop.memory_consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
await loop.process_direct("hello", session_key="cli:test")
@@ -1,478 +0,0 @@
"""Test MemoryStore.consolidate() handles non-string tool call arguments.
Regression test for https://github.com/HKUDS/nanobot/issues/1042
When memory consolidation receives dict values instead of strings from the LLM
tool call response, it should serialize them to JSON instead of raising TypeError.
"""
import json
from pathlib import Path
from unittest.mock import AsyncMock
import pytest
from nanobot.agent.memory import MemoryStore
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
def _make_messages(message_count: int = 30):
"""Create a list of mock messages."""
return [
{"role": "user", "content": f"msg{i}", "timestamp": "2026-01-01 00:00"}
for i in range(message_count)
]
def _make_tool_response(history_entry, memory_update):
"""Create an LLMResponse with a save_memory tool call."""
return LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="save_memory",
arguments={
"history_entry": history_entry,
"memory_update": memory_update,
},
)
],
)
class ScriptedProvider(LLMProvider):
def __init__(self, responses: list[LLMResponse]):
super().__init__()
self._responses = list(responses)
self.calls = 0
async def chat(self, *args, **kwargs) -> LLMResponse:
self.calls += 1
if self._responses:
return self._responses.pop(0)
return LLMResponse(content="", tool_calls=[])
def get_default_model(self) -> str:
return "test-model"
class TestMemoryConsolidationTypeHandling:
"""Test that consolidation handles various argument types correctly."""
@pytest.mark.asyncio
async def test_string_arguments_work(self, tmp_path: Path) -> None:
"""Normal case: LLM returns string arguments."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat = AsyncMock(
return_value=_make_tool_response(
history_entry="[2026-01-01] User discussed testing.",
memory_update="# Memory\nUser likes testing.",
)
)
provider.chat_with_retry = provider.chat
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is True
assert store.history_file.exists()
assert "[2026-01-01] User discussed testing." in store.history_file.read_text()
assert "User likes testing." in store.memory_file.read_text()
@pytest.mark.asyncio
async def test_dict_arguments_serialized_to_json(self, tmp_path: Path) -> None:
"""Issue #1042: LLM returns dict instead of string — must not raise TypeError."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat = AsyncMock(
return_value=_make_tool_response(
history_entry={"timestamp": "2026-01-01", "summary": "User discussed testing."},
memory_update={"facts": ["User likes testing"], "topics": ["testing"]},
)
)
provider.chat_with_retry = provider.chat
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is True
assert store.history_file.exists()
history_content = store.history_file.read_text()
parsed = json.loads(history_content.strip())
assert parsed["summary"] == "User discussed testing."
memory_content = store.memory_file.read_text()
parsed_mem = json.loads(memory_content)
assert "User likes testing" in parsed_mem["facts"]
@pytest.mark.asyncio
async def test_string_arguments_as_raw_json(self, tmp_path: Path) -> None:
"""Some providers return arguments as a JSON string instead of parsed dict."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
response = LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="save_memory",
arguments=json.dumps({
"history_entry": "[2026-01-01] User discussed testing.",
"memory_update": "# Memory\nUser likes testing.",
}),
)
],
)
provider.chat = AsyncMock(return_value=response)
provider.chat_with_retry = provider.chat
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is True
assert "User discussed testing." in store.history_file.read_text()
@pytest.mark.asyncio
async def test_no_tool_call_returns_false(self, tmp_path: Path) -> None:
"""When LLM doesn't use the save_memory tool, return False."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat = AsyncMock(
return_value=LLMResponse(content="I summarized the conversation.", tool_calls=[])
)
provider.chat_with_retry = provider.chat
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is False
assert not store.history_file.exists()
@pytest.mark.asyncio
async def test_skips_when_message_chunk_is_empty(self, tmp_path: Path) -> None:
"""Consolidation should be a no-op when the selected chunk is empty."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat_with_retry = provider.chat
messages: list[dict] = []
result = await store.consolidate(messages, provider, "test-model")
assert result is True
provider.chat.assert_not_called()
@pytest.mark.asyncio
async def test_list_arguments_extracts_first_dict(self, tmp_path: Path) -> None:
"""Some providers return arguments as a list - extract first element if it's a dict."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
response = LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="save_memory",
arguments=[{
"history_entry": "[2026-01-01] User discussed testing.",
"memory_update": "# Memory\nUser likes testing.",
}],
)
],
)
provider.chat = AsyncMock(return_value=response)
provider.chat_with_retry = provider.chat
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is True
assert "User discussed testing." in store.history_file.read_text()
assert "User likes testing." in store.memory_file.read_text()
@pytest.mark.asyncio
async def test_list_arguments_empty_list_returns_false(self, tmp_path: Path) -> None:
"""Empty list arguments should return False."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
response = LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="save_memory",
arguments=[],
)
],
)
provider.chat = AsyncMock(return_value=response)
provider.chat_with_retry = provider.chat
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is False
@pytest.mark.asyncio
async def test_list_arguments_non_dict_content_returns_false(self, tmp_path: Path) -> None:
"""List with non-dict content should return False."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
response = LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="save_memory",
arguments=["string", "content"],
)
],
)
provider.chat = AsyncMock(return_value=response)
provider.chat_with_retry = provider.chat
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is False
@pytest.mark.asyncio
async def test_missing_history_entry_returns_false_without_writing(self, tmp_path: Path) -> None:
"""Do not persist partial results when required fields are missing."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="save_memory",
arguments={"memory_update": "# Memory\nOnly memory update"},
)
],
)
)
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is False
assert not store.history_file.exists()
assert not store.memory_file.exists()
@pytest.mark.asyncio
async def test_missing_memory_update_returns_false_without_writing(self, tmp_path: Path) -> None:
"""Do not append history if memory_update is missing."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(
content=None,
tool_calls=[
ToolCallRequest(
id="call_1",
name="save_memory",
arguments={"history_entry": "[2026-01-01] Partial output."},
)
],
)
)
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is False
assert not store.history_file.exists()
assert not store.memory_file.exists()
@pytest.mark.asyncio
async def test_null_required_field_returns_false_without_writing(self, tmp_path: Path) -> None:
"""Null required fields should be rejected before persistence."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(
return_value=_make_tool_response(
history_entry=None,
memory_update="# Memory\nUser likes testing.",
)
)
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is False
assert not store.history_file.exists()
assert not store.memory_file.exists()
@pytest.mark.asyncio
async def test_empty_history_entry_returns_false_without_writing(self, tmp_path: Path) -> None:
"""Empty history entries should be rejected to avoid blank archival records."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(
return_value=_make_tool_response(
history_entry=" ",
memory_update="# Memory\nUser likes testing.",
)
)
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is False
assert not store.history_file.exists()
assert not store.memory_file.exists()
@pytest.mark.asyncio
async def test_retries_transient_error_then_succeeds(self, tmp_path: Path, monkeypatch) -> None:
store = MemoryStore(tmp_path)
provider = ScriptedProvider([
LLMResponse(content="503 server error", finish_reason="error"),
_make_tool_response(
history_entry="[2026-01-01] User discussed testing.",
memory_update="# Memory\nUser likes testing.",
),
])
messages = _make_messages(message_count=60)
delays: list[int] = []
async def _fake_sleep(delay: int) -> None:
delays.append(delay)
monkeypatch.setattr("nanobot.providers.base.asyncio.sleep", _fake_sleep)
result = await store.consolidate(messages, provider, "test-model")
assert result is True
assert provider.calls == 2
assert delays == [1]
@pytest.mark.asyncio
async def test_consolidation_delegates_to_provider_defaults(self, tmp_path: Path) -> None:
"""Consolidation no longer passes generation params — the provider owns them."""
store = MemoryStore(tmp_path)
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(
return_value=_make_tool_response(
history_entry="[2026-01-01] User discussed testing.",
memory_update="# Memory\nUser likes testing.",
)
)
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is True
provider.chat_with_retry.assert_awaited_once()
_, kwargs = provider.chat_with_retry.await_args
assert kwargs["model"] == "test-model"
assert "temperature" not in kwargs
assert "max_tokens" not in kwargs
assert "reasoning_effort" not in kwargs
@pytest.mark.asyncio
async def test_tool_choice_fallback_on_unsupported_error(self, tmp_path: Path) -> None:
"""Forced tool_choice rejected by provider -> retry with auto and succeed."""
store = MemoryStore(tmp_path)
error_resp = LLMResponse(
content="Error calling LLM: BadRequestError: "
"The tool_choice parameter does not support being set to required or object",
finish_reason="error",
tool_calls=[],
)
ok_resp = _make_tool_response(
history_entry="[2026-01-01] Fallback worked.",
memory_update="# Memory\nFallback OK.",
)
call_log: list[dict] = []
async def _tracking_chat(**kwargs):
call_log.append(kwargs)
return error_resp if len(call_log) == 1 else ok_resp
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(side_effect=_tracking_chat)
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is True
assert len(call_log) == 2
assert isinstance(call_log[0]["tool_choice"], dict)
assert call_log[1]["tool_choice"] == "auto"
assert "Fallback worked." in store.history_file.read_text()
@pytest.mark.asyncio
async def test_tool_choice_fallback_auto_no_tool_call(self, tmp_path: Path) -> None:
"""Forced rejected, auto retry also produces no tool call -> return False."""
store = MemoryStore(tmp_path)
error_resp = LLMResponse(
content="Error: tool_choice must be none or auto",
finish_reason="error",
tool_calls=[],
)
no_tool_resp = LLMResponse(
content="Here is a summary.",
finish_reason="stop",
tool_calls=[],
)
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(side_effect=[error_resp, no_tool_resp])
messages = _make_messages(message_count=60)
result = await store.consolidate(messages, provider, "test-model")
assert result is False
assert not store.history_file.exists()
@pytest.mark.asyncio
async def test_raw_archive_after_consecutive_failures(self, tmp_path: Path) -> None:
"""After 3 consecutive failures, raw-archive messages and return True."""
store = MemoryStore(tmp_path)
no_tool = LLMResponse(content="No tool call.", finish_reason="stop", tool_calls=[])
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(return_value=no_tool)
messages = _make_messages(message_count=10)
assert await store.consolidate(messages, provider, "m") is False
assert await store.consolidate(messages, provider, "m") is False
assert await store.consolidate(messages, provider, "m") is True
assert store.history_file.exists()
content = store.history_file.read_text()
assert "[RAW]" in content
assert "10 messages" in content
assert "msg0" in content
assert not store.memory_file.exists()
@pytest.mark.asyncio
async def test_raw_archive_counter_resets_on_success(self, tmp_path: Path) -> None:
"""A successful consolidation resets the failure counter."""
store = MemoryStore(tmp_path)
no_tool = LLMResponse(content="Nope.", finish_reason="stop", tool_calls=[])
ok_resp = _make_tool_response(
history_entry="[2026-01-01] OK.",
memory_update="# Memory\nOK.",
)
messages = _make_messages(message_count=10)
provider = AsyncMock()
provider.chat_with_retry = AsyncMock(return_value=no_tool)
assert await store.consolidate(messages, provider, "m") is False
assert await store.consolidate(messages, provider, "m") is False
assert store._consecutive_failures == 2
provider.chat_with_retry = AsyncMock(return_value=ok_resp)
assert await store.consolidate(messages, provider, "m") is True
assert store._consecutive_failures == 0
provider.chat_with_retry = AsyncMock(return_value=no_tool)
assert await store.consolidate(messages, provider, "m") is False
assert store._consecutive_failures == 1
+133
View File
@@ -0,0 +1,133 @@
"""Tests for the restructured MemoryStore — pure file I/O layer."""
import json
import pytest
from pathlib import Path
from nanobot.agent.memory import MemoryStore
@pytest.fixture
def store(tmp_path):
return MemoryStore(tmp_path)
class TestMemoryStoreBasicIO:
def test_read_memory_returns_empty_when_missing(self, store):
assert store.read_memory() == ""
def test_write_and_read_memory(self, store):
store.write_memory("hello")
assert store.read_memory() == "hello"
def test_read_soul_returns_empty_when_missing(self, store):
assert store.read_soul() == ""
def test_write_and_read_soul(self, store):
store.write_soul("soul content")
assert store.read_soul() == "soul content"
def test_read_user_returns_empty_when_missing(self, store):
assert store.read_user() == ""
def test_write_and_read_user(self, store):
store.write_user("user content")
assert store.read_user() == "user content"
def test_get_memory_context_returns_empty_when_missing(self, store):
assert store.get_memory_context() == ""
def test_get_memory_context_returns_formatted_content(self, store):
store.write_memory("important fact")
ctx = store.get_memory_context()
assert "Long-term Memory" in ctx
assert "important fact" in ctx
class TestHistoryWithCursor:
def test_append_history_returns_cursor(self, store):
cursor = store.append_history("event 1")
assert cursor == 1
cursor2 = store.append_history("event 2")
assert cursor2 == 2
def test_append_history_includes_cursor_in_file(self, store):
store.append_history("event 1")
content = store.read_file(store.history_file)
data = json.loads(content)
assert data["cursor"] == 1
def test_cursor_persists_across_appends(self, store):
store.append_history("event 1")
store.append_history("event 2")
cursor = store.append_history("event 3")
assert cursor == 3
def test_read_unprocessed_history(self, store):
store.append_history("event 1")
store.append_history("event 2")
store.append_history("event 3")
entries = store.read_unprocessed_history(since_cursor=1)
assert len(entries) == 2
assert entries[0]["cursor"] == 2
def test_read_unprocessed_history_returns_all_when_cursor_zero(self, store):
store.append_history("event 1")
store.append_history("event 2")
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 2
def test_compact_history_drops_oldest(self, tmp_path):
store = MemoryStore(tmp_path, max_history_entries=2)
store.append_history("event 1")
store.append_history("event 2")
store.append_history("event 3")
store.append_history("event 4")
store.append_history("event 5")
store.compact_history()
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 2
assert entries[0]["cursor"] in {4, 5}
class TestDreamCursor:
def test_initial_cursor_is_zero(self, store):
assert store.get_last_dream_cursor() == 0
def test_set_and_get_cursor(self, store):
store.set_last_dream_cursor(5)
assert store.get_last_dream_cursor() == 5
def test_cursor_persists(self, store):
store.set_last_dream_cursor(3)
store2 = MemoryStore(store.workspace)
assert store2.get_last_dream_cursor() == 3
class TestDreamLog:
def test_read_dream_log_returns_empty_when_missing(self, store):
assert store.read_dream_log() == ""
def test_append_dream_log(self, store):
store.append_dream_log("## 2026-03-30\nProcessed entries #1-#5")
log = store.read_dream_log()
assert "Processed entries #1-#5" in log
def test_append_dream_log_is_additive(self, store):
store.append_dream_log("first run")
store.append_dream_log("second run")
log = store.read_dream_log()
assert "first run" in log
assert "second run" in log
class TestLegacyHistoryMigration:
def test_read_unprocessed_history_handles_entries_without_cursor(self, store):
"""JSONL entries with cursor=1 are correctly parsed and returned."""
store.history_file.write_text(
'{"cursor": 1, "timestamp": "2026-03-30 14:30", "content": "Old event"}\n',
encoding="utf-8")
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1
assert entries[0]["cursor"] == 1
+79
View File
@@ -333,3 +333,82 @@ async def test_subagent_max_iterations_announces_existing_fallback(tmp_path, mon
args = mgr._announce_result.await_args.args
assert args[3] == "Task completed but no final response was generated."
assert args[5] == "ok"
@pytest.mark.asyncio
async def test_runner_accumulates_usage_and_preserves_cached_tokens():
"""Runner should accumulate prompt/completion tokens across iterations
and preserve cached_tokens from provider responses."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner
provider = MagicMock()
call_count = {"n": 0}
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
return LLMResponse(
content="thinking",
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})],
usage={"prompt_tokens": 100, "completion_tokens": 10, "cached_tokens": 80},
)
return LLMResponse(
content="done",
tool_calls=[],
usage={"prompt_tokens": 200, "completion_tokens": 20, "cached_tokens": 150},
)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
tools.execute = AsyncMock(return_value="file content")
runner = AgentRunner(provider)
result = await runner.run(AgentRunSpec(
initial_messages=[{"role": "user", "content": "do task"}],
tools=tools,
model="test-model",
max_iterations=3,
))
# Usage should be accumulated across iterations
assert result.usage["prompt_tokens"] == 300 # 100 + 200
assert result.usage["completion_tokens"] == 30 # 10 + 20
assert result.usage["cached_tokens"] == 230 # 80 + 150
@pytest.mark.asyncio
async def test_runner_passes_cached_tokens_to_hook_context():
"""Hook context.usage should contain cached_tokens."""
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.runner import AgentRunSpec, AgentRunner
provider = MagicMock()
captured_usage: list[dict] = []
class UsageHook(AgentHook):
async def after_iteration(self, context: AgentHookContext) -> None:
captured_usage.append(dict(context.usage))
async def chat_with_retry(**kwargs):
return LLMResponse(
content="done",
tool_calls=[],
usage={"prompt_tokens": 200, "completion_tokens": 20, "cached_tokens": 150},
)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
runner = AgentRunner(provider)
await runner.run(AgentRunSpec(
initial_messages=[],
tools=tools,
model="test-model",
max_iterations=1,
hook=UsageHook(),
))
assert len(captured_usage) == 1
assert captured_usage[0]["cached_tokens"] == 150
-37
View File
@@ -117,43 +117,6 @@ class TestDispatch:
out = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
assert out.content == "hi"
@pytest.mark.asyncio
async def test_dispatch_streaming_preserves_message_metadata(self):
from nanobot.bus.events import InboundMessage
loop, bus = _make_loop()
msg = InboundMessage(
channel="matrix",
sender_id="u1",
chat_id="!room:matrix.org",
content="hello",
metadata={
"_wants_stream": True,
"thread_root_event_id": "$root1",
"thread_reply_to_event_id": "$reply1",
},
)
async def fake_process(_msg, *, on_stream=None, on_stream_end=None, **kwargs):
assert on_stream is not None
assert on_stream_end is not None
await on_stream("hi")
await on_stream_end(resuming=False)
return None
loop._process_message = fake_process
await loop._dispatch(msg)
first = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
second = await asyncio.wait_for(bus.consume_outbound(), timeout=1.0)
assert first.metadata["thread_root_event_id"] == "$root1"
assert first.metadata["thread_reply_to_event_id"] == "$reply1"
assert first.metadata["_stream_delta"] is True
assert second.metadata["thread_root_event_id"] == "$root1"
assert second.metadata["thread_reply_to_event_id"] == "$reply1"
assert second.metadata["_stream_end"] is True
@pytest.mark.asyncio
async def test_processing_lock_serializes(self):
from nanobot.bus.events import InboundMessage, OutboundMessage
+1 -1
View File
@@ -4,8 +4,8 @@ import asyncio
from pathlib import Path
from types import SimpleNamespace
import discord
import pytest
discord = pytest.importorskip("discord")
from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus
+238
View File
@@ -0,0 +1,238 @@
"""Tests for Feishu reaction add/remove and auto-cleanup on stream end."""
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from nanobot.bus.queue import MessageBus
from nanobot.channels.feishu import FeishuChannel, FeishuConfig, _FeishuStreamBuf
def _make_channel() -> FeishuChannel:
config = FeishuConfig(
enabled=True,
app_id="cli_test",
app_secret="secret",
allow_from=["*"],
)
ch = FeishuChannel(config, MessageBus())
ch._client = MagicMock()
ch._loop = None
return ch
def _mock_reaction_create_response(reaction_id: str = "reaction_001", success: bool = True):
resp = MagicMock()
resp.success.return_value = success
resp.code = 0 if success else 99999
resp.msg = "ok" if success else "error"
if success:
resp.data = SimpleNamespace(reaction_id=reaction_id)
else:
resp.data = None
return resp
# ── _add_reaction_sync ──────────────────────────────────────────────────────
class TestAddReactionSync:
def test_returns_reaction_id_on_success(self):
ch = _make_channel()
ch._client.im.v1.message_reaction.create.return_value = _mock_reaction_create_response("rx_42")
result = ch._add_reaction_sync("om_001", "THUMBSUP")
assert result == "rx_42"
def test_returns_none_when_response_fails(self):
ch = _make_channel()
ch._client.im.v1.message_reaction.create.return_value = _mock_reaction_create_response(success=False)
assert ch._add_reaction_sync("om_001", "THUMBSUP") is None
def test_returns_none_when_response_data_is_none(self):
ch = _make_channel()
resp = MagicMock()
resp.success.return_value = True
resp.data = None
ch._client.im.v1.message_reaction.create.return_value = resp
assert ch._add_reaction_sync("om_001", "THUMBSUP") is None
def test_returns_none_on_exception(self):
ch = _make_channel()
ch._client.im.v1.message_reaction.create.side_effect = RuntimeError("network error")
assert ch._add_reaction_sync("om_001", "THUMBSUP") is None
# ── _add_reaction (async) ───────────────────────────────────────────────────
class TestAddReactionAsync:
@pytest.mark.asyncio
async def test_returns_reaction_id(self):
ch = _make_channel()
ch._add_reaction_sync = MagicMock(return_value="rx_99")
result = await ch._add_reaction("om_001", "EYES")
assert result == "rx_99"
@pytest.mark.asyncio
async def test_returns_none_when_no_client(self):
ch = _make_channel()
ch._client = None
result = await ch._add_reaction("om_001", "THUMBSUP")
assert result is None
# ── _remove_reaction_sync ───────────────────────────────────────────────────
class TestRemoveReactionSync:
def test_calls_delete_on_success(self):
ch = _make_channel()
resp = MagicMock()
resp.success.return_value = True
ch._client.im.v1.message_reaction.delete.return_value = resp
ch._remove_reaction_sync("om_001", "rx_42")
ch._client.im.v1.message_reaction.delete.assert_called_once()
def test_handles_failure_gracefully(self):
ch = _make_channel()
resp = MagicMock()
resp.success.return_value = False
resp.code = 99999
resp.msg = "not found"
ch._client.im.v1.message_reaction.delete.return_value = resp
# Should not raise
ch._remove_reaction_sync("om_001", "rx_42")
def test_handles_exception_gracefully(self):
ch = _make_channel()
ch._client.im.v1.message_reaction.delete.side_effect = RuntimeError("network error")
# Should not raise
ch._remove_reaction_sync("om_001", "rx_42")
# ── _remove_reaction (async) ────────────────────────────────────────────────
class TestRemoveReactionAsync:
@pytest.mark.asyncio
async def test_calls_sync_helper(self):
ch = _make_channel()
ch._remove_reaction_sync = MagicMock()
await ch._remove_reaction("om_001", "rx_42")
ch._remove_reaction_sync.assert_called_once_with("om_001", "rx_42")
@pytest.mark.asyncio
async def test_noop_when_no_client(self):
ch = _make_channel()
ch._client = None
ch._remove_reaction_sync = MagicMock()
await ch._remove_reaction("om_001", "rx_42")
ch._remove_reaction_sync.assert_not_called()
@pytest.mark.asyncio
async def test_noop_when_reaction_id_is_empty(self):
ch = _make_channel()
ch._remove_reaction_sync = MagicMock()
await ch._remove_reaction("om_001", "")
ch._remove_reaction_sync.assert_not_called()
@pytest.mark.asyncio
async def test_noop_when_reaction_id_is_none(self):
ch = _make_channel()
ch._remove_reaction_sync = MagicMock()
await ch._remove_reaction("om_001", None)
ch._remove_reaction_sync.assert_not_called()
# ── send_delta stream end: reaction auto-cleanup ────────────────────────────
class TestStreamEndReactionCleanup:
@pytest.mark.asyncio
async def test_removes_reaction_on_stream_end(self):
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Done", card_id="card_1", sequence=3, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = MagicMock(success=MagicMock(return_value=True))
ch._client.cardkit.v1.card.settings.return_value = MagicMock(success=MagicMock(return_value=True))
ch._remove_reaction = AsyncMock()
await ch.send_delta(
"oc_chat1", "",
metadata={"_stream_end": True, "message_id": "om_001", "reaction_id": "rx_42"},
)
ch._remove_reaction.assert_called_once_with("om_001", "rx_42")
@pytest.mark.asyncio
async def test_no_removal_when_message_id_missing(self):
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Done", card_id="card_1", sequence=3, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = MagicMock(success=MagicMock(return_value=True))
ch._client.cardkit.v1.card.settings.return_value = MagicMock(success=MagicMock(return_value=True))
ch._remove_reaction = AsyncMock()
await ch.send_delta(
"oc_chat1", "",
metadata={"_stream_end": True, "reaction_id": "rx_42"},
)
ch._remove_reaction.assert_not_called()
@pytest.mark.asyncio
async def test_no_removal_when_reaction_id_missing(self):
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Done", card_id="card_1", sequence=3, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = MagicMock(success=MagicMock(return_value=True))
ch._client.cardkit.v1.card.settings.return_value = MagicMock(success=MagicMock(return_value=True))
ch._remove_reaction = AsyncMock()
await ch.send_delta(
"oc_chat1", "",
metadata={"_stream_end": True, "message_id": "om_001"},
)
ch._remove_reaction.assert_not_called()
@pytest.mark.asyncio
async def test_no_removal_when_both_ids_missing(self):
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Done", card_id="card_1", sequence=3, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = MagicMock(success=MagicMock(return_value=True))
ch._client.cardkit.v1.card.settings.return_value = MagicMock(success=MagicMock(return_value=True))
ch._remove_reaction = AsyncMock()
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True})
ch._remove_reaction.assert_not_called()
@pytest.mark.asyncio
async def test_no_removal_when_not_stream_end(self):
ch = _make_channel()
ch._remove_reaction = AsyncMock()
await ch.send_delta(
"oc_chat1", "more text",
metadata={"message_id": "om_001", "reaction_id": "rx_42"},
)
ch._remove_reaction.assert_not_called()
-82
View File
@@ -1367,23 +1367,6 @@ def test_build_matrix_text_content_with_event_id() -> None:
assert result["m.relates_to"]["event_id"] == event_id
def test_build_matrix_text_content_with_event_id_preserves_thread_relation() -> None:
"""Thread relations for edits should stay inside m.new_content."""
relates_to = {
"rel_type": "m.thread",
"event_id": "$root1",
"m.in_reply_to": {"event_id": "$reply1"},
"is_falling_back": True,
}
result = _build_matrix_text_content("Updated message", "event-1", relates_to)
assert result["m.relates_to"] == {
"rel_type": "m.replace",
"event_id": "event-1",
}
assert result["m.new_content"]["m.relates_to"] == relates_to
def test_build_matrix_text_content_no_event_id() -> None:
"""Test that when event_id is not provided, no extra properties are added."""
result = _build_matrix_text_content("Regular message")
@@ -1517,71 +1500,6 @@ async def test_send_delta_stream_end_replaces_existing_message() -> None:
}
@pytest.mark.asyncio
async def test_send_delta_starts_threaded_stream_inside_thread() -> None:
channel = MatrixChannel(_make_config(), MessageBus())
client = _FakeAsyncClient("", "", "", None)
channel.client = client
client.room_send_response.event_id = "event-1"
metadata = {
"thread_root_event_id": "$root1",
"thread_reply_to_event_id": "$reply1",
}
await channel.send_delta("!room:matrix.org", "Hello", metadata)
assert client.room_send_calls[0]["content"]["m.relates_to"] == {
"rel_type": "m.thread",
"event_id": "$root1",
"m.in_reply_to": {"event_id": "$reply1"},
"is_falling_back": True,
}
@pytest.mark.asyncio
async def test_send_delta_threaded_edit_keeps_replace_and_thread_relation(monkeypatch) -> None:
channel = MatrixChannel(_make_config(), MessageBus())
client = _FakeAsyncClient("", "", "", None)
channel.client = client
client.room_send_response.event_id = "event-1"
times = [100.0, 102.0, 104.0]
times.reverse()
monkeypatch.setattr(channel, "monotonic_time", lambda: times and times.pop())
metadata = {
"thread_root_event_id": "$root1",
"thread_reply_to_event_id": "$reply1",
}
await channel.send_delta("!room:matrix.org", "Hello", metadata)
await channel.send_delta("!room:matrix.org", " world", metadata)
await channel.send_delta("!room:matrix.org", "", {"_stream_end": True, **metadata})
edit_content = client.room_send_calls[1]["content"]
final_content = client.room_send_calls[2]["content"]
assert edit_content["m.relates_to"] == {
"rel_type": "m.replace",
"event_id": "event-1",
}
assert edit_content["m.new_content"]["m.relates_to"] == {
"rel_type": "m.thread",
"event_id": "$root1",
"m.in_reply_to": {"event_id": "$reply1"},
"is_falling_back": True,
}
assert final_content["m.relates_to"] == {
"rel_type": "m.replace",
"event_id": "event-1",
}
assert final_content["m.new_content"]["m.relates_to"] == {
"rel_type": "m.thread",
"event_id": "$root1",
"m.in_reply_to": {"event_id": "$reply1"},
"is_falling_back": True,
}
@pytest.mark.asyncio
async def test_send_delta_stream_end_noop_when_buffer_missing() -> None:
channel = MatrixChannel(_make_config(), MessageBus())
+172
View File
@@ -0,0 +1,172 @@
"""Tests for QQ channel ack_message feature.
Covers the four verification points from the PR:
1. C2C message: ack appears instantly
2. Group message: ack appears instantly
3. ack_message set to "": no ack sent
4. Custom ack_message text: correct text delivered
Each test also verifies that normal message processing is not blocked.
"""
from types import SimpleNamespace
import pytest
try:
from nanobot.channels import qq
QQ_AVAILABLE = getattr(qq, "QQ_AVAILABLE", False)
except ImportError:
QQ_AVAILABLE = False
if not QQ_AVAILABLE:
pytest.skip("QQ dependencies not installed (qq-botpy)", allow_module_level=True)
from nanobot.bus.queue import MessageBus
from nanobot.channels.qq import QQChannel, QQConfig
class _FakeApi:
def __init__(self) -> None:
self.c2c_calls: list[dict] = []
self.group_calls: list[dict] = []
async def post_c2c_message(self, **kwargs) -> None:
self.c2c_calls.append(kwargs)
async def post_group_message(self, **kwargs) -> None:
self.group_calls.append(kwargs)
class _FakeClient:
def __init__(self) -> None:
self.api = _FakeApi()
@pytest.mark.asyncio
async def test_ack_sent_on_c2c_message() -> None:
"""Ack is sent immediately for C2C messages, then normal processing continues."""
channel = QQChannel(
QQConfig(
app_id="app",
secret="secret",
allow_from=["*"],
ack_message="⏳ Processing...",
),
MessageBus(),
)
channel._client = _FakeClient()
data = SimpleNamespace(
id="msg1",
content="hello",
author=SimpleNamespace(user_openid="user1"),
attachments=[],
)
await channel._on_message(data, is_group=False)
assert len(channel._client.api.c2c_calls) >= 1
ack_call = channel._client.api.c2c_calls[0]
assert ack_call["content"] == "⏳ Processing..."
assert ack_call["openid"] == "user1"
assert ack_call["msg_id"] == "msg1"
assert ack_call["msg_type"] == 0
msg = await channel.bus.consume_inbound()
assert msg.content == "hello"
assert msg.sender_id == "user1"
@pytest.mark.asyncio
async def test_ack_sent_on_group_message() -> None:
"""Ack is sent immediately for group messages, then normal processing continues."""
channel = QQChannel(
QQConfig(
app_id="app",
secret="secret",
allow_from=["*"],
ack_message="⏳ Processing...",
),
MessageBus(),
)
channel._client = _FakeClient()
data = SimpleNamespace(
id="msg2",
content="hello group",
group_openid="group123",
author=SimpleNamespace(member_openid="user1"),
attachments=[],
)
await channel._on_message(data, is_group=True)
assert len(channel._client.api.group_calls) >= 1
ack_call = channel._client.api.group_calls[0]
assert ack_call["content"] == "⏳ Processing..."
assert ack_call["group_openid"] == "group123"
assert ack_call["msg_id"] == "msg2"
assert ack_call["msg_type"] == 0
msg = await channel.bus.consume_inbound()
assert msg.content == "hello group"
assert msg.chat_id == "group123"
@pytest.mark.asyncio
async def test_no_ack_when_ack_message_empty() -> None:
"""Setting ack_message to empty string disables the ack entirely."""
channel = QQChannel(
QQConfig(
app_id="app",
secret="secret",
allow_from=["*"],
ack_message="",
),
MessageBus(),
)
channel._client = _FakeClient()
data = SimpleNamespace(
id="msg3",
content="hello",
author=SimpleNamespace(user_openid="user1"),
attachments=[],
)
await channel._on_message(data, is_group=False)
assert len(channel._client.api.c2c_calls) == 0
assert len(channel._client.api.group_calls) == 0
msg = await channel.bus.consume_inbound()
assert msg.content == "hello"
@pytest.mark.asyncio
async def test_custom_ack_message_text() -> None:
"""Custom Chinese ack_message text is delivered correctly."""
custom = "正在处理中,请稍候..."
channel = QQChannel(
QQConfig(
app_id="app",
secret="secret",
allow_from=["*"],
ack_message=custom,
),
MessageBus(),
)
channel._client = _FakeClient()
data = SimpleNamespace(
id="msg4",
content="test input",
author=SimpleNamespace(user_openid="user1"),
attachments=[],
)
await channel._on_message(data, is_group=False)
assert len(channel._client.api.c2c_calls) >= 1
ack_call = channel._client.api.c2c_calls[0]
assert ack_call["content"] == custom
msg = await channel.bus.consume_inbound()
assert msg.content == "test input"
+26 -13
View File
@@ -647,43 +647,56 @@ async def test_group_policy_open_accepts_plain_group_message() -> None:
assert channel._app.bot.get_me_calls == 0
def test_extract_reply_context_no_reply() -> None:
@pytest.mark.asyncio
async def test_extract_reply_context_no_reply() -> None:
"""When there is no reply_to_message, _extract_reply_context returns None."""
channel = TelegramChannel(TelegramConfig(enabled=True, token="123:abc"), MessageBus())
message = SimpleNamespace(reply_to_message=None)
assert TelegramChannel._extract_reply_context(message) is None
assert await channel._extract_reply_context(message) is None
def test_extract_reply_context_with_text() -> None:
@pytest.mark.asyncio
async def test_extract_reply_context_with_text() -> None:
"""When reply has text, return prefixed string."""
reply = SimpleNamespace(text="Hello world", caption=None)
channel = TelegramChannel(TelegramConfig(enabled=True, token="123:abc"), MessageBus())
channel._app = _FakeApp(lambda: None)
reply = SimpleNamespace(text="Hello world", caption=None, from_user=SimpleNamespace(id=2, username="testuser", first_name="Test"))
message = SimpleNamespace(reply_to_message=reply)
assert TelegramChannel._extract_reply_context(message) == "[Reply to: Hello world]"
assert await channel._extract_reply_context(message) == "[Reply to @testuser: Hello world]"
def test_extract_reply_context_with_caption_only() -> None:
@pytest.mark.asyncio
async def test_extract_reply_context_with_caption_only() -> None:
"""When reply has only caption (no text), caption is used."""
reply = SimpleNamespace(text=None, caption="Photo caption")
channel = TelegramChannel(TelegramConfig(enabled=True, token="123:abc"), MessageBus())
channel._app = _FakeApp(lambda: None)
reply = SimpleNamespace(text=None, caption="Photo caption", from_user=SimpleNamespace(id=2, username=None, first_name="Test"))
message = SimpleNamespace(reply_to_message=reply)
assert TelegramChannel._extract_reply_context(message) == "[Reply to: Photo caption]"
assert await channel._extract_reply_context(message) == "[Reply to Test: Photo caption]"
def test_extract_reply_context_truncation() -> None:
@pytest.mark.asyncio
async def test_extract_reply_context_truncation() -> None:
"""Reply text is truncated at TELEGRAM_REPLY_CONTEXT_MAX_LEN."""
channel = TelegramChannel(TelegramConfig(enabled=True, token="123:abc"), MessageBus())
channel._app = _FakeApp(lambda: None)
long_text = "x" * (TELEGRAM_REPLY_CONTEXT_MAX_LEN + 100)
reply = SimpleNamespace(text=long_text, caption=None)
reply = SimpleNamespace(text=long_text, caption=None, from_user=SimpleNamespace(id=2, username=None, first_name=None))
message = SimpleNamespace(reply_to_message=reply)
result = TelegramChannel._extract_reply_context(message)
result = await channel._extract_reply_context(message)
assert result is not None
assert result.startswith("[Reply to: ")
assert result.endswith("...]")
assert len(result) == len("[Reply to: ]") + TELEGRAM_REPLY_CONTEXT_MAX_LEN + len("...")
def test_extract_reply_context_no_text_returns_none() -> None:
@pytest.mark.asyncio
async def test_extract_reply_context_no_text_returns_none() -> None:
"""When reply has no text/caption, _extract_reply_context returns None (media handled separately)."""
channel = TelegramChannel(TelegramConfig(enabled=True, token="123:abc"), MessageBus())
reply = SimpleNamespace(text=None, caption=None)
message = SimpleNamespace(reply_to_message=reply)
assert TelegramChannel._extract_reply_context(message) is None
assert await channel._extract_reply_context(message) is None
@pytest.mark.asyncio
+5 -651
View File
@@ -1,22 +1,17 @@
import asyncio
import json
import tempfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import httpx
import nanobot.channels.weixin as weixin_mod
from nanobot.bus.queue import MessageBus
from nanobot.channels.weixin import (
ITEM_IMAGE,
ITEM_TEXT,
MESSAGE_TYPE_BOT,
WEIXIN_CHANNEL_VERSION,
_decrypt_aes_ecb,
_encrypt_aes_ecb,
WeixinChannel,
WeixinConfig,
)
@@ -47,12 +42,10 @@ def test_make_headers_includes_route_tag_when_configured() -> None:
assert headers["Authorization"] == "Bearer token"
assert headers["SKRouteTag"] == "123"
assert headers["iLink-App-Id"] == "bot"
assert headers["iLink-App-ClientVersion"] == str((2 << 16) | (1 << 8) | 1)
def test_channel_version_matches_reference_plugin_version() -> None:
assert WEIXIN_CHANNEL_VERSION == "2.1.1"
assert WEIXIN_CHANNEL_VERSION == "1.0.3"
def test_save_and_load_state_persists_context_tokens(tmp_path) -> None:
@@ -176,120 +169,6 @@ async def test_process_message_extracts_media_and_preserves_paths() -> None:
assert inbound.media == ["/tmp/test.jpg"]
@pytest.mark.asyncio
async def test_process_message_falls_back_to_referenced_media_when_no_top_level_media() -> None:
channel, bus = _make_channel()
channel._download_media_item = AsyncMock(return_value="/tmp/ref.jpg")
await channel._process_message(
{
"message_type": 1,
"message_id": "m3-ref-fallback",
"from_user_id": "wx-user",
"context_token": "ctx-3-ref-fallback",
"item_list": [
{
"type": ITEM_TEXT,
"text_item": {"text": "reply to image"},
"ref_msg": {
"message_item": {
"type": ITEM_IMAGE,
"image_item": {"media": {"encrypt_query_param": "ref-enc"}},
},
},
},
],
}
)
inbound = await asyncio.wait_for(bus.consume_inbound(), timeout=1.0)
channel._download_media_item.assert_awaited_once_with(
{"media": {"encrypt_query_param": "ref-enc"}},
"image",
)
assert inbound.media == ["/tmp/ref.jpg"]
assert "reply to image" in inbound.content
assert "[image]" in inbound.content
@pytest.mark.asyncio
async def test_process_message_does_not_use_referenced_fallback_when_top_level_media_exists() -> None:
channel, bus = _make_channel()
channel._download_media_item = AsyncMock(side_effect=["/tmp/top.jpg", "/tmp/ref.jpg"])
await channel._process_message(
{
"message_type": 1,
"message_id": "m3-ref-no-fallback",
"from_user_id": "wx-user",
"context_token": "ctx-3-ref-no-fallback",
"item_list": [
{"type": ITEM_IMAGE, "image_item": {"media": {"encrypt_query_param": "top-enc"}}},
{
"type": ITEM_TEXT,
"text_item": {"text": "has top-level media"},
"ref_msg": {
"message_item": {
"type": ITEM_IMAGE,
"image_item": {"media": {"encrypt_query_param": "ref-enc"}},
},
},
},
],
}
)
inbound = await asyncio.wait_for(bus.consume_inbound(), timeout=1.0)
channel._download_media_item.assert_awaited_once_with(
{"media": {"encrypt_query_param": "top-enc"}},
"image",
)
assert inbound.media == ["/tmp/top.jpg"]
assert "/tmp/ref.jpg" not in inbound.content
@pytest.mark.asyncio
async def test_process_message_does_not_fallback_when_top_level_media_exists_but_download_fails() -> None:
channel, bus = _make_channel()
# Top-level image download fails (None), referenced image would succeed if fallback were triggered.
channel._download_media_item = AsyncMock(side_effect=[None, "/tmp/ref.jpg"])
await channel._process_message(
{
"message_type": 1,
"message_id": "m3-ref-no-fallback-on-failure",
"from_user_id": "wx-user",
"context_token": "ctx-3-ref-no-fallback-on-failure",
"item_list": [
{"type": ITEM_IMAGE, "image_item": {"media": {"encrypt_query_param": "top-enc"}}},
{
"type": ITEM_TEXT,
"text_item": {"text": "quoted has media"},
"ref_msg": {
"message_item": {
"type": ITEM_IMAGE,
"image_item": {"media": {"encrypt_query_param": "ref-enc"}},
},
},
},
],
}
)
inbound = await asyncio.wait_for(bus.consume_inbound(), timeout=1.0)
# Should only attempt top-level media item; reference fallback must not activate.
channel._download_media_item.assert_awaited_once_with(
{"media": {"encrypt_query_param": "top-enc"}},
"image",
)
assert inbound.media == []
assert "[image]" in inbound.content
assert "/tmp/ref.jpg" not in inbound.content
@pytest.mark.asyncio
async def test_send_without_context_token_does_not_send_text() -> None:
channel, _bus = _make_channel()
@@ -320,70 +199,6 @@ async def test_send_does_not_send_when_session_is_paused() -> None:
channel._send_text.assert_not_awaited()
@pytest.mark.asyncio
async def test_get_typing_ticket_fetches_and_caches_per_user() -> None:
channel, _bus = _make_channel()
channel._client = object()
channel._token = "token"
channel._api_post = AsyncMock(return_value={"ret": 0, "typing_ticket": "ticket-1"})
first = await channel._get_typing_ticket("wx-user", "ctx-1")
second = await channel._get_typing_ticket("wx-user", "ctx-2")
assert first == "ticket-1"
assert second == "ticket-1"
channel._api_post.assert_awaited_once_with(
"ilink/bot/getconfig",
{"ilink_user_id": "wx-user", "context_token": "ctx-1", "base_info": weixin_mod.BASE_INFO},
)
@pytest.mark.asyncio
async def test_send_uses_typing_start_and_cancel_when_ticket_available() -> None:
channel, _bus = _make_channel()
channel._client = object()
channel._token = "token"
channel._context_tokens["wx-user"] = "ctx-typing"
channel._send_text = AsyncMock()
channel._api_post = AsyncMock(
side_effect=[
{"ret": 0, "typing_ticket": "ticket-typing"},
{"ret": 0},
{"ret": 0},
]
)
await channel.send(
type("Msg", (), {"chat_id": "wx-user", "content": "pong", "media": [], "metadata": {}})()
)
channel._send_text.assert_awaited_once_with("wx-user", "pong", "ctx-typing")
assert channel._api_post.await_count == 3
assert channel._api_post.await_args_list[0].args[0] == "ilink/bot/getconfig"
assert channel._api_post.await_args_list[1].args[0] == "ilink/bot/sendtyping"
assert channel._api_post.await_args_list[1].args[1]["status"] == 1
assert channel._api_post.await_args_list[2].args[0] == "ilink/bot/sendtyping"
assert channel._api_post.await_args_list[2].args[1]["status"] == 2
@pytest.mark.asyncio
async def test_send_still_sends_text_when_typing_ticket_missing() -> None:
channel, _bus = _make_channel()
channel._client = object()
channel._token = "token"
channel._context_tokens["wx-user"] = "ctx-no-ticket"
channel._send_text = AsyncMock()
channel._api_post = AsyncMock(return_value={"ret": 1, "errmsg": "no config"})
await channel.send(
type("Msg", (), {"chat_id": "wx-user", "content": "pong", "media": [], "metadata": {}})()
)
channel._send_text.assert_awaited_once_with("wx-user", "pong", "ctx-no-ticket")
channel._api_post.assert_awaited_once()
assert channel._api_post.await_args_list[0].args[0] == "ilink/bot/getconfig"
@pytest.mark.asyncio
async def test_poll_once_pauses_session_on_expired_errcode() -> None:
channel, _bus = _make_channel()
@@ -405,12 +220,8 @@ async def test_qr_login_refreshes_expired_qr_and_then_succeeds() -> None:
channel._api_get = AsyncMock(
side_effect=[
{"qrcode": "qr-1", "qrcode_img_content": "url-1"},
{"qrcode": "qr-2", "qrcode_img_content": "url-2"},
]
)
channel._api_get_with_base = AsyncMock(
side_effect=[
{"status": "expired"},
{"qrcode": "qr-2", "qrcode_img_content": "url-2"},
{
"status": "confirmed",
"bot_token": "token-2",
@@ -436,16 +247,12 @@ async def test_qr_login_returns_false_after_too_many_expired_qr_codes() -> None:
channel._api_get = AsyncMock(
side_effect=[
{"qrcode": "qr-1", "qrcode_img_content": "url-1"},
{"status": "expired"},
{"qrcode": "qr-2", "qrcode_img_content": "url-2"},
{"status": "expired"},
{"qrcode": "qr-3", "qrcode_img_content": "url-3"},
{"status": "expired"},
{"qrcode": "qr-4", "qrcode_img_content": "url-4"},
]
)
channel._api_get_with_base = AsyncMock(
side_effect=[
{"status": "expired"},
{"status": "expired"},
{"status": "expired"},
{"status": "expired"},
]
)
@@ -455,105 +262,6 @@ async def test_qr_login_returns_false_after_too_many_expired_qr_codes() -> None:
assert ok is False
@pytest.mark.asyncio
async def test_qr_login_switches_polling_base_url_on_redirect_status() -> None:
channel, _bus = _make_channel()
channel._running = True
channel._save_state = lambda: None
channel._print_qr_code = lambda url: None
channel._fetch_qr_code = AsyncMock(return_value=("qr-1", "url-1"))
status_side_effect = [
{"status": "scaned_but_redirect", "redirect_host": "idc.redirect.test"},
{
"status": "confirmed",
"bot_token": "token-3",
"ilink_bot_id": "bot-3",
"baseurl": "https://example.test",
"ilink_user_id": "wx-user",
},
]
channel._api_get = AsyncMock(side_effect=list(status_side_effect))
channel._api_get_with_base = AsyncMock(side_effect=list(status_side_effect))
ok = await channel._qr_login()
assert ok is True
assert channel._token == "token-3"
assert channel._api_get_with_base.await_count == 2
first_call = channel._api_get_with_base.await_args_list[0]
second_call = channel._api_get_with_base.await_args_list[1]
assert first_call.kwargs["base_url"] == "https://ilinkai.weixin.qq.com"
assert second_call.kwargs["base_url"] == "https://idc.redirect.test"
@pytest.mark.asyncio
async def test_qr_login_redirect_without_host_keeps_current_polling_base_url() -> None:
channel, _bus = _make_channel()
channel._running = True
channel._save_state = lambda: None
channel._print_qr_code = lambda url: None
channel._fetch_qr_code = AsyncMock(return_value=("qr-1", "url-1"))
status_side_effect = [
{"status": "scaned_but_redirect"},
{
"status": "confirmed",
"bot_token": "token-4",
"ilink_bot_id": "bot-4",
"baseurl": "https://example.test",
"ilink_user_id": "wx-user",
},
]
channel._api_get = AsyncMock(side_effect=list(status_side_effect))
channel._api_get_with_base = AsyncMock(side_effect=list(status_side_effect))
ok = await channel._qr_login()
assert ok is True
assert channel._token == "token-4"
assert channel._api_get_with_base.await_count == 2
first_call = channel._api_get_with_base.await_args_list[0]
second_call = channel._api_get_with_base.await_args_list[1]
assert first_call.kwargs["base_url"] == "https://ilinkai.weixin.qq.com"
assert second_call.kwargs["base_url"] == "https://ilinkai.weixin.qq.com"
@pytest.mark.asyncio
async def test_qr_login_resets_redirect_base_url_after_qr_refresh() -> None:
channel, _bus = _make_channel()
channel._running = True
channel._save_state = lambda: None
channel._print_qr_code = lambda url: None
channel._fetch_qr_code = AsyncMock(side_effect=[("qr-1", "url-1"), ("qr-2", "url-2")])
channel._api_get_with_base = AsyncMock(
side_effect=[
{"status": "scaned_but_redirect", "redirect_host": "idc.redirect.test"},
{"status": "expired"},
{
"status": "confirmed",
"bot_token": "token-5",
"ilink_bot_id": "bot-5",
"baseurl": "https://example.test",
"ilink_user_id": "wx-user",
},
]
)
ok = await channel._qr_login()
assert ok is True
assert channel._token == "token-5"
assert channel._api_get_with_base.await_count == 3
first_call = channel._api_get_with_base.await_args_list[0]
second_call = channel._api_get_with_base.await_args_list[1]
third_call = channel._api_get_with_base.await_args_list[2]
assert first_call.kwargs["base_url"] == "https://ilinkai.weixin.qq.com"
assert second_call.kwargs["base_url"] == "https://idc.redirect.test"
assert third_call.kwargs["base_url"] == "https://ilinkai.weixin.qq.com"
@pytest.mark.asyncio
async def test_process_message_skips_bot_messages() -> None:
channel, bus = _make_channel()
@@ -570,357 +278,3 @@ async def test_process_message_skips_bot_messages() -> None:
)
assert bus.inbound_size == 0
class _DummyHttpResponse:
def __init__(self, *, headers: dict[str, str] | None = None, status_code: int = 200) -> None:
self.headers = headers or {}
self.status_code = status_code
def raise_for_status(self) -> None:
return None
@pytest.mark.asyncio
async def test_send_media_uses_upload_full_url_when_present(tmp_path) -> None:
channel, _bus = _make_channel()
media_file = tmp_path / "photo.jpg"
media_file.write_bytes(b"hello-weixin")
cdn_post = AsyncMock(return_value=_DummyHttpResponse(headers={"x-encrypted-param": "dl-param"}))
channel._client = SimpleNamespace(post=cdn_post)
channel._api_post = AsyncMock(
side_effect=[
{
"upload_full_url": "https://upload-full.example.test/path?foo=bar",
"upload_param": "should-not-be-used",
},
{"ret": 0},
]
)
await channel._send_media_file("wx-user", str(media_file), "ctx-1")
# first POST call is CDN upload
cdn_url = cdn_post.await_args_list[0].args[0]
assert cdn_url == "https://upload-full.example.test/path?foo=bar"
@pytest.mark.asyncio
async def test_send_media_falls_back_to_upload_param_url(tmp_path) -> None:
channel, _bus = _make_channel()
media_file = tmp_path / "photo.jpg"
media_file.write_bytes(b"hello-weixin")
cdn_post = AsyncMock(return_value=_DummyHttpResponse(headers={"x-encrypted-param": "dl-param"}))
channel._client = SimpleNamespace(post=cdn_post)
channel._api_post = AsyncMock(
side_effect=[
{"upload_param": "enc-need-fallback"},
{"ret": 0},
]
)
await channel._send_media_file("wx-user", str(media_file), "ctx-1")
cdn_url = cdn_post.await_args_list[0].args[0]
assert cdn_url.startswith(f"{channel.config.cdn_base_url}/upload?encrypted_query_param=enc-need-fallback")
assert "&filekey=" in cdn_url
@pytest.mark.asyncio
async def test_send_media_voice_file_uses_voice_item_and_voice_upload_type(tmp_path) -> None:
channel, _bus = _make_channel()
media_file = tmp_path / "voice.mp3"
media_file.write_bytes(b"voice-bytes")
cdn_post = AsyncMock(return_value=_DummyHttpResponse(headers={"x-encrypted-param": "voice-dl-param"}))
channel._client = SimpleNamespace(post=cdn_post)
channel._api_post = AsyncMock(
side_effect=[
{"upload_full_url": "https://upload-full.example.test/voice?foo=bar"},
{"ret": 0},
]
)
await channel._send_media_file("wx-user", str(media_file), "ctx-voice")
getupload_body = channel._api_post.await_args_list[0].args[1]
assert getupload_body["media_type"] == 4
sendmessage_body = channel._api_post.await_args_list[1].args[1]
item = sendmessage_body["msg"]["item_list"][0]
assert item["type"] == 3
assert "voice_item" in item
assert "file_item" not in item
assert item["voice_item"]["media"]["encrypt_query_param"] == "voice-dl-param"
@pytest.mark.asyncio
async def test_send_typing_uses_keepalive_until_send_finishes() -> None:
channel, _bus = _make_channel()
channel._client = object()
channel._token = "token"
channel._context_tokens["wx-user"] = "ctx-typing-loop"
async def _api_post_side_effect(endpoint: str, _body: dict | None = None, *, auth: bool = True):
if endpoint == "ilink/bot/getconfig":
return {"ret": 0, "typing_ticket": "ticket-keepalive"}
return {"ret": 0}
channel._api_post = AsyncMock(side_effect=_api_post_side_effect)
async def _slow_send_text(*_args, **_kwargs) -> None:
await asyncio.sleep(0.03)
channel._send_text = AsyncMock(side_effect=_slow_send_text)
old_interval = weixin_mod.TYPING_KEEPALIVE_INTERVAL_S
weixin_mod.TYPING_KEEPALIVE_INTERVAL_S = 0.01
try:
await channel.send(
type("Msg", (), {"chat_id": "wx-user", "content": "pong", "media": [], "metadata": {}})()
)
finally:
weixin_mod.TYPING_KEEPALIVE_INTERVAL_S = old_interval
status_calls = [
c.args[1]["status"]
for c in channel._api_post.await_args_list
if c.args and c.args[0] == "ilink/bot/sendtyping"
]
assert status_calls.count(1) >= 2
assert status_calls[-1] == 2
@pytest.mark.asyncio
async def test_get_typing_ticket_failure_uses_backoff_and_cached_ticket(monkeypatch) -> None:
channel, _bus = _make_channel()
channel._client = object()
channel._token = "token"
now = {"value": 1000.0}
monkeypatch.setattr(weixin_mod.time, "time", lambda: now["value"])
monkeypatch.setattr(weixin_mod.random, "random", lambda: 0.5)
channel._api_post = AsyncMock(return_value={"ret": 0, "typing_ticket": "ticket-ok"})
first = await channel._get_typing_ticket("wx-user", "ctx-1")
assert first == "ticket-ok"
# force refresh window reached
now["value"] = now["value"] + (12 * 60 * 60) + 1
channel._api_post = AsyncMock(return_value={"ret": 1, "errmsg": "temporary failure"})
# On refresh failure, should still return cached ticket and apply backoff.
second = await channel._get_typing_ticket("wx-user", "ctx-2")
assert second == "ticket-ok"
assert channel._api_post.await_count == 1
# Before backoff expiry, no extra fetch should happen.
now["value"] += 1
third = await channel._get_typing_ticket("wx-user", "ctx-3")
assert third == "ticket-ok"
assert channel._api_post.await_count == 1
@pytest.mark.asyncio
async def test_qr_login_treats_temporary_connect_error_as_wait_and_recovers() -> None:
channel, _bus = _make_channel()
channel._running = True
channel._save_state = lambda: None
channel._print_qr_code = lambda url: None
channel._fetch_qr_code = AsyncMock(return_value=("qr-1", "url-1"))
request = httpx.Request("GET", "https://ilinkai.weixin.qq.com/ilink/bot/get_qrcode_status")
channel._api_get_with_base = AsyncMock(
side_effect=[
httpx.ConnectError("temporary network", request=request),
{
"status": "confirmed",
"bot_token": "token-net-ok",
"ilink_bot_id": "bot-id",
"baseurl": "https://example.test",
"ilink_user_id": "wx-user",
},
]
)
ok = await channel._qr_login()
assert ok is True
assert channel._token == "token-net-ok"
@pytest.mark.asyncio
async def test_qr_login_treats_5xx_gateway_response_error_as_wait_and_recovers() -> None:
channel, _bus = _make_channel()
channel._running = True
channel._save_state = lambda: None
channel._print_qr_code = lambda url: None
channel._fetch_qr_code = AsyncMock(return_value=("qr-1", "url-1"))
request = httpx.Request("GET", "https://ilinkai.weixin.qq.com/ilink/bot/get_qrcode_status")
response = httpx.Response(status_code=524, request=request)
channel._api_get_with_base = AsyncMock(
side_effect=[
httpx.HTTPStatusError("gateway timeout", request=request, response=response),
{
"status": "confirmed",
"bot_token": "token-5xx-ok",
"ilink_bot_id": "bot-id",
"baseurl": "https://example.test",
"ilink_user_id": "wx-user",
},
]
)
ok = await channel._qr_login()
assert ok is True
assert channel._token == "token-5xx-ok"
def test_decrypt_aes_ecb_strips_valid_pkcs7_padding() -> None:
key_b64 = "MDEyMzQ1Njc4OWFiY2RlZg==" # base64("0123456789abcdef")
plaintext = b"hello-weixin-padding"
ciphertext = _encrypt_aes_ecb(plaintext, key_b64)
decrypted = _decrypt_aes_ecb(ciphertext, key_b64)
assert decrypted == plaintext
class _DummyDownloadResponse:
def __init__(self, content: bytes, status_code: int = 200) -> None:
self.content = content
self.status_code = status_code
def raise_for_status(self) -> None:
return None
class _DummyErrorDownloadResponse(_DummyDownloadResponse):
def __init__(self, url: str, status_code: int) -> None:
super().__init__(content=b"", status_code=status_code)
self._url = url
def raise_for_status(self) -> None:
request = httpx.Request("GET", self._url)
response = httpx.Response(self.status_code, request=request)
raise httpx.HTTPStatusError(
f"download failed with status {self.status_code}",
request=request,
response=response,
)
@pytest.mark.asyncio
async def test_download_media_item_uses_full_url_when_present(tmp_path) -> None:
channel, _bus = _make_channel()
weixin_mod.get_media_dir = lambda _name: tmp_path
full_url = "https://cdn.example.test/download/full"
channel._client = SimpleNamespace(
get=AsyncMock(return_value=_DummyDownloadResponse(content=b"raw-image-bytes"))
)
item = {
"media": {
"full_url": full_url,
"encrypt_query_param": "enc-fallback-should-not-be-used",
},
}
saved_path = await channel._download_media_item(item, "image")
assert saved_path is not None
assert Path(saved_path).read_bytes() == b"raw-image-bytes"
channel._client.get.assert_awaited_once_with(full_url)
@pytest.mark.asyncio
async def test_download_media_item_falls_back_when_full_url_returns_retryable_error(tmp_path) -> None:
channel, _bus = _make_channel()
weixin_mod.get_media_dir = lambda _name: tmp_path
full_url = "https://cdn.example.test/download/full?taskid=123"
channel._client = SimpleNamespace(
get=AsyncMock(
side_effect=[
_DummyErrorDownloadResponse(full_url, 500),
_DummyDownloadResponse(content=b"fallback-bytes"),
]
)
)
item = {
"media": {
"full_url": full_url,
"encrypt_query_param": "enc-fallback",
},
}
saved_path = await channel._download_media_item(item, "image")
assert saved_path is not None
assert Path(saved_path).read_bytes() == b"fallback-bytes"
assert channel._client.get.await_count == 2
assert channel._client.get.await_args_list[0].args[0] == full_url
fallback_url = channel._client.get.await_args_list[1].args[0]
assert fallback_url.startswith(f"{channel.config.cdn_base_url}/download?encrypted_query_param=enc-fallback")
@pytest.mark.asyncio
async def test_download_media_item_falls_back_to_encrypt_query_param(tmp_path) -> None:
channel, _bus = _make_channel()
weixin_mod.get_media_dir = lambda _name: tmp_path
channel._client = SimpleNamespace(
get=AsyncMock(return_value=_DummyDownloadResponse(content=b"fallback-bytes"))
)
item = {"media": {"encrypt_query_param": "enc-fallback"}}
saved_path = await channel._download_media_item(item, "image")
assert saved_path is not None
assert Path(saved_path).read_bytes() == b"fallback-bytes"
called_url = channel._client.get.await_args_list[0].args[0]
assert called_url.startswith(f"{channel.config.cdn_base_url}/download?encrypted_query_param=enc-fallback")
@pytest.mark.asyncio
async def test_download_media_item_does_not_retry_when_full_url_fails_without_fallback(tmp_path) -> None:
channel, _bus = _make_channel()
weixin_mod.get_media_dir = lambda _name: tmp_path
full_url = "https://cdn.example.test/download/full"
channel._client = SimpleNamespace(
get=AsyncMock(return_value=_DummyErrorDownloadResponse(full_url, 500))
)
item = {"media": {"full_url": full_url}}
saved_path = await channel._download_media_item(item, "image")
assert saved_path is None
channel._client.get.assert_awaited_once_with(full_url)
@pytest.mark.asyncio
async def test_download_media_item_non_image_requires_aes_key_even_with_full_url(tmp_path) -> None:
channel, _bus = _make_channel()
weixin_mod.get_media_dir = lambda _name: tmp_path
full_url = "https://cdn.example.test/download/voice"
channel._client = SimpleNamespace(
get=AsyncMock(return_value=_DummyDownloadResponse(content=b"ciphertext-or-unknown"))
)
item = {
"media": {
"full_url": full_url,
},
}
saved_path = await channel._download_media_item(item, "voice")
assert saved_path is None
channel._client.get.assert_not_awaited()
+26
View File
@@ -145,3 +145,29 @@ def test_response_renderable_without_metadata_keeps_markdown_path():
renderable = commands._response_renderable(help_text, render_markdown=True)
assert renderable.__class__.__name__ == "Markdown"
def test_stream_renderer_stop_for_input_stops_spinner():
"""stop_for_input should stop the active spinner to avoid prompt_toolkit conflicts."""
spinner = MagicMock()
mock_console = MagicMock()
mock_console.status.return_value = spinner
# Create renderer with mocked console
with patch.object(stream_mod, "_make_console", return_value=mock_console):
renderer = stream_mod.StreamRenderer(show_spinner=True)
# Verify spinner started
spinner.start.assert_called_once()
# Stop for input
renderer.stop_for_input()
# Verify spinner stopped
spinner.stop.assert_called_once()
def test_make_console_uses_force_terminal():
"""Console should be created with force_terminal=True for proper ANSI handling."""
console = stream_mod._make_console()
assert console._force_terminal is True
+6 -4
View File
@@ -127,7 +127,7 @@ class TestRestartCommand:
loop.sessions.get_or_create.return_value = session
loop._start_time = time.time() - 125
loop._last_usage = {"prompt_tokens": 0, "completion_tokens": 0}
loop.memory_consolidator.estimate_session_prompt_tokens = MagicMock(
loop.consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(20500, "tiktoken")
)
@@ -152,10 +152,12 @@ class TestRestartCommand:
])
await loop._run_agent_loop([])
assert loop._last_usage == {"prompt_tokens": 9, "completion_tokens": 4}
assert loop._last_usage["prompt_tokens"] == 9
assert loop._last_usage["completion_tokens"] == 4
await loop._run_agent_loop([])
assert loop._last_usage == {"prompt_tokens": 0, "completion_tokens": 0}
assert loop._last_usage["prompt_tokens"] == 0
assert loop._last_usage["completion_tokens"] == 0
@pytest.mark.asyncio
async def test_status_falls_back_to_last_usage_when_context_estimate_missing(self):
@@ -164,7 +166,7 @@ class TestRestartCommand:
session.get_history.return_value = [{"role": "user"}]
loop.sessions.get_or_create.return_value = session
loop._last_usage = {"prompt_tokens": 1200, "completion_tokens": 34}
loop.memory_consolidator.estimate_session_prompt_tokens = MagicMock(
loop.consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(0, "none")
)
+231
View File
@@ -0,0 +1,231 @@
"""Tests for cached token extraction from OpenAI-compatible providers."""
from __future__ import annotations
from nanobot.providers.openai_compat_provider import OpenAICompatProvider
class FakeUsage:
"""Mimics an OpenAI SDK usage object (has attributes, not dict keys)."""
def __init__(self, **kwargs):
for k, v in kwargs.items():
setattr(self, k, v)
class FakePromptDetails:
"""Mimics prompt_tokens_details sub-object."""
def __init__(self, cached_tokens=0):
self.cached_tokens = cached_tokens
class _FakeSpec:
supports_prompt_caching = False
model_id_prefix = None
strip_model_prefix = False
max_completion_tokens = False
reasoning_effort = None
def _provider():
from unittest.mock import MagicMock
p = OpenAICompatProvider.__new__(OpenAICompatProvider)
p.client = MagicMock()
p.spec = _FakeSpec()
return p
# Minimal valid choice so _parse reaches _extract_usage.
_DICT_CHOICE = {"message": {"content": "Hello"}}
class _FakeMessage:
content = "Hello"
tool_calls = None
class _FakeChoice:
message = _FakeMessage()
finish_reason = "stop"
# --- dict-based response (raw JSON / mapping) ---
def test_extract_usage_openai_cached_tokens_dict():
"""prompt_tokens_details.cached_tokens from a dict response."""
p = _provider()
response = {
"choices": [_DICT_CHOICE],
"usage": {
"prompt_tokens": 2000,
"completion_tokens": 300,
"total_tokens": 2300,
"prompt_tokens_details": {"cached_tokens": 1200},
}
}
result = p._parse(response)
assert result.usage["cached_tokens"] == 1200
assert result.usage["prompt_tokens"] == 2000
def test_extract_usage_deepseek_cached_tokens_dict():
"""prompt_cache_hit_tokens from a DeepSeek dict response."""
p = _provider()
response = {
"choices": [_DICT_CHOICE],
"usage": {
"prompt_tokens": 1500,
"completion_tokens": 200,
"total_tokens": 1700,
"prompt_cache_hit_tokens": 1200,
"prompt_cache_miss_tokens": 300,
}
}
result = p._parse(response)
assert result.usage["cached_tokens"] == 1200
def test_extract_usage_no_cached_tokens_dict():
"""Response without any cache fields -> no cached_tokens key."""
p = _provider()
response = {
"choices": [_DICT_CHOICE],
"usage": {
"prompt_tokens": 1000,
"completion_tokens": 200,
"total_tokens": 1200,
}
}
result = p._parse(response)
assert "cached_tokens" not in result.usage
def test_extract_usage_openai_cached_zero_dict():
"""cached_tokens=0 should NOT be included (same as existing fields)."""
p = _provider()
response = {
"choices": [_DICT_CHOICE],
"usage": {
"prompt_tokens": 2000,
"completion_tokens": 300,
"total_tokens": 2300,
"prompt_tokens_details": {"cached_tokens": 0},
}
}
result = p._parse(response)
assert "cached_tokens" not in result.usage
# --- object-based response (OpenAI SDK Pydantic model) ---
def test_extract_usage_openai_cached_tokens_obj():
"""prompt_tokens_details.cached_tokens from an SDK object response."""
p = _provider()
usage_obj = FakeUsage(
prompt_tokens=2000,
completion_tokens=300,
total_tokens=2300,
prompt_tokens_details=FakePromptDetails(cached_tokens=1200),
)
response = FakeUsage(choices=[_FakeChoice()], usage=usage_obj)
result = p._parse(response)
assert result.usage["cached_tokens"] == 1200
def test_extract_usage_deepseek_cached_tokens_obj():
"""prompt_cache_hit_tokens from a DeepSeek SDK object response."""
p = _provider()
usage_obj = FakeUsage(
prompt_tokens=1500,
completion_tokens=200,
total_tokens=1700,
prompt_cache_hit_tokens=1200,
)
response = FakeUsage(choices=[_FakeChoice()], usage=usage_obj)
result = p._parse(response)
assert result.usage["cached_tokens"] == 1200
def test_extract_usage_stepfun_top_level_cached_tokens_dict():
"""StepFun/Moonshot: usage.cached_tokens at top level (not nested)."""
p = _provider()
response = {
"choices": [_DICT_CHOICE],
"usage": {
"prompt_tokens": 591,
"completion_tokens": 120,
"total_tokens": 711,
"cached_tokens": 512,
}
}
result = p._parse(response)
assert result.usage["cached_tokens"] == 512
def test_extract_usage_stepfun_top_level_cached_tokens_obj():
"""StepFun/Moonshot: usage.cached_tokens as SDK object attribute."""
p = _provider()
usage_obj = FakeUsage(
prompt_tokens=591,
completion_tokens=120,
total_tokens=711,
cached_tokens=512,
)
response = FakeUsage(choices=[_FakeChoice()], usage=usage_obj)
result = p._parse(response)
assert result.usage["cached_tokens"] == 512
def test_extract_usage_priority_nested_over_top_level_dict():
"""When both nested and top-level cached_tokens exist, nested wins."""
p = _provider()
response = {
"choices": [_DICT_CHOICE],
"usage": {
"prompt_tokens": 2000,
"completion_tokens": 300,
"total_tokens": 2300,
"prompt_tokens_details": {"cached_tokens": 100},
"cached_tokens": 500,
}
}
result = p._parse(response)
assert result.usage["cached_tokens"] == 100
def test_anthropic_maps_cache_fields_to_cached_tokens():
"""Anthropic's cache_read_input_tokens should map to cached_tokens."""
from nanobot.providers.anthropic_provider import AnthropicProvider
usage_obj = FakeUsage(
input_tokens=800,
output_tokens=200,
cache_creation_input_tokens=0,
cache_read_input_tokens=1200,
)
content_block = FakeUsage(type="text", text="hello")
response = FakeUsage(
id="msg_1",
type="message",
stop_reason="end_turn",
content=[content_block],
usage=usage_obj,
)
result = AnthropicProvider._parse_response(response)
assert result.usage["cached_tokens"] == 1200
assert result.usage["prompt_tokens"] == 800
def test_anthropic_no_cache_fields():
"""Anthropic response without cache fields should not have cached_tokens."""
from nanobot.providers.anthropic_provider import AnthropicProvider
usage_obj = FakeUsage(input_tokens=800, output_tokens=200)
content_block = FakeUsage(type="text", text="hello")
response = FakeUsage(
id="msg_1",
type="message",
stop_reason="end_turn",
content=[content_block],
usage=usage_obj,
)
result = AnthropicProvider._parse_response(response)
assert "cached_tokens" not in result.usage
@@ -0,0 +1,34 @@
"""Regression tests for max_completion_tokens selection in OpenAI-compatible providers."""
from unittest.mock import patch
from nanobot.providers.openai_compat_provider import OpenAICompatProvider
from nanobot.providers.registry import find_by_name
def test_openai_provider_uses_max_completion_tokens_when_supported():
"""OpenAI registry spec should drive max_completion_tokens payload selection."""
spec = find_by_name("openai")
assert spec is not None
assert spec.supports_max_completion_tokens is True
with patch("nanobot.providers.openai_compat_provider.AsyncOpenAI"):
provider = OpenAICompatProvider(
api_key="test-key",
api_base=None,
default_model="gpt-4.1",
spec=spec,
)
payload = provider._build_kwargs(
messages=[{"role": "user", "content": "Hello"}],
tools=None,
model="gpt-4.1",
max_tokens=1234,
temperature=0.2,
reasoning_effort=None,
tool_choice=None,
)
assert payload["max_completion_tokens"] == 1234
assert "max_tokens" not in payload
+59
View File
@@ -0,0 +1,59 @@
"""Tests for build_status_content cache hit rate display."""
from nanobot.utils.helpers import build_status_content
def test_status_shows_cache_hit_rate():
content = build_status_content(
version="0.1.0",
model="glm-4-plus",
start_time=1000000.0,
last_usage={"prompt_tokens": 2000, "completion_tokens": 300, "cached_tokens": 1200},
context_window_tokens=128000,
session_msg_count=10,
context_tokens_estimate=5000,
)
assert "60% cached" in content
assert "2000 in / 300 out" in content
def test_status_no_cache_info():
"""Without cached_tokens, display should not show cache percentage."""
content = build_status_content(
version="0.1.0",
model="glm-4-plus",
start_time=1000000.0,
last_usage={"prompt_tokens": 2000, "completion_tokens": 300},
context_window_tokens=128000,
session_msg_count=10,
context_tokens_estimate=5000,
)
assert "cached" not in content.lower()
assert "2000 in / 300 out" in content
def test_status_zero_cached_tokens():
"""cached_tokens=0 should not show cache percentage."""
content = build_status_content(
version="0.1.0",
model="glm-4-plus",
start_time=1000000.0,
last_usage={"prompt_tokens": 2000, "completion_tokens": 300, "cached_tokens": 0},
context_window_tokens=128000,
session_msg_count=10,
context_tokens_estimate=5000,
)
assert "cached" not in content.lower()
def test_status_100_percent_cached():
content = build_status_content(
version="0.1.0",
model="glm-4-plus",
start_time=1000000.0,
last_usage={"prompt_tokens": 1000, "completion_tokens": 100, "cached_tokens": 1000},
context_window_tokens=128000,
session_msg_count=5,
context_tokens_estimate=3000,
)
assert "100% cached" in content
+7 -7
View File
@@ -235,15 +235,10 @@ async def test_followup_requests_share_same_session_key(aiohttp_client) -> None:
@pytest.mark.asyncio
async def test_fixed_session_requests_are_serialized(aiohttp_client) -> None:
order: list[str] = []
barrier = asyncio.Event()
async def slow_process(content, session_key="", channel="", chat_id=""):
order.append(f"start:{content}")
if content == "first":
barrier.set()
await asyncio.sleep(0.1)
else:
await barrier.wait()
await asyncio.sleep(0.1)
order.append(f"end:{content}")
return content
@@ -264,7 +259,12 @@ async def test_fixed_session_requests_are_serialized(aiohttp_client) -> None:
r1, r2 = await asyncio.gather(send("first"), send("second"))
assert r1.status == 200
assert r2.status == 200
assert order.index("end:first") < order.index("start:second")
# Verify serialization: one process must fully finish before the other starts
assert "end:second" in order or "end:first" in order
if order[0] == "start:first":
assert order.index("end:first") < order.index("start:second")
else:
assert order.index("end:second") < order.index("start:first")
@pytest.mark.skipif(not HAS_AIOHTTP, reason="aiohttp not installed")
-50
View File
@@ -408,56 +408,6 @@ async def test_exec_timeout_capped_at_max() -> None:
assert "Exit code: 0" in result
async def test_exec_command_wrapper_applied() -> None:
"""command_wrapper should wrap the original command."""
tool = ExecTool(command_wrapper="echo WRAPPED: {command}")
result = await tool.execute(command="hello")
assert "WRAPPED: hello" in result
assert "Exit code: 0" in result
async def test_exec_command_wrapper_with_cwd(tmp_path) -> None:
"""command_wrapper should substitute {cwd} with the absolute working directory."""
tool = ExecTool(command_wrapper="echo CWD:{cwd} CMD:{command}")
result = await tool.execute(command="hi", working_dir=str(tmp_path))
assert str(tmp_path) in result
assert "CMD:hi" in result
async def test_exec_command_wrapper_empty_noop() -> None:
"""Empty command_wrapper should leave the command unchanged."""
tool = ExecTool(command_wrapper="")
result = await tool.execute(command="echo direct")
assert "direct" in result
assert "Exit code: 0" in result
async def test_exec_command_wrapper_guard_runs_before_wrapper() -> None:
"""Safety guard should run before wrapper substitution."""
tool = ExecTool(command_wrapper="echo WRAPPED:{command}")
result = await tool.execute(command="rm -rf /")
assert "blocked by safety guard" in result
assert "WRAPPED:" not in result
async def test_exec_command_wrapper_ignores_unknown_placeholders() -> None:
"""Unknown {placeholders} in the wrapper should be left as-is, not raise KeyError."""
tool = ExecTool(command_wrapper="echo {command} {unknown}")
result = await tool.execute(command="hello")
assert "Exit code: 0" in result
assert "{unknown}" in result
async def test_exec_command_wrapper_does_not_leak_attributes() -> None:
"""Wrapper should not expose Python internals via attribute access."""
tool = ExecTool(command_wrapper="echo {command.__class__}")
result = await tool.execute(command="hello")
assert "Exit code: 0" in result
# {command.__class__} is not a valid placeholder; {command} gets replaced
# leaving {.__class__} as a literal string — no Python object is leaked.
assert "<class" not in result
# --- _resolve_type and nullable param tests ---