mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-08 21:38:40 +03:00
Compare commits
28
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
59152877af | ||
|
|
791282fc75 | ||
|
|
a0979d17ee | ||
|
|
a6aa0b7932 | ||
|
|
a662ace8dd | ||
|
|
475ea06294 | ||
|
|
b2598270bf | ||
|
|
9d07093e6d | ||
|
|
4fd4658146 | ||
|
|
c93bf114d2 | ||
|
|
e747d32dda | ||
|
|
5d2a13db6a | ||
|
|
9a468ab69b | ||
|
|
c41ab9e52e | ||
|
|
2780292c9d | ||
|
|
ebf6201b42 | ||
|
|
94e6d569b3 | ||
|
|
973fcced2f | ||
|
|
0d6deb9197 | ||
|
|
5da86258cc | ||
|
|
cda627f956 | ||
|
|
0334fa9944 | ||
|
|
4741026538 | ||
|
|
7c34910fd1 | ||
|
|
dd73d5d8df | ||
|
|
0d10c129e8 | ||
|
|
80403d352b | ||
|
|
5118ea824a |
@@ -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. |
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
@@ -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",
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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
@@ -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
@@ -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
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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("&", "&").replace("<", "<").replace(">", ">")
|
||||
|
||||
|
||||
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("&", "&").replace("<", "<").replace(">", ">")
|
||||
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("&", "&").replace("<", "<").replace(">", ">")
|
||||
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("&", "&").replace("<", "<").replace(">", ">")
|
||||
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
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
@@ -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)."""
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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())
|
||||
|
||||
@@ -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"
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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")
|
||||
|
||||
@@ -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 ---
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user