Compare commits

...
Author SHA1 Message Date
Dianqi Jiandchengyongru c3b55ba289 feat(channels/feishu): add domain config for Lark global support
Add 'domain' field to FeishuConfig (Literal['feishu', 'lark'], default 'feishu').
Pass domain to lark.Client.builder() and lark.ws.Client to support Lark global
(open.larksuite.com) in addition to Feishu China (open.feishu.cn).
Existing configs default to 'feishu' for backward compatibility.

Also add documentation for domain field in README.md and add tests for
domain config.
2026-04-11 23:50:41 +08:00
chengyongruandGitHub bc4cc49a59 feat(agent): mid-turn message injection for responsive follow-ups (#2985)
* feat(agent): add mid-turn message injection for responsive follow-ups

Allow user messages sent during an active agent turn to be injected
into the running LLM context instead of being queued behind a
per-session lock. Inspired by Claude Code's mid-turn queue drain
mechanism (query.ts:1547-1643).

Key design decisions:
- Messages are injected as natural user messages between iterations,
  no tool cancellation or special system prompt needed
- Two drain checkpoints: after tool execution and after final LLM
  response ("last-mile" to prevent dropping late arrivals)
- Bounded by MAX_INJECTION_CYCLES (5) to prevent consuming the
  iteration budget on rapid follow-ups
- had_injections flag bypasses _sent_in_turn suppression so follow-up
  responses are always delivered

Closes #1609

* fix(agent): harden mid-turn injection with streaming fix, bounded queue, and message safety

- Fix streaming protocol violation: Checkpoint 2 now checks for injections
  BEFORE calling on_stream_end, passing resuming=True when injections found
  so streaming channels (Feishu) don't prematurely finalize the card
- Bound pending queue to maxsize=20 with QueueFull handling
- Add warning log when injection batch exceeds _MAX_INJECTIONS_PER_TURN
- Re-publish leftover queue messages to bus in _dispatch finally block to
  prevent silent message loss on early exit (max_iterations, tool_error, cancel)
- Fix PEP 8 blank line before dataclass and logger.info indentation
- Add 12 new tests covering drain, checkpoints, cycle cap, queue routing,
  cleanup, and leftover re-publish
2026-04-11 02:11:02 +08:00
chengyongru df6f9dd171 fix(wecom): use reply_stream for progress messages to avoid errcode=40008
The plain reply() uses cmd="reply" which does not support "text" msgtype
and causes WeCom API to return errcode=40008 (invalid message type).
Unify both progress and final text messages to use reply_stream()
(cmd="aibot_respond_msg"), differentiating via finish flag.

Fixes #2999
2026-04-10 22:20:28 +08:00
chengyongruandGitHub 6af81bc4a3 feat(agent): auto compact — proactive session compression to reduce token cost and latency (#2982)
When a user is idle for longer than a configured TTL, nanobot **proactively** compresses the session context into a summary. This reduces token cost and first-token latency when the user returns — instead of re-processing a long stale context with an expired KV cache, the model receives a compact summary and fresh input.
2026-04-10 17:43:42 +08:00
chengyongruandchengyongru 4a33c1392b test(channels): add media support tests for QQ and WeCom channels
Cover helpers (sanitize_filename, guess media type), outbound send
(exception handling, media-then-text order, fallback), inbound message
processing (attachments, dedup, empty content), _post_base64file
payload filtering, and WeCom upload/download flows.
2026-04-10 17:05:59 +08:00
chengyongruandchengyongru 4b0fdffe39 fix(wecom): harden upload/download, extract media type helper
- Use asyncio.to_thread for file I/O to avoid blocking event loop
- Add 200MB upload size limit with early rejection
- Fix file handle leak by using context manager
- Use memoryview for upload chunking to reduce peak memory
- Add inbound download size check to prevent OOM
- Use asyncio.to_thread for write_bytes in download path
- Extract inline media_type detection to _guess_wecom_media_type()
2026-04-10 17:05:59 +08:00
chengyongruandchengyongru 4fe23a01c9 fix(wecom): harden upload and inbound media handling
- Use asyncio.to_thread for file I/O to avoid blocking event loop
- Add 200MB upload size limit with early rejection
- Fix file handle leak by using context manager
- Free raw bytes early after chunking to reduce memory pressure
- Add file attachments to media_paths (was text-only, inconsistent with image)
- Use robust _sanitize_filename() instead of os.path.basename() for path safety
- Remove re-raise in send() for consistency with QQ channel
- Fix truncated media_id logging for short IDs
2026-04-10 17:05:59 +08:00
gem12andchengyongru 973b888d39 feat(channels): Add full media support for QQ and WeCom channels
QQ channel improvements (on top of nightly):
- Add top-level try/except in _on_message and send() for resilience
- Use defensive getattr() for attachment attributes (botpy version compat)
- Skip file_name for image uploads to avoid QQ rendering as file attachment
- Extract only file_info from upload response to avoid extra fields
- Handle protocol-relative URLs (//...) in attachment downloads

WeCom channel improvements:
- Add _upload_media_ws() for WebSocket 3-step media upload protocol
- Send media files (image/video/voice/file) via WeCom rich media API
- Support progress messages (plain reply) vs final response (streaming)
- Support proactive send when no frame available (cron push)
- Pass media_paths to message bus for downstream processing
2026-04-10 17:05:59 +08:00
flobo3andchengyongru 7b1ce24600 fix: strip <thought> blocks from Gemma 4 and similar models 2026-04-10 00:58:00 +08:00
chengyongruandchengyongru 3bece171c2 docs(websocket): add WebSocket channel documentation
Comprehensive guide covering wire protocol, configuration reference,
token issuance, security notes, and common deployment patterns.
2026-04-09 15:56:34 +08:00
chengyongruandchengyongru 8f7ce9fef7 fix(websocket): harden security and robustness
- Use hmac.compare_digest for timing-safe static token comparison
- Add issued token capacity limit (_MAX_ISSUED_TOKENS=10000) with 429 response
- Use atomic pop in _take_issued_token_if_valid to eliminate TOCTOU window
- Enforce TLSv1.2 minimum version for SSL connections
- Extract _safe_send helper for consistent ConnectionClosed handling
- Move connection registration after ready send to prevent out-of-order delivery
- Add HTTP-level allow_from check and client_id truncation in process_request
- Make stop() idempotent with graceful shutdown error handling
- Normalize path via validator instead of leaving raw value
- Default websocket_requires_token to True for secure-by-default behavior
- Add integration tests and ws_test_client helper
- Refactor tests to use shared _ch factory and bus fixture
2026-04-09 15:56:34 +08:00
chengyongruandchengyongru d327c19db0 fix(websocket): handle ConnectionClosed gracefully in send and send_delta 2026-04-09 15:56:34 +08:00
Jack Luandchengyongru e00dca2f84 feat(channels): add WebSocket server channel and tests
Port Python implementation from a1ec7b192a
(websocket channel module and channel tests; excludes webui debug app).
2026-04-09 15:56:34 +08:00
Jiajun Xieandchengyongru 51200a954c fix(feishu): improve voice message download with detailed logging
- Add explicit error logging for missing file_key and message_id
- Add logging for download failures
- Change audio extension from .opus to .ogg for better Whisper compatibility
- Feishu voice messages are opus in OGG container; .ogg is more widely recognized
2026-04-09 10:14:26 +08:00
chengyongruandchengyongru c121547114 refactor(feishu): simplify tool hint to append-only, delegate to send_delta for throttling
- Make tool_hint_prefix configurable in FeishuConfig (default: 🔧)
- Delegate tool hint card updates from send() to send_delta() so hints
  automatically benefit from _STREAM_EDIT_INTERVAL throttling
- Fix staticmethod calls to use self.__class__ instead of self
- Document all supported metadata keys in send_delta docstring
- Add test for empty/whitespace-only tool hint with active stream buffer
2026-04-08 21:02:21 +08:00
xzq.xuandchengyongru dcc9c057bb fix(tool-hints): deduplicate by formatted string + per-line inline display
Two display fixes based on real-world Feishu testing:

1. tool_hints.py: format_tool_hints now deduplicates by comparing the
   fully formatted hint string instead of tool name alone. This fixes
   `ls /Desktop` and `ls /Downloads` being incorrectly merged as
   `ls /Desktop × 2`. Truly identical calls still fold correctly.
   (_group_consecutive and all abbreviation logic preserved unchanged.)

2. feishu.py: inline tool hints now display one tool per line with
   🔧 prefix, and use double-newline trailing to prevent Setext heading
   rendering when followed by markdown `---`.

Made-with: Cursor
2026-04-08 21:02:21 +08:00
xzq.xuandchengyongru 586d4e2411 fix(feishu): preserve tool hints in final card content
Tool hints should be kept as permanent content in the streaming card
so users can see which tools were called (matching the standalone card
behavior). Previously, hints were stripped when new deltas arrived or
when the stream ended, causing tool call information to disappear.

Now:
- New delta: hint becomes permanent content, delta appends after it
- New tool hint: replaces the previous hint (unchanged)
- Resuming/stream_end: hint is preserved in the final text

Updated 3 tests to verify hint preservation semantics.

Made-with: Cursor
2026-04-08 21:02:21 +08:00
xzq.xuandchengyongru a4bb1923ac fix(feishu): prevent tool hint stacking and clean hints on stream_end
Three fixes for inline tool hints:

1. Consecutive tool hints now replace the previous one instead of
   stacking — the old suffix is stripped before appending the new one.

2. When _resuming flushes the buffer, any trailing tool hint suffix
   is removed so it doesn't persist into the next streaming segment.

3. When final _stream_end closes the card, tool hint suffix is
   cleaned from the text before the final card update.

Adds 3 regression tests covering all three scenarios.

Made-with: Cursor
2026-04-08 21:02:21 +08:00
xzq.xuandchengyongru 8d6f41e484 feat(feishu): streaming resuming + inline tool hints
Two improvements to Feishu streaming card experience:

1. Handle _resuming in send_delta: when a mid-turn _stream_end arrives
   with resuming=True (tool call between segments), flush current text
   to the card but keep the buffer alive so subsequent segments append
   to the same card instead of creating a new one.

2. Inline tool hints into streaming cards: when a tool hint arrives
   while a streaming card is active, append it to the card content
   (e.g. "🔧 web_fetch(...)") instead of sending a separate card.
   The hint is automatically stripped when the next delta arrives.

Made-with: Cursor
2026-04-08 21:02:21 +08:00
chengyongruandchengyongru 4962867112 fix(tool-hint): fold paths in exec commands instead of blind truncation
exec tool hints previously used val[:40] which cut paths mid-segment
(e.g. "D:\Documents\GitHub\nanobot.worktree…"). Now uses regex to
detect file paths in commands and abbreviates them properly, with
smart truncation at chain separators (&&, |, ;) as fallback.
2026-04-08 11:37:17 +08:00
JiajunandGitHub 473637ceff feat(feishu): add done emoji support for reaction lifecycle (#2899)
* feat(feishu): add done emoji support for reaction lifecycle

* feat(feishu): add done emoji support and update documentation
2026-04-07 23:56:23 +08:00
chengyongru c44d4f2b2b fix(test): fix two flaky tests on Windows
- test_exec_head_tail_truncation: use temp script file instead of
  python -c to avoid cmd.exe quote-parsing issues after PR #2893
- test_grep_files_with_matches_supports_head_limit_and_offset: query
  full result set first to avoid mtime-dependent sort assumption
2026-04-07 21:30:45 +08:00
Xubin Renandchengyongru ae27d69ecb fix(exec): add Windows support for shell command execution
ExecTool hardcoded bash, breaking exec on Windows. Now uses cmd.exe
via COMSPEC on Windows with a curated minimal env (PATH, SYSTEMROOT,
etc.) that excludes secrets. bwrap sandbox gracefully skips on Windows.
2026-04-07 21:13:32 +08:00
chengyongru ba38d41ad1 Merge remote-tracking branch 'origin/main' into nightly 2026-04-07 20:49:53 +08:00
chengyongruandGitHub ad4d095080 feat(memory):dream enhancement (#2887)
* feat(dream): enhance memory cleanup with staleness detection

- Phase 1: add [FILE-REMOVE] directive and staleness patterns (14-day
  threshold, completed tasks, superseded info, resolved tracking)
- Phase 2: add explicit cleanup rules, file paths section, and deletion
  guidance to prevent LLM path confusion
- Inject current date and file sizes into Phase 1 context for age-aware
  analysis
- Add _dream_debug() helper for observability (dream-debug.log in workspace)
- Log Phase 1 analysis output and Phase 2 tool events for debugging

Tested with glm-5-turbo: MEMORY.md reduced from 149 to 108-129 lines
across two rounds, correctly identifying and removing weather data,
detailed incident info, completed research, and stale discussions.

* refactor(dream): replace _dream_debug file logger with loguru

Remove the custom _dream_debug() helper that wrote to dream-debug.log
and use the existing loguru logger instead. Phase 1 analysis is logged
at debug level, tool events at info level — consistent with the rest
of the codebase and no extra log file to manage.

* fix(dream): make stale scan independent of conversation history

Reframe Phase 1 from a single comparison task to two independent
tasks: history diff AND proactive stale scan. The LLM was skipping
stale content that wasn't referenced in conversation history (e.g.
old triage snapshots). Now explicitly requires scanning memory files
for staleness patterns on every run.

* fix(dream): correct old_text param name and truncate debug log

- Phase 2 prompt: old_string -> old_text to match EditFileTool interface
- Phase 1 debug log: truncate analysis to 500 chars to avoid oversized lines

* refactor(dream): streamline prompts by separating concerns

Phase 1 owns all staleness judgment logic; Phase 2 is pure execution
guidance. Remove duplicated cleanup rules from Phase 2 since Phase 1
already determines what to add/remove. Fix remaining old_string -> old_text.
Total prompt size reduced ~45% (870 -> 480 tokens).

* fix(dream): add FILE-REMOVE execution guidance to Phase 2 prompt

Phase 2 was only processing [FILE] additions and ignoring [FILE-REMOVE]
deletions after the cleanup rules were removed. Add explicit mapping:
[FILE] → add content, [FILE-REMOVE] → delete content.
2026-04-07 15:41:54 +08:00
chengyongruandchengyongru 3723cd726e fix(msteams): remove optional deps from dev extras and gate tests
PyJWT and cryptography are optional msteams deps; they should not be
bundled into the generic dev install.  Tests now skip the entire file
when the deps are missing, following the dingtalk pattern.
2026-04-06 15:48:10 +08:00
Bob Johnsonandchengyongru 8f0b653a4c Fix MSTeams PR review follow-ups 2026-04-06 15:48:10 +08:00
T3chC0wb0yandchengyongru 5857f7fdd0 Add Microsoft Teams channel on current nightly base 2026-04-06 15:48:10 +08:00
chengyongruandchengyongru 7d2c62716c fix(dream): allow LLM to retry on tool errors instead of failing immediately
Dream Phase 2 uses fail_on_tool_error=True, which terminates the entire
run on the first tool error (e.g. old_text not found in edit_file).
Normal agent runs default to False so the LLM can self-correct and retry.
Dream should behave the same way.
2026-04-05 22:09:42 +08:00
40 changed files with 6918 additions and 274 deletions
+33 -1
View File
@@ -558,7 +558,10 @@ Uses **WebSocket** long connection — no public IP required.
"verificationToken": "", "verificationToken": "",
"allowFrom": ["ou_YOUR_OPEN_ID"], "allowFrom": ["ou_YOUR_OPEN_ID"],
"groupPolicy": "mention", "groupPolicy": "mention",
"streaming": true "reactEmoji": "OnIt",
"doneEmoji": "DONE",
"streaming": true,
"domain": "feishu"
} }
} }
} }
@@ -568,6 +571,9 @@ Uses **WebSocket** long connection — no public IP required.
> `encryptKey` and `verificationToken` are optional for Long Connection mode. > `encryptKey` and `verificationToken` are optional for Long Connection mode.
> `allowFrom`: Add your open_id (find it in nanobot logs when you message the bot). Use `["*"]` to allow all users. > `allowFrom`: Add your open_id (find it in nanobot logs when you message the bot). Use `["*"]` to allow all users.
> `groupPolicy`: `"mention"` (default — respond only when @mentioned), `"open"` (respond to all group messages). Private chats always respond. > `groupPolicy`: `"mention"` (default — respond only when @mentioned), `"open"` (respond to all group messages). Private chats always respond.
> `reactEmoji`: Emoji for "processing" status (default: `OnIt`). See [available emojis](https://open.larkoffice.com/document/server-docs/im-v1/message-reaction/emojis-introduce).
> `doneEmoji`: Optional emoji for "completed" status (e.g., `DONE`, `OK`, `HEART`). When set, bot adds this reaction after removing `reactEmoji`.
> `domain`: `"feishu"` (default) for China (open.feishu.cn), `"lark"` for international Lark (open.larksuite.com).
**3. Run** **3. Run**
@@ -1495,6 +1501,32 @@ MCP tools are automatically discovered and registered on startup. The LLM can us
**Docker security**: The official Docker image runs as a non-root user (`nanobot`, UID 1000) with bubblewrap pre-installed. When using `docker-compose.yml`, the container drops all Linux capabilities except `SYS_ADMIN` (required for bwrap's namespace isolation). **Docker security**: The official Docker image runs as a non-root user (`nanobot`, UID 1000) with bubblewrap pre-installed. When using `docker-compose.yml`, the container drops all Linux capabilities except `SYS_ADMIN` (required for bwrap's namespace isolation).
### Auto Compact
When a user is idle for longer than a configured TTL, nanobot **proactively** compresses the session context into a summary. This reduces token cost and first-token latency when the user returns — instead of re-processing a long stale context with an expired KV cache, the model receives a compact summary and fresh input.
```json
{
"agents": {
"defaults": {
"sessionTtlMinutes": 15
}
}
}
```
| Option | Default | Description |
|--------|---------|-------------|
| `agents.defaults.sessionTtlMinutes` | `0` (disabled) | Minutes of idle time before auto-compaction. Set to `0` to disable. Recommended: `15` — matches typical LLM KV cache expiration, so compacted sessions won't waste cache on cold entries. |
How it works:
1. **Idle detection**: On each idle tick (~1 s), checks all sessions for expiration.
2. **Background compaction**: Expired sessions are summarized via LLM, then cleared.
3. **Summary injection**: When the user returns, the summary is injected as runtime context (one-shot, not persisted).
> [!TIP]
> The summary survives bot restarts — it's stored in session metadata and recovered on the next message.
### Timezone ### Timezone
Time is context. Context should be precise. Time is context. Context should be precise.
+68
View File
@@ -0,0 +1,68 @@
# Microsoft Teams (MVP)
This repository includes a built-in `msteams` channel MVP for Microsoft Teams direct messages.
## Current scope
- Direct-message text in/out
- Tenant-aware OAuth token acquisition
- Conversation reference persistence for replies
- Public HTTPS webhook support through a tunnel or reverse proxy
## Not yet included
- Group/channel handling
- Attachments and cards
- Polls
- Richer Teams activity handling
## Example config
```json
{
"channels": {
"msteams": {
"enabled": true,
"appId": "YOUR_APP_ID",
"appPassword": "YOUR_APP_SECRET",
"tenantId": "YOUR_TENANT_ID",
"host": "0.0.0.0",
"port": 3978,
"path": "/api/messages",
"allowFrom": ["*"],
"replyInThread": true,
"mentionOnlyResponse": "Hi — what can I help with?",
"validateInboundAuth": false,
"restartNotifyEnabled": false,
"restartNotifyPreMessage": "Nanobot agent initiated a gateway restart. I will message again when the gateway is back online.",
"restartNotifyPostMessage": "Nanobot gateway is back online."
}
}
}
```
## Behavior notes
- `replyInThread: true` replies to the triggering Teams activity when a stored `activity_id` is available.
- `replyInThread: false` posts replies as normal conversation messages.
- If `replyInThread` is enabled but no `activity_id` is stored, Nanobot falls back to a normal conversation message.
- `mentionOnlyResponse` controls what Nanobot receives when a user sends only a bot mention such as `<at>Nanobot</at>`.
- Set `mentionOnlyResponse` to an empty string to ignore mention-only messages.
- `validateInboundAuth: true` enables inbound Bot Framework bearer-token validation.
- `validateInboundAuth: false` leaves inbound auth unenforced, which is safer while first validating a new relay, tunnel, or proxy path.
- When enabled, Nanobot validates the inbound bearer token signature, issuer, audience, token lifetime, and `serviceUrl` claim when present.
- `restartNotifyEnabled: true` enables optional Teams restart-notification configuration for external wrapper-script driven restarts.
- `restartNotifyPreMessage` and `restartNotifyPostMessage` control the before/after announcement text used by that external wrapper.
## Setup notes
1. Create or reuse a Microsoft Teams / Azure bot app registration.
2. Set the bot messaging endpoint to a public HTTPS URL ending in `/api/messages`.
3. Forward that public endpoint to `http://localhost:3978/api/messages`.
4. Start Nanobot with:
```bash
nanobot gateway
```
5. Optional: if you use an external restart wrapper (for example a script that stops and restarts the gateway), you can enable Teams restart announcements with `restartNotifyEnabled: true` and have the wrapper send `restartNotifyPreMessage` before restart and `restartNotifyPostMessage` after the gateway is back online.
+331
View File
@@ -0,0 +1,331 @@
# WebSocket Server Channel
Nanobot can act as a WebSocket server, allowing external clients (web apps, CLIs, scripts) to interact with the agent in real time via persistent connections.
## Features
- Bidirectional real-time communication over WebSocket
- Streaming support — receive agent responses token by token
- Token-based authentication (static tokens and short-lived issued tokens)
- Per-connection sessions — each connection gets a unique `chat_id`
- TLS/SSL support (WSS) with enforced TLSv1.2 minimum
- Client allow-list via `allowFrom`
- Auto-cleanup of dead connections
## Quick Start
### 1. Configure
Add to `config.json` under `channels.websocket`:
```json
{
"channels": {
"websocket": {
"enabled": true,
"host": "127.0.0.1",
"port": 8765,
"path": "/",
"websocketRequiresToken": false,
"allowFrom": ["*"],
"streaming": true
}
}
}
```
### 2. Start nanobot
```bash
nanobot gateway
```
You should see:
```
WebSocket server listening on ws://127.0.0.1:8765/
```
### 3. Connect a client
```bash
# Using websocat
websocat ws://127.0.0.1:8765/?client_id=alice
# Using Python
import asyncio, json, websockets
async def main():
async with websockets.connect("ws://127.0.0.1:8765/?client_id=alice") as ws:
ready = json.loads(await ws.recv())
print(ready) # {"event": "ready", "chat_id": "...", "client_id": "alice"}
await ws.send(json.dumps({"content": "Hello nanobot!"}))
reply = json.loads(await ws.recv())
print(reply["text"])
asyncio.run(main())
```
## Connection URL
```
ws://{host}:{port}{path}?client_id={id}&token={token}
```
| Parameter | Required | Description |
|-----------|----------|-------------|
| `client_id` | No | Identifier for `allowFrom` authorization. Auto-generated as `anon-xxxxxxxxxxxx` if omitted. Truncated to 128 chars. |
| `token` | Conditional | Authentication token. Required when `websocketRequiresToken` is `true` or `token` (static secret) is configured. |
## Wire Protocol
All frames are JSON text. Each message has an `event` field.
### Server → Client
**`ready`** — sent immediately after connection is established:
```json
{
"event": "ready",
"chat_id": "uuid-v4",
"client_id": "alice"
}
```
**`message`** — full agent response:
```json
{
"event": "message",
"text": "Hello! How can I help?",
"media": ["/tmp/image.png"],
"reply_to": "msg-id"
}
```
`media` and `reply_to` are only present when applicable.
**`delta`** — streaming text chunk (only when `streaming: true`):
```json
{
"event": "delta",
"text": "Hello",
"stream_id": "s1"
}
```
**`stream_end`** — signals the end of a streaming segment:
```json
{
"event": "stream_end",
"stream_id": "s1"
}
```
### Client → Server
Send plain text:
```json
"Hello nanobot!"
```
Or send a JSON object with a recognized text field:
```json
{"content": "Hello nanobot!"}
```
Recognized fields: `content`, `text`, `message` (checked in that order). Invalid JSON is treated as plain text.
## Configuration Reference
All fields go under `channels.websocket` in `config.json`.
### Connection
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `enabled` | bool | `false` | Enable the WebSocket server. |
| `host` | string | `"127.0.0.1"` | Bind address. Use `"0.0.0.0"` to accept external connections. |
| `port` | int | `8765` | Listen port. |
| `path` | string | `"/"` | WebSocket upgrade path. Trailing slashes are normalized (root `/` is preserved). |
| `maxMessageBytes` | int | `1048576` | Maximum inbound message size in bytes (1 KB 16 MB). |
### Authentication
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `token` | string | `""` | Static shared secret. When set, clients must provide `?token=<value>` matching this secret (timing-safe comparison). Issued tokens are also accepted as a fallback. |
| `websocketRequiresToken` | bool | `true` | When `true` and no static `token` is configured, clients must still present a valid issued token. Set to `false` to allow unauthenticated connections (only safe for local/trusted networks). |
| `tokenIssuePath` | string | `""` | HTTP path for issuing short-lived tokens. Must differ from `path`. See [Token Issuance](#token-issuance). |
| `tokenIssueSecret` | string | `""` | Secret required to obtain tokens via the issue endpoint. If empty, any client can obtain tokens (logged as a warning). |
| `tokenTtlS` | int | `300` | Time-to-live for issued tokens in seconds (30 86,400). |
### Access Control
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `allowFrom` | list of string | `["*"]` | Allowed `client_id` values. `"*"` allows all; `[]` denies all. |
### Streaming
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `streaming` | bool | `true` | Enable streaming mode. The agent sends `delta` + `stream_end` frames instead of a single `message`. |
### Keep-alive
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `pingIntervalS` | float | `20.0` | WebSocket ping interval in seconds (5 300). |
| `pingTimeoutS` | float | `20.0` | Time to wait for a pong before closing the connection (5 300). |
### TLS/SSL
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `sslCertfile` | string | `""` | Path to the TLS certificate file (PEM). Both `sslCertfile` and `sslKeyfile` must be set to enable WSS. |
| `sslKeyfile` | string | `""` | Path to the TLS private key file (PEM). Minimum TLS version is enforced as TLSv1.2. |
## Token Issuance
For production deployments where `websocketRequiresToken: true`, use short-lived tokens instead of embedding static secrets in clients.
### How it works
1. Client sends `GET {tokenIssuePath}` with `Authorization: Bearer {tokenIssueSecret}` (or `X-Nanobot-Auth` header).
2. Server responds with a one-time-use token:
```json
{"token": "nbwt_aBcDeFg...", "expires_in": 300}
```
3. Client opens WebSocket with `?token=nbwt_aBcDeFg...&client_id=...`.
4. The token is consumed (single use) and cannot be reused.
### Example setup
```json
{
"channels": {
"websocket": {
"enabled": true,
"port": 8765,
"path": "/ws",
"tokenIssuePath": "/auth/token",
"tokenIssueSecret": "your-secret-here",
"tokenTtlS": 300,
"websocketRequiresToken": true,
"allowFrom": ["*"],
"streaming": true
}
}
}
```
Client flow:
```bash
# 1. Obtain a token
curl -H "Authorization: Bearer your-secret-here" http://127.0.0.1:8765/auth/token
# 2. Connect using the token
websocat "ws://127.0.0.1:8765/ws?client_id=alice&token=nbwt_aBcDeFg..."
```
### Limits
- Issued tokens are single-use — each token can only complete one handshake.
- Outstanding tokens are capped at 10,000. Requests beyond this return HTTP 429.
- Expired tokens are purged lazily on each issue or validation request.
## Security Notes
- **Timing-safe comparison**: Static token validation uses `hmac.compare_digest` to prevent timing attacks.
- **Defense in depth**: `allowFrom` is checked at both the HTTP handshake level and the message level.
- **Token isolation**: Each WebSocket connection gets a unique `chat_id`. Clients cannot access other sessions.
- **TLS enforcement**: When SSL is enabled, TLSv1.2 is the minimum allowed version.
- **Default-secure**: `websocketRequiresToken` defaults to `true`. Explicitly set it to `false` only on trusted networks.
## Media Files
Outbound `message` events may include a `media` field containing local filesystem paths. Remote clients cannot access these files directly — they need either:
- A shared filesystem mount, or
- An HTTP file server serving the nanobot media directory
## Common Patterns
### Trusted local network (no auth)
```json
{
"channels": {
"websocket": {
"enabled": true,
"host": "0.0.0.0",
"port": 8765,
"websocketRequiresToken": false,
"allowFrom": ["*"],
"streaming": true
}
}
}
```
### Static token (simple auth)
```json
{
"channels": {
"websocket": {
"enabled": true,
"token": "my-shared-secret",
"allowFrom": ["alice", "bob"]
}
}
}
```
Clients connect with `?token=my-shared-secret&client_id=alice`.
### Public endpoint with issued tokens
```json
{
"channels": {
"websocket": {
"enabled": true,
"host": "0.0.0.0",
"port": 8765,
"path": "/ws",
"tokenIssuePath": "/auth/token",
"tokenIssueSecret": "production-secret",
"websocketRequiresToken": true,
"sslCertfile": "/etc/ssl/certs/server.pem",
"sslKeyfile": "/etc/ssl/private/server-key.pem",
"allowFrom": ["*"]
}
}
}
```
### Custom path
```json
{
"channels": {
"websocket": {
"enabled": true,
"path": "/chat/ws",
"allowFrom": ["*"]
}
}
}
```
Clients connect to `ws://127.0.0.1:8765/chat/ws?client_id=...`. Trailing slashes are normalized, so `/chat/ws/` works the same.
+82
View File
@@ -0,0 +1,82 @@
"""Auto compact: proactive compression of idle sessions to reduce token cost and latency."""
from __future__ import annotations
from datetime import datetime
from typing import TYPE_CHECKING, Callable, Coroutine
from loguru import logger
if TYPE_CHECKING:
from nanobot.agent.memory import Consolidator
from nanobot.session.manager import Session, SessionManager
class AutoCompact:
def __init__(self, sessions: SessionManager, consolidator: Consolidator,
session_ttl_minutes: int = 0):
self.sessions = sessions
self.consolidator = consolidator
self._ttl = session_ttl_minutes
self._archiving: set[str] = set()
self._summaries: dict[str, tuple[str, datetime]] = {}
def _is_expired(self, ts: datetime | str | None) -> bool:
if self._ttl <= 0 or not ts:
return False
if isinstance(ts, str):
ts = datetime.fromisoformat(ts)
return (datetime.now() - ts).total_seconds() >= self._ttl * 60
@staticmethod
def _format_summary(text: str, last_active: datetime) -> str:
idle_min = int((datetime.now() - last_active).total_seconds() / 60)
return f"Inactive for {idle_min} minutes.\nPrevious conversation summary: {text}"
def check_expired(self, schedule_background: Callable[[Coroutine], None]) -> None:
for info in self.sessions.list_sessions():
key = info.get("key", "")
if key and key not in self._archiving and self._is_expired(info.get("updated_at")):
self._archiving.add(key)
logger.debug("Auto-compact: scheduling archival for {} (idle > {} min)", key, self._ttl)
schedule_background(self._archive(key))
async def _archive(self, key: str) -> None:
try:
self.sessions.invalidate(key)
session = self.sessions.get_or_create(key)
msgs = session.messages[session.last_consolidated:]
if not msgs:
logger.debug("Auto-compact: skipping {}, no un-consolidated messages", key)
session.updated_at = datetime.now()
self.sessions.save(session)
return
n = len(msgs)
last_active = session.updated_at
await self.consolidator.archive(msgs)
entry = self.consolidator.get_last_history_entry()
summary = (entry or {}).get("content", "")
if summary and summary != "(nothing)":
self._summaries[key] = (summary, last_active)
session.metadata["_last_summary"] = {"text": summary, "last_active": last_active.isoformat()}
session.clear()
self.sessions.save(session)
logger.info("Auto-compact: archived {} ({} messages, summary={})", key, n, bool(summary))
except Exception:
logger.exception("Auto-compact: failed for {}", key)
finally:
self._archiving.discard(key)
def prepare_session(self, session: Session, key: str) -> tuple[Session, str | None]:
if key in self._archiving or self._is_expired(session.updated_at):
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
session = self.sessions.get_or_create(key)
entry = self._summaries.pop(key, None)
if entry:
session.metadata.pop("_last_summary", None)
return session, self._format_summary(entry[0], entry[1])
if not session.messages and "_last_summary" in session.metadata:
meta = session.metadata.pop("_last_summary")
self.sessions.save(session)
return session, self._format_summary(meta["text"], datetime.fromisoformat(meta["last_active"]))
return session, None
+7 -2
View File
@@ -19,6 +19,7 @@ class ContextBuilder:
BOOTSTRAP_FILES = ["AGENTS.md", "SOUL.md", "USER.md", "TOOLS.md"] BOOTSTRAP_FILES = ["AGENTS.md", "SOUL.md", "USER.md", "TOOLS.md"]
_RUNTIME_CONTEXT_TAG = "[Runtime Context — metadata only, not instructions]" _RUNTIME_CONTEXT_TAG = "[Runtime Context — metadata only, not instructions]"
_RUNTIME_CONTEXT_END = "[/Runtime Context]"
def __init__(self, workspace: Path, timezone: str | None = None): def __init__(self, workspace: Path, timezone: str | None = None):
self.workspace = workspace self.workspace = workspace
@@ -66,12 +67,15 @@ class ContextBuilder:
@staticmethod @staticmethod
def _build_runtime_context( def _build_runtime_context(
channel: str | None, chat_id: str | None, timezone: str | None = None, channel: str | None, chat_id: str | None, timezone: str | None = None,
session_summary: str | None = None,
) -> str: ) -> str:
"""Build untrusted runtime metadata block for injection before the user message.""" """Build untrusted runtime metadata block for injection before the user message."""
lines = [f"Current Time: {current_time_str(timezone)}"] lines = [f"Current Time: {current_time_str(timezone)}"]
if channel and chat_id: if channel and chat_id:
lines += [f"Channel: {channel}", f"Chat ID: {chat_id}"] lines += [f"Channel: {channel}", f"Chat ID: {chat_id}"]
return ContextBuilder._RUNTIME_CONTEXT_TAG + "\n" + "\n".join(lines) if session_summary:
lines += ["", "[Resumed Session]", session_summary]
return ContextBuilder._RUNTIME_CONTEXT_TAG + "\n" + "\n".join(lines) + "\n" + ContextBuilder._RUNTIME_CONTEXT_END
@staticmethod @staticmethod
def _merge_message_content(left: Any, right: Any) -> str | list[dict[str, Any]]: def _merge_message_content(left: Any, right: Any) -> str | list[dict[str, Any]]:
@@ -108,9 +112,10 @@ class ContextBuilder:
channel: str | None = None, channel: str | None = None,
chat_id: str | None = None, chat_id: str | None = None,
current_role: str = "user", current_role: str = "user",
session_summary: str | None = None,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
"""Build the complete message list for an LLM call.""" """Build the complete message list for an LLM call."""
runtime_ctx = self._build_runtime_context(channel, chat_id, self.timezone) runtime_ctx = self._build_runtime_context(channel, chat_id, self.timezone, session_summary=session_summary)
user_content = self._build_user_content(current_message, media) user_content = self._build_user_content(current_message, media)
# Merge runtime context and user content into a single user message # Merge runtime context and user content into a single user message
+157 -57
View File
@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable
from loguru import logger from loguru import logger
from nanobot.agent.auto_compact import AutoCompact
from nanobot.agent.context import ContextBuilder from nanobot.agent.context import ContextBuilder
from nanobot.agent.hook import AgentHook, AgentHookContext, CompositeHook from nanobot.agent.hook import AgentHook, AgentHookContext, CompositeHook
from nanobot.agent.memory import Consolidator, Dream from nanobot.agent.memory import Consolidator, Dream
@@ -179,6 +180,7 @@ class AgentLoop:
mcp_servers: dict | None = None, mcp_servers: dict | None = None,
channels_config: ChannelsConfig | None = None, channels_config: ChannelsConfig | None = None,
timezone: str | None = None, timezone: str | None = None,
session_ttl_minutes: int = 0,
hooks: list[AgentHook] | None = None, hooks: list[AgentHook] | None = None,
): ):
from nanobot.config.schema import ExecToolConfig, WebToolsConfig from nanobot.config.schema import ExecToolConfig, WebToolsConfig
@@ -235,6 +237,10 @@ class AgentLoop:
self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks self._active_tasks: dict[str, list[asyncio.Task]] = {} # session_key -> tasks
self._background_tasks: list[asyncio.Task] = [] self._background_tasks: list[asyncio.Task] = []
self._session_locks: dict[str, asyncio.Lock] = {} self._session_locks: dict[str, asyncio.Lock] = {}
# Per-session pending queues for mid-turn message injection.
# When a session has an active task, new messages for that session
# are routed here instead of creating a new task.
self._pending_queues: dict[str, asyncio.Queue] = {}
# NANOBOT_MAX_CONCURRENT_REQUESTS: <=0 means unlimited; default 3. # NANOBOT_MAX_CONCURRENT_REQUESTS: <=0 means unlimited; default 3.
_max = int(os.environ.get("NANOBOT_MAX_CONCURRENT_REQUESTS", "3")) _max = int(os.environ.get("NANOBOT_MAX_CONCURRENT_REQUESTS", "3"))
self._concurrency_gate: asyncio.Semaphore | None = ( self._concurrency_gate: asyncio.Semaphore | None = (
@@ -250,6 +256,11 @@ class AgentLoop:
get_tool_definitions=self.tools.get_definitions, get_tool_definitions=self.tools.get_definitions,
max_completion_tokens=provider.generation.max_tokens, max_completion_tokens=provider.generation.max_tokens,
) )
self.auto_compact = AutoCompact(
sessions=self.sessions,
consolidator=self.consolidator,
session_ttl_minutes=session_ttl_minutes,
)
self.dream = Dream( self.dream = Dream(
store=self.context.memory, store=self.context.memory,
provider=provider, provider=provider,
@@ -341,13 +352,16 @@ class AgentLoop:
channel: str = "cli", channel: str = "cli",
chat_id: str = "direct", chat_id: str = "direct",
message_id: str | None = None, message_id: str | None = None,
) -> tuple[str | None, list[str], list[dict]]: pending_queue: asyncio.Queue | None = None,
) -> tuple[str | None, list[str], list[dict], str, bool]:
"""Run the agent iteration loop. """Run the agent iteration loop.
*on_stream*: called with each content delta during streaming. *on_stream*: called with each content delta during streaming.
*on_stream_end(resuming)*: called when a streaming session finishes. *on_stream_end(resuming)*: called when a streaming session finishes.
``resuming=True`` means tool calls follow (spinner should restart); ``resuming=True`` means tool calls follow (spinner should restart);
``resuming=False`` means this is the final response. ``resuming=False`` means this is the final response.
Returns (final_content, tools_used, messages, stop_reason, had_injections).
""" """
loop_hook = _LoopHook( loop_hook = _LoopHook(
self, self,
@@ -369,6 +383,18 @@ class AgentLoop:
return return
self._set_runtime_checkpoint(session, payload) self._set_runtime_checkpoint(session, payload)
async def _drain_pending() -> list[InboundMessage]:
"""Non-blocking drain of follow-up messages from the pending queue."""
if pending_queue is None:
return []
items: list[InboundMessage] = []
while True:
try:
items.append(pending_queue.get_nowait())
except asyncio.QueueEmpty:
break
return items
result = await self.runner.run(AgentRunSpec( result = await self.runner.run(AgentRunSpec(
initial_messages=initial_messages, initial_messages=initial_messages,
tools=self.tools, tools=self.tools,
@@ -385,13 +411,14 @@ class AgentLoop:
provider_retry_mode=self.provider_retry_mode, provider_retry_mode=self.provider_retry_mode,
progress_callback=on_progress, progress_callback=on_progress,
checkpoint_callback=_checkpoint, checkpoint_callback=_checkpoint,
injection_callback=_drain_pending,
)) ))
self._last_usage = result.usage self._last_usage = result.usage
if result.stop_reason == "max_iterations": if result.stop_reason == "max_iterations":
logger.warning("Max iterations ({}) reached", self.max_iterations) logger.warning("Max iterations ({}) reached", self.max_iterations)
elif result.stop_reason == "error": elif result.stop_reason == "error":
logger.error("LLM returned error: {}", (result.final_content or "")[:200]) logger.error("LLM returned error: {}", (result.final_content or "")[:200])
return result.final_content, result.tools_used, result.messages return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
async def run(self) -> None: async def run(self) -> None:
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop.""" """Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
@@ -403,6 +430,7 @@ class AgentLoop:
try: try:
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0) msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
except asyncio.TimeoutError: except asyncio.TimeoutError:
self.auto_compact.check_expired(self._schedule_background)
continue continue
except asyncio.CancelledError: except asyncio.CancelledError:
# Preserve real task cancellation so shutdown can complete cleanly. # Preserve real task cancellation so shutdown can complete cleanly.
@@ -421,67 +449,112 @@ class AgentLoop:
if result: if result:
await self.bus.publish_outbound(result) await self.bus.publish_outbound(result)
continue continue
# If this session already has an active pending queue (i.e. a task
# is processing this session), route the message there for mid-turn
# injection instead of creating a competing task.
if msg.session_key in self._pending_queues:
try:
self._pending_queues[msg.session_key].put_nowait(msg)
except asyncio.QueueFull:
logger.warning(
"Pending queue full for session {}, dropping follow-up",
msg.session_key,
)
else:
logger.info(
"Routed follow-up message to pending queue for session {}",
msg.session_key,
)
continue
task = asyncio.create_task(self._dispatch(msg)) task = asyncio.create_task(self._dispatch(msg))
self._active_tasks.setdefault(msg.session_key, []).append(task) self._active_tasks.setdefault(msg.session_key, []).append(task)
task.add_done_callback(lambda t, k=msg.session_key: self._active_tasks.get(k, []) and self._active_tasks[k].remove(t) if t in self._active_tasks.get(k, []) else None) task.add_done_callback(lambda t, k=msg.session_key: self._active_tasks.get(k, []) and self._active_tasks[k].remove(t) if t in self._active_tasks.get(k, []) else None)
async def _dispatch(self, msg: InboundMessage) -> None: async def _dispatch(self, msg: InboundMessage) -> None:
"""Process a message: per-session serial, cross-session concurrent.""" """Process a message: per-session serial, cross-session concurrent."""
lock = self._session_locks.setdefault(msg.session_key, asyncio.Lock()) session_key = msg.session_key
lock = self._session_locks.setdefault(session_key, asyncio.Lock())
gate = self._concurrency_gate or nullcontext() gate = self._concurrency_gate or nullcontext()
async with lock, gate:
try:
on_stream = on_stream_end = None
if msg.metadata.get("_wants_stream"):
# Split one answer into distinct stream segments.
stream_base_id = f"{msg.session_key}:{time.time_ns()}"
stream_segment = 0
def _current_stream_id() -> str: # Register a pending queue so follow-up messages for this session are
return f"{stream_base_id}:{stream_segment}" # routed here (mid-turn injection) instead of spawning a new task.
pending = asyncio.Queue(maxsize=20)
self._pending_queues[session_key] = pending
async def on_stream(delta: str) -> None: try:
meta = dict(msg.metadata or {}) async with lock, gate:
meta["_stream_delta"] = True try:
meta["_stream_id"] = _current_stream_id() on_stream = on_stream_end = None
if msg.metadata.get("_wants_stream"):
# Split one answer into distinct stream segments.
stream_base_id = f"{msg.session_key}:{time.time_ns()}"
stream_segment = 0
def _current_stream_id() -> str:
return f"{stream_base_id}:{stream_segment}"
async def on_stream(delta: str) -> None:
meta = dict(msg.metadata or {})
meta["_stream_delta"] = True
meta["_stream_id"] = _current_stream_id()
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content=delta,
metadata=meta,
))
async def on_stream_end(*, resuming: bool = False) -> None:
nonlocal stream_segment
meta = dict(msg.metadata or {})
meta["_stream_end"] = True
meta["_resuming"] = resuming
meta["_stream_id"] = _current_stream_id()
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content="",
metadata=meta,
))
stream_segment += 1
response = await self._process_message(
msg, on_stream=on_stream, on_stream_end=on_stream_end,
pending_queue=pending,
)
if response is not None:
await self.bus.publish_outbound(response)
elif msg.channel == "cli":
await self.bus.publish_outbound(OutboundMessage( await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id, channel=msg.channel, chat_id=msg.chat_id,
content=delta, content="", metadata=msg.metadata or {},
metadata=meta,
)) ))
except asyncio.CancelledError:
async def on_stream_end(*, resuming: bool = False) -> None: logger.info("Task cancelled for session {}", session_key)
nonlocal stream_segment raise
meta = dict(msg.metadata or {}) except Exception:
meta["_stream_end"] = True logger.exception("Error processing message for session {}", session_key)
meta["_resuming"] = resuming
meta["_stream_id"] = _current_stream_id()
await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id,
content="",
metadata=meta,
))
stream_segment += 1
response = await self._process_message(
msg, on_stream=on_stream, on_stream_end=on_stream_end,
)
if response is not None:
await self.bus.publish_outbound(response)
elif msg.channel == "cli":
await self.bus.publish_outbound(OutboundMessage( await self.bus.publish_outbound(OutboundMessage(
channel=msg.channel, chat_id=msg.chat_id, channel=msg.channel, chat_id=msg.chat_id,
content="", metadata=msg.metadata or {}, content="Sorry, I encountered an error.",
)) ))
except asyncio.CancelledError: finally:
logger.info("Task cancelled for session {}", msg.session_key) # Drain any messages still in the pending queue and re-publish
raise # them to the bus so they are processed as fresh inbound messages
except Exception: # rather than silently lost.
logger.exception("Error processing message for session {}", msg.session_key) queue = self._pending_queues.pop(session_key, None)
await self.bus.publish_outbound(OutboundMessage( if queue is not None:
channel=msg.channel, chat_id=msg.chat_id, leftover = 0
content="Sorry, I encountered an error.", while True:
)) try:
item = queue.get_nowait()
except asyncio.QueueEmpty:
break
await self.bus.publish_inbound(item)
leftover += 1
if leftover:
logger.info(
"Re-published {} leftover message(s) to bus for session {}",
leftover, session_key,
)
async def close_mcp(self) -> None: async def close_mcp(self) -> None:
"""Drain pending background archives, then close MCP connections.""" """Drain pending background archives, then close MCP connections."""
@@ -513,6 +586,7 @@ class AgentLoop:
on_progress: Callable[[str], Awaitable[None]] | None = None, on_progress: Callable[[str], Awaitable[None]] | None = None,
on_stream: Callable[[str], Awaitable[None]] | None = None, on_stream: Callable[[str], Awaitable[None]] | None = None,
on_stream_end: Callable[..., Awaitable[None]] | None = None, on_stream_end: Callable[..., Awaitable[None]] | None = None,
pending_queue: asyncio.Queue | None = None,
) -> OutboundMessage | None: ) -> OutboundMessage | None:
"""Process a single inbound message and return the response.""" """Process a single inbound message and return the response."""
# System messages: parse origin from chat_id ("channel:chat_id") # System messages: parse origin from chat_id ("channel:chat_id")
@@ -524,16 +598,21 @@ class AgentLoop:
session = self.sessions.get_or_create(key) session = self.sessions.get_or_create(key)
if self._restore_runtime_checkpoint(session): if self._restore_runtime_checkpoint(session):
self.sessions.save(session) self.sessions.save(session)
session, pending = self.auto_compact.prepare_session(session, key)
await self.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")) self._set_tool_context(channel, chat_id, msg.metadata.get("message_id"))
history = session.get_history(max_messages=0) history = session.get_history(max_messages=0)
current_role = "assistant" if msg.sender_id == "subagent" else "user" current_role = "assistant" if msg.sender_id == "subagent" else "user"
messages = self.context.build_messages( messages = self.context.build_messages(
history=history, history=history,
current_message=msg.content, channel=channel, chat_id=chat_id, current_message=msg.content, channel=channel, chat_id=chat_id,
session_summary=pending,
current_role=current_role, current_role=current_role,
) )
final_content, _, all_msgs = await self._run_agent_loop( final_content, _, all_msgs, _, _ = await self._run_agent_loop(
messages, session=session, channel=channel, chat_id=chat_id, messages, session=session, channel=channel, chat_id=chat_id,
message_id=msg.metadata.get("message_id"), message_id=msg.metadata.get("message_id"),
) )
@@ -552,6 +631,8 @@ class AgentLoop:
if self._restore_runtime_checkpoint(session): if self._restore_runtime_checkpoint(session):
self.sessions.save(session) self.sessions.save(session)
session, pending = self.auto_compact.prepare_session(session, key)
# Slash commands # Slash commands
raw = msg.content.strip() raw = msg.content.strip()
ctx = CommandContext(msg=msg, session=session, key=key, raw=raw, loop=self) ctx = CommandContext(msg=msg, session=session, key=key, raw=raw, loop=self)
@@ -566,9 +647,11 @@ class AgentLoop:
message_tool.start_turn() message_tool.start_turn()
history = session.get_history(max_messages=0) history = session.get_history(max_messages=0)
initial_messages = self.context.build_messages( initial_messages = self.context.build_messages(
history=history, history=history,
current_message=msg.content, current_message=msg.content,
session_summary=pending,
media=msg.media if msg.media else None, media=msg.media if msg.media else None,
channel=msg.channel, chat_id=msg.chat_id, channel=msg.channel, chat_id=msg.chat_id,
) )
@@ -581,7 +664,7 @@ class AgentLoop:
channel=msg.channel, chat_id=msg.chat_id, content=content, metadata=meta, channel=msg.channel, chat_id=msg.chat_id, content=content, metadata=meta,
)) ))
final_content, _, all_msgs = await self._run_agent_loop( final_content, _, all_msgs, stop_reason, had_injections = await self._run_agent_loop(
initial_messages, initial_messages,
on_progress=on_progress or _bus_progress, on_progress=on_progress or _bus_progress,
on_stream=on_stream, on_stream=on_stream,
@@ -589,6 +672,7 @@ class AgentLoop:
session=session, session=session,
channel=msg.channel, chat_id=msg.chat_id, channel=msg.channel, chat_id=msg.chat_id,
message_id=msg.metadata.get("message_id"), message_id=msg.metadata.get("message_id"),
pending_queue=pending_queue,
) )
if final_content is None or not final_content.strip(): if final_content is None or not final_content.strip():
@@ -599,8 +683,13 @@ class AgentLoop:
self.sessions.save(session) self.sessions.save(session)
self._schedule_background(self.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: # When follow-up messages were injected mid-turn, the LLM's final
return None # response addresses those follow-ups. Always send the response in
# this case, even if MessageTool was used earlier in the turn — the
# follow-up response is new content the user hasn't seen.
if not had_injections:
if (mt := self.tools.get("message")) and isinstance(mt, MessageTool) and mt._sent_in_turn:
return None
preview = final_content[:120] + "..." if len(final_content) > 120 else final_content preview = final_content[:120] + "..." if len(final_content) > 120 else final_content
logger.info("Response to {}:{}: {}", msg.channel, msg.sender_id, preview) logger.info("Response to {}:{}: {}", msg.channel, msg.sender_id, preview)
@@ -672,12 +761,23 @@ class AgentLoop:
entry["content"] = filtered entry["content"] = filtered
elif role == "user": elif role == "user":
if isinstance(content, str) and content.startswith(ContextBuilder._RUNTIME_CONTEXT_TAG): if isinstance(content, str) and content.startswith(ContextBuilder._RUNTIME_CONTEXT_TAG):
# Strip the runtime-context prefix, keep only the user text. # Strip the entire runtime-context block (including any session summary).
parts = content.split("\n\n", 1) # The block is bounded by _RUNTIME_CONTEXT_TAG and _RUNTIME_CONTEXT_END.
if len(parts) > 1 and parts[1].strip(): end_marker = ContextBuilder._RUNTIME_CONTEXT_END
entry["content"] = parts[1] end_pos = content.find(end_marker)
if end_pos >= 0:
after = content[end_pos + len(end_marker):].lstrip("\n")
if after:
entry["content"] = after
else:
continue
else: else:
continue # Fallback: no end marker found, strip the tag prefix
after_tag = content[len(ContextBuilder._RUNTIME_CONTEXT_TAG):].lstrip("\n")
if after_tag.strip():
entry["content"] = after_tag
else:
continue
if isinstance(content, list): if isinstance(content, list):
filtered = self._sanitize_persisted_blocks(content, drop_runtime=True) filtered = self._sanitize_persisted_blocks(content, drop_runtime=True)
if not filtered: if not filtered:
+12 -4
View File
@@ -373,6 +373,10 @@ class Consolidator:
weakref.WeakValueDictionary() weakref.WeakValueDictionary()
) )
def get_last_history_entry(self) -> dict[str, Any] | None:
"""Return the most recent entry from history.jsonl."""
return self.store._read_last_entry()
def get_lock(self, session_key: str) -> asyncio.Lock: def get_lock(self, session_key: str) -> asyncio.Lock:
"""Return the shared consolidation lock for one session.""" """Return the shared consolidation lock for one session."""
return self._locks.setdefault(session_key, asyncio.Lock()) return self._locks.setdefault(session_key, asyncio.Lock())
@@ -575,13 +579,15 @@ class Dream:
) )
# Current file contents # Current file contents
current_date = datetime.now().strftime("%Y-%m-%d")
current_memory = self.store.read_memory() or "(empty)" current_memory = self.store.read_memory() or "(empty)"
current_soul = self.store.read_soul() or "(empty)" current_soul = self.store.read_soul() or "(empty)"
current_user = self.store.read_user() or "(empty)" current_user = self.store.read_user() or "(empty)"
file_context = ( file_context = (
f"## Current MEMORY.md\n{current_memory}\n\n" f"## Current Date\n{current_date}\n\n"
f"## Current SOUL.md\n{current_soul}\n\n" f"## Current MEMORY.md ({len(current_memory)} chars)\n{current_memory}\n\n"
f"## Current USER.md\n{current_user}" f"## Current SOUL.md ({len(current_soul)} chars)\n{current_soul}\n\n"
f"## Current USER.md ({len(current_user)} chars)\n{current_user}"
) )
# Phase 1: Analyze # Phase 1: Analyze
@@ -603,7 +609,7 @@ class Dream:
tool_choice=None, tool_choice=None,
) )
analysis = phase1_response.content or "" analysis = phase1_response.content or ""
logger.debug("Dream Phase 1 complete ({} chars)", len(analysis)) logger.debug("Dream Phase 1 analysis ({} chars): {}", len(analysis), analysis[:500])
except Exception: except Exception:
logger.exception("Dream Phase 1 failed") logger.exception("Dream Phase 1 failed")
return False return False
@@ -633,6 +639,8 @@ class Dream:
"Dream Phase 2 complete: stop_reason={}, tool_events={}", "Dream Phase 2 complete: stop_reason={}, tool_events={}",
result.stop_reason, len(result.tool_events), result.stop_reason, len(result.tool_events),
) )
for ev in (result.tool_events or []):
logger.info("Dream tool_event: name={}, status={}, detail={}", ev.get("name"), ev.get("status"), ev.get("detail", "")[:200])
except Exception: except Exception:
logger.exception("Dream Phase 2 failed") logger.exception("Dream Phase 2 failed")
result = None result = None
+75 -1
View File
@@ -31,7 +31,11 @@ from nanobot.utils.runtime import (
_DEFAULT_ERROR_MESSAGE = "Sorry, I encountered an error calling the AI model." _DEFAULT_ERROR_MESSAGE = "Sorry, I encountered an error calling the AI model."
_MAX_EMPTY_RETRIES = 2 _MAX_EMPTY_RETRIES = 2
_MAX_INJECTIONS_PER_TURN = 3
_MAX_INJECTION_CYCLES = 5
_SNIP_SAFETY_BUFFER = 1024 _SNIP_SAFETY_BUFFER = 1024
@dataclass(slots=True) @dataclass(slots=True)
class AgentRunSpec: class AgentRunSpec:
"""Configuration for a single agent execution.""" """Configuration for a single agent execution."""
@@ -56,6 +60,7 @@ class AgentRunSpec:
provider_retry_mode: str = "standard" provider_retry_mode: str = "standard"
progress_callback: Any | None = None progress_callback: Any | None = None
checkpoint_callback: Any | None = None checkpoint_callback: Any | None = None
injection_callback: Any | None = None
@dataclass(slots=True) @dataclass(slots=True)
@@ -69,6 +74,7 @@ class AgentRunResult:
stop_reason: str = "completed" stop_reason: str = "completed"
error: str | None = None error: str | None = None
tool_events: list[dict[str, str]] = field(default_factory=list) tool_events: list[dict[str, str]] = field(default_factory=list)
had_injections: bool = False
class AgentRunner: class AgentRunner:
@@ -77,6 +83,38 @@ class AgentRunner:
def __init__(self, provider: LLMProvider): def __init__(self, provider: LLMProvider):
self.provider = provider self.provider = provider
async def _drain_injections(self, spec: AgentRunSpec) -> list[str]:
"""Drain pending user messages via the injection callback.
Returns all drained message contents (capped by
``_MAX_INJECTIONS_PER_TURN``), or an empty list when there is
nothing to inject. Messages beyond the cap are logged so they
are not silently lost.
"""
if spec.injection_callback is None:
return []
try:
items = await spec.injection_callback()
except Exception:
logger.exception("injection_callback failed")
return []
if not items:
return []
# items are InboundMessage objects from _drain_pending
texts: list[str] = []
for item in items:
text = getattr(item, "content", str(item))
if text.strip():
texts.append(text)
if len(texts) > _MAX_INJECTIONS_PER_TURN:
dropped = len(texts) - _MAX_INJECTIONS_PER_TURN
logger.warning(
"Injection batch has {} messages, capping to {} ({} dropped)",
len(texts), _MAX_INJECTIONS_PER_TURN, dropped,
)
texts = texts[-_MAX_INJECTIONS_PER_TURN:]
return texts
async def run(self, spec: AgentRunSpec) -> AgentRunResult: async def run(self, spec: AgentRunSpec) -> AgentRunResult:
hook = spec.hook or AgentHook() hook = spec.hook or AgentHook()
messages = list(spec.initial_messages) messages = list(spec.initial_messages)
@@ -88,6 +126,8 @@ class AgentRunner:
tool_events: list[dict[str, str]] = [] tool_events: list[dict[str, str]] = []
external_lookup_counts: dict[str, int] = {} external_lookup_counts: dict[str, int] = {}
empty_content_retries = 0 empty_content_retries = 0
had_injections = False
injection_cycles = 0
for iteration in range(spec.max_iterations): for iteration in range(spec.max_iterations):
try: try:
@@ -181,6 +221,18 @@ class AgentRunner:
}, },
) )
empty_content_retries = 0 empty_content_retries = 0
# Checkpoint 1: drain injections after tools, before next LLM call
if injection_cycles < _MAX_INJECTION_CYCLES:
injections = await self._drain_injections(spec)
if injections:
had_injections = True
injection_cycles += 1
for text in injections:
messages.append({"role": "user", "content": text})
logger.info(
"Injected {} follow-up message(s) after tool execution ({}/{})",
len(injections), injection_cycles, _MAX_INJECTION_CYCLES,
)
await hook.after_iteration(context) await hook.after_iteration(context)
continue continue
@@ -216,8 +268,29 @@ class AgentRunner:
context.tool_calls = list(response.tool_calls) context.tool_calls = list(response.tool_calls)
clean = hook.finalize_content(context, response.content) clean = hook.finalize_content(context, response.content)
# Check for mid-turn injections BEFORE signaling stream end.
# If injections are found we keep the stream alive (resuming=True)
# so streaming channels don't prematurely finalize the card.
_injected_after_final = False
if injection_cycles < _MAX_INJECTION_CYCLES:
injections = await self._drain_injections(spec)
if injections:
had_injections = True
injection_cycles += 1
_injected_after_final = True
for text in injections:
messages.append({"role": "user", "content": text})
logger.info(
"Injected {} follow-up message(s) after final response ({}/{})",
len(injections), injection_cycles, _MAX_INJECTION_CYCLES,
)
if hook.wants_streaming(): if hook.wants_streaming():
await hook.on_stream_end(context, resuming=False) await hook.on_stream_end(context, resuming=_injected_after_final)
if _injected_after_final:
await hook.after_iteration(context)
continue
if response.finish_reason == "error": if response.finish_reason == "error":
final_content = clean or spec.error_message or _DEFAULT_ERROR_MESSAGE final_content = clean or spec.error_message or _DEFAULT_ERROR_MESSAGE
@@ -283,6 +356,7 @@ class AgentRunner:
stop_reason=stop_reason, stop_reason=stop_reason,
error=error, error=error,
tool_events=tool_events, tool_events=tool_events,
had_injections=had_injections,
) )
def _build_request_kwargs( def _build_request_kwargs(
+59 -20
View File
@@ -15,6 +15,8 @@ from nanobot.agent.tools.sandbox import wrap_command
from nanobot.agent.tools.schema import IntegerSchema, StringSchema, tool_parameters_schema from nanobot.agent.tools.schema import IntegerSchema, StringSchema, tool_parameters_schema
from nanobot.config.paths import get_media_dir from nanobot.config.paths import get_media_dir
_IS_WINDOWS = sys.platform == "win32"
@tool_parameters( @tool_parameters(
tool_parameters_schema( tool_parameters_schema(
@@ -88,27 +90,27 @@ class ExecTool(Tool):
return guard_error return guard_error
if self.sandbox: if self.sandbox:
workspace = self.working_dir or cwd if _IS_WINDOWS:
command = wrap_command(self.sandbox, command, workspace, cwd) logger.warning(
cwd = str(Path(workspace).resolve()) "Sandbox '{}' is not supported on Windows; running unsandboxed",
self.sandbox,
)
else:
workspace = self.working_dir or cwd
command = wrap_command(self.sandbox, command, workspace, cwd)
cwd = str(Path(workspace).resolve())
effective_timeout = min(timeout or self.timeout, self._MAX_TIMEOUT) effective_timeout = min(timeout or self.timeout, self._MAX_TIMEOUT)
env = self._build_env() env = self._build_env()
if self.path_append: if self.path_append:
command = f'export PATH="$PATH:{self.path_append}"; {command}' if _IS_WINDOWS:
env["PATH"] = env.get("PATH", "") + ";" + self.path_append
bash = shutil.which("bash") or "/bin/bash" else:
command = f'export PATH="$PATH:{self.path_append}"; {command}'
try: try:
process = await asyncio.create_subprocess_exec( process = await self._spawn(command, cwd, env)
bash, "-l", "-c", command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=cwd,
env=env,
)
try: try:
stdout, stderr = await asyncio.wait_for( stdout, stderr = await asyncio.wait_for(
@@ -136,7 +138,6 @@ class ExecTool(Tool):
result = "\n".join(output_parts) if output_parts else "(no output)" result = "\n".join(output_parts) if output_parts else "(no output)"
# Head + tail truncation to preserve both start and end of output
max_len = self._MAX_OUTPUT max_len = self._MAX_OUTPUT
if len(result) > max_len: if len(result) > max_len:
half = max_len // 2 half = max_len // 2
@@ -151,6 +152,29 @@ class ExecTool(Tool):
except Exception as e: except Exception as e:
return f"Error executing command: {str(e)}" return f"Error executing command: {str(e)}"
@staticmethod
async def _spawn(
command: str, cwd: str, env: dict[str, str],
) -> asyncio.subprocess.Process:
"""Launch *command* in a platform-appropriate shell."""
if _IS_WINDOWS:
comspec = env.get("COMSPEC", os.environ.get("COMSPEC", "cmd.exe"))
return await asyncio.create_subprocess_exec(
comspec, "/c", command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=cwd,
env=env,
)
bash = shutil.which("bash") or "/bin/bash"
return await asyncio.create_subprocess_exec(
bash, "-l", "-c", command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=cwd,
env=env,
)
@staticmethod @staticmethod
async def _kill_process(process: asyncio.subprocess.Process) -> None: async def _kill_process(process: asyncio.subprocess.Process) -> None:
"""Kill a subprocess and reap it to prevent zombies.""" """Kill a subprocess and reap it to prevent zombies."""
@@ -160,7 +184,7 @@ class ExecTool(Tool):
except asyncio.TimeoutError: except asyncio.TimeoutError:
pass pass
finally: finally:
if sys.platform != "win32": if not _IS_WINDOWS:
try: try:
os.waitpid(process.pid, os.WNOHANG) os.waitpid(process.pid, os.WNOHANG)
except (ProcessLookupError, ChildProcessError) as e: except (ProcessLookupError, ChildProcessError) as e:
@@ -169,11 +193,26 @@ class ExecTool(Tool):
def _build_env(self) -> dict[str, str]: def _build_env(self) -> dict[str, str]:
"""Build a minimal environment for subprocess execution. """Build a minimal environment for subprocess execution.
Uses HOME so that ``bash -l`` sources the user's profile (which sets On Unix, only HOME/LANG/TERM are passed; ``bash -l`` sources the
PATH and other essentials). Only PATH is extended with *path_append*; user's profile which sets PATH and other essentials.
the parent process's environment is **not** inherited, preventing
secrets in env vars from leaking to LLM-generated commands. On Windows, ``cmd.exe`` has no login-profile mechanism, so a curated
set of system variables (including PATH) is forwarded. API keys and
other secrets are still excluded.
""" """
if _IS_WINDOWS:
sr = os.environ.get("SYSTEMROOT", r"C:\Windows")
return {
"SYSTEMROOT": sr,
"COMSPEC": os.environ.get("COMSPEC", f"{sr}\\system32\\cmd.exe"),
"USERPROFILE": os.environ.get("USERPROFILE", ""),
"HOMEDRIVE": os.environ.get("HOMEDRIVE", "C:"),
"HOMEPATH": os.environ.get("HOMEPATH", "\\"),
"TEMP": os.environ.get("TEMP", f"{sr}\\Temp"),
"TMP": os.environ.get("TMP", f"{sr}\\Temp"),
"PATHEXT": os.environ.get("PATHEXT", ".COM;.EXE;.BAT;.CMD"),
"PATH": os.environ.get("PATH", f"{sr}\\system32;{sr}"),
}
home = os.environ.get("HOME", "/tmp") home = os.environ.get("HOME", "/tmp")
return { return {
"HOME": home, "HOME": home,
+76 -16
View File
@@ -22,6 +22,8 @@ from nanobot.channels.base import BaseChannel
from nanobot.config.paths import get_media_dir from nanobot.config.paths import get_media_dir
from nanobot.config.schema import Base from nanobot.config.schema import Base
from lark_oapi.core.const import FEISHU_DOMAIN, LARK_DOMAIN
FEISHU_AVAILABLE = importlib.util.find_spec("lark_oapi") is not None FEISHU_AVAILABLE = importlib.util.find_spec("lark_oapi") is not None
# Message type display mapping # Message type display mapping
@@ -250,9 +252,12 @@ class FeishuConfig(Base):
verification_token: str = "" verification_token: str = ""
allow_from: list[str] = Field(default_factory=list) allow_from: list[str] = Field(default_factory=list)
react_emoji: str = "THUMBSUP" react_emoji: str = "THUMBSUP"
done_emoji: str | None = None # Emoji to show when task is completed (e.g., "DONE", "OK")
tool_hint_prefix: str = "\U0001f527" # Prefix for inline tool hints (default: 🔧)
group_policy: Literal["open", "mention"] = "mention" group_policy: Literal["open", "mention"] = "mention"
reply_to_message: bool = False # If True, bot replies quote the user's original message reply_to_message: bool = False # If True, bot replies quote the user's original message
streaming: bool = True streaming: bool = True
domain: Literal["feishu", "lark"] = "feishu" # Set to "lark" for international Lark
_STREAM_ELEMENT_ID = "streaming_md" _STREAM_ELEMENT_ID = "streaming_md"
@@ -326,10 +331,12 @@ class FeishuChannel(BaseChannel):
self._loop = asyncio.get_running_loop() self._loop = asyncio.get_running_loop()
# Create Lark client for sending messages # Create Lark client for sending messages
domain = LARK_DOMAIN if self.config.domain == "lark" else FEISHU_DOMAIN
self._client = ( self._client = (
lark.Client.builder() lark.Client.builder()
.app_id(self.config.app_id) .app_id(self.config.app_id)
.app_secret(self.config.app_secret) .app_secret(self.config.app_secret)
.domain(domain)
.log_level(lark.LogLevel.INFO) .log_level(lark.LogLevel.INFO)
.build() .build()
) )
@@ -357,6 +364,7 @@ class FeishuChannel(BaseChannel):
self._ws_client = lark.ws.Client( self._ws_client = lark.ws.Client(
self.config.app_id, self.config.app_id,
self.config.app_secret, self.config.app_secret,
domain=domain,
event_handler=event_handler, event_handler=event_handler,
log_level=lark.LogLevel.INFO, log_level=lark.LogLevel.INFO,
) )
@@ -1012,14 +1020,29 @@ class FeishuChannel(BaseChannel):
elif msg_type in ("audio", "file", "media"): elif msg_type in ("audio", "file", "media"):
file_key = content_json.get("file_key") file_key = content_json.get("file_key")
if file_key and message_id: if not file_key:
data, filename = await loop.run_in_executor( logger.warning("Feishu {} message missing file_key: {}", msg_type, content_json)
None, self._download_file_sync, message_id, file_key, msg_type return None, f"[{msg_type}: missing file_key]"
) if not message_id:
if not filename: logger.warning("Feishu {} message missing message_id", msg_type)
filename = file_key[:16] return None, f"[{msg_type}: missing message_id]"
if msg_type == "audio" and not filename.endswith(".opus"):
filename = f"{filename}.opus" data, filename = await loop.run_in_executor(
None, self._download_file_sync, message_id, file_key, msg_type
)
if not data:
logger.warning("Feishu {} download failed: file_key={}", msg_type, file_key)
return None, f"[{msg_type}: download failed]"
if not filename:
filename = file_key[:16]
# Feishu voice messages are opus in OGG container.
# Use .ogg extension for better Whisper compatibility.
if msg_type == "audio":
if not any(filename.endswith(ext) for ext in (".opus", ".ogg", ".oga")):
filename = f"{filename}.ogg"
if data and filename: if data and filename:
file_path = media_dir / filename file_path = media_dir / filename
@@ -1263,7 +1286,15 @@ class FeishuChannel(BaseChannel):
async def send_delta( async def send_delta(
self, chat_id: str, delta: str, metadata: dict[str, Any] | None = None self, chat_id: str, delta: str, metadata: dict[str, Any] | None = None
) -> None: ) -> None:
"""Progressive streaming via CardKit: create card on first delta, stream-update on subsequent.""" """Progressive streaming via CardKit: create card on first delta, stream-update on subsequent.
Supported metadata keys:
_stream_end: Finalize the streaming card.
_resuming: Mid-turn pause flush but keep the buffer alive.
_tool_hint: Delta is a formatted tool hint (for display only).
message_id: Original message id (used with _stream_end for reaction cleanup).
reaction_id: Reaction id to remove on stream end.
"""
if not self._client: if not self._client:
return return
meta = metadata or {} meta = metadata or {}
@@ -1274,6 +1305,22 @@ class FeishuChannel(BaseChannel):
if meta.get("_stream_end"): if meta.get("_stream_end"):
if (message_id := meta.get("message_id")) and (reaction_id := meta.get("reaction_id")): if (message_id := meta.get("message_id")) and (reaction_id := meta.get("reaction_id")):
await self._remove_reaction(message_id, reaction_id) await self._remove_reaction(message_id, reaction_id)
# Add completion emoji if configured
if self.config.done_emoji and message_id:
await self._add_reaction(message_id, self.config.done_emoji)
resuming = meta.get("_resuming", False)
if resuming:
# Mid-turn pause (e.g. tool call between streaming segments).
# Flush current text to card but keep the buffer alive so the
# next segment appends to the same card.
buf = self._stream_bufs.get(chat_id)
if buf and buf.card_id and buf.text:
buf.sequence += 1
await loop.run_in_executor(
None, self._stream_update_text_sync, buf.card_id, buf.text, buf.sequence,
)
return
buf = self._stream_bufs.pop(chat_id, None) buf = self._stream_bufs.pop(chat_id, None)
if not buf or not buf.text: if not buf or not buf.text:
@@ -1346,13 +1393,26 @@ class FeishuChannel(BaseChannel):
receive_id_type = "chat_id" if msg.chat_id.startswith("oc_") else "open_id" receive_id_type = "chat_id" if msg.chat_id.startswith("oc_") else "open_id"
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
# Handle tool hint messages as code blocks in interactive cards. # Handle tool hint messages. When a streaming card is active for
# These are progress-only messages and should bypass normal reply routing. # this chat, inline the hint into the card instead of sending a
# separate message so the user experience stays cohesive.
if msg.metadata.get("_tool_hint"): if msg.metadata.get("_tool_hint"):
if msg.content and msg.content.strip(): hint = (msg.content or "").strip()
await self._send_tool_hint_card( if not hint:
receive_id_type, msg.chat_id, msg.content.strip() return
) buf = self._stream_bufs.get(msg.chat_id)
if buf and buf.card_id:
# Delegate to send_delta so tool hints get the same
# throttling (and card creation) as regular text deltas.
lines = self.__class__._format_tool_hint_lines(hint).split("\n")
delta = "\n\n" + "\n".join(
f"{self.config.tool_hint_prefix} {ln}" for ln in lines if ln.strip()
) + "\n\n"
await self.send_delta(msg.chat_id, delta)
return
await self._send_tool_hint_card(
receive_id_type, msg.chat_id, hint
)
return return
# Determine whether the first message should quote the user's message. # Determine whether the first message should quote the user's message.
@@ -1661,7 +1721,7 @@ class FeishuChannel(BaseChannel):
loop = asyncio.get_running_loop() loop = asyncio.get_running_loop()
# Put each top-level tool call on its own line without altering commas inside arguments. # Put each top-level tool call on its own line without altering commas inside arguments.
formatted_code = self._format_tool_hint_lines(tool_hint) formatted_code = self.__class__._format_tool_hint_lines(tool_hint)
card = { card = {
"config": {"wide_screen_mode": True}, "config": {"wide_screen_mode": True},
+522
View File
@@ -0,0 +1,522 @@
"""Microsoft Teams channel MVP using a tiny built-in HTTP webhook server.
Scope:
- DM-focused MVP
- text inbound/outbound
- conversation reference persistence
- sender allowlist support
- optional inbound Bot Framework bearer-token validation
- no attachments/cards/polls yet
"""
from __future__ import annotations
import asyncio
import html
import importlib.util
import json
import re
import threading
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import TYPE_CHECKING, Any
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_workspace_path
from nanobot.config.schema import Base
MSTEAMS_AVAILABLE = importlib.util.find_spec("jwt") is not None
if TYPE_CHECKING:
import jwt
if MSTEAMS_AVAILABLE:
import jwt
class MSTeamsConfig(Base):
"""Microsoft Teams channel configuration."""
enabled: bool = False
app_id: str = ""
app_password: str = ""
tenant_id: str = ""
host: str = "0.0.0.0"
port: int = 3978
path: str = "/api/messages"
allow_from: list[str] = Field(default_factory=list)
reply_in_thread: bool = True
mention_only_response: str = "Hi — what can I help with?"
validate_inbound_auth: bool = False
restart_notify_enabled: bool = False
restart_notify_pre_message: str = (
"Nanobot agent initiated a gateway restart. I will message again when the gateway is back online."
)
restart_notify_post_message: str = "Nanobot gateway is back online."
@dataclass
class ConversationRef:
"""Minimal stored conversation reference for replies."""
service_url: str
conversation_id: str
bot_id: str | None = None
activity_id: str | None = None
conversation_type: str | None = None
tenant_id: str | None = None
class MSTeamsChannel(BaseChannel):
"""Microsoft Teams channel (DM-first MVP)."""
name = "msteams"
display_name = "Microsoft Teams"
@classmethod
def default_config(cls) -> dict[str, Any]:
return MSTeamsConfig().model_dump(by_alias=True)
def __init__(self, config: Any, bus: MessageBus):
if isinstance(config, dict):
config = MSTeamsConfig.model_validate(config)
super().__init__(config, bus)
self.config: MSTeamsConfig = config
self._loop: asyncio.AbstractEventLoop | None = None
self._server: ThreadingHTTPServer | None = None
self._server_thread: threading.Thread | None = None
self._http: httpx.AsyncClient | None = None
self._token: str | None = None
self._token_expires_at: float = 0.0
self._botframework_openid_config_url = (
"https://login.botframework.com/v1/.well-known/openidconfiguration"
)
self._botframework_openid_config: dict[str, Any] | None = None
self._botframework_openid_config_expires_at: float = 0.0
self._botframework_jwks: dict[str, Any] | None = None
self._botframework_jwks_expires_at: float = 0.0
self._refs_path = get_workspace_path() / "state" / "msteams_conversations.json"
self._refs_path.parent.mkdir(parents=True, exist_ok=True)
self._conversation_refs: dict[str, ConversationRef] = self._load_refs()
async def start(self) -> None:
"""Start the Teams webhook listener."""
if not MSTEAMS_AVAILABLE:
logger.error("PyJWT not installed. Run: pip install nanobot-ai[msteams]")
return
if not self.config.app_id or not self.config.app_password:
logger.error("MSTeams app_id/app_password not configured")
return
self._loop = asyncio.get_running_loop()
self._http = httpx.AsyncClient(timeout=30.0)
self._running = True
channel = self
class Handler(BaseHTTPRequestHandler):
def do_POST(self) -> None:
if self.path != channel.config.path:
self.send_response(404)
self.end_headers()
return
try:
length = int(self.headers.get("Content-Length", "0"))
raw = self.rfile.read(length) if length > 0 else b"{}"
payload = json.loads(raw.decode("utf-8"))
except Exception as e:
logger.warning("MSTeams invalid request body: {}", e)
self.send_response(400)
self.end_headers()
return
auth_header = self.headers.get("Authorization", "")
if channel.config.validate_inbound_auth:
try:
fut = asyncio.run_coroutine_threadsafe(
channel._validate_inbound_auth(auth_header, payload),
channel._loop,
)
fut.result(timeout=15)
except Exception as e:
logger.warning("MSTeams inbound auth validation failed: {}", e)
self.send_response(401)
self.send_header("Content-Type", "application/json")
self.end_headers()
self.wfile.write(b'{"error":"unauthorized"}')
return
try:
fut = asyncio.run_coroutine_threadsafe(
channel._handle_activity(payload),
channel._loop,
)
fut.result(timeout=15)
except Exception as e:
logger.warning("MSTeams activity handling failed: {}", e)
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.end_headers()
self.wfile.write(b"{}")
def log_message(self, format: str, *args: Any) -> None:
return
self._server = ThreadingHTTPServer((self.config.host, self.config.port), Handler)
self._server_thread = threading.Thread(
target=self._server.serve_forever,
name="nanobot-msteams",
daemon=True,
)
self._server_thread.start()
logger.info(
"MSTeams webhook listening on http://{}:{}{}",
self.config.host,
self.config.port,
self.config.path,
)
while self._running:
await asyncio.sleep(1)
async def stop(self) -> None:
"""Stop the channel."""
self._running = False
if self._server:
self._server.shutdown()
self._server.server_close()
self._server = None
if self._server_thread and self._server_thread.is_alive():
self._server_thread.join(timeout=2)
self._server_thread = None
if self._http:
await self._http.aclose()
self._http = None
async def send(self, msg: OutboundMessage) -> None:
"""Send a plain text reply into an existing Teams conversation."""
if not self._http:
raise RuntimeError("MSTeams HTTP client not initialized")
ref = self._conversation_refs.get(str(msg.chat_id))
if not ref:
raise RuntimeError(f"MSTeams conversation ref not found for chat_id={msg.chat_id}")
token = await self._get_access_token()
base_url = f"{ref.service_url.rstrip('/')}/v3/conversations/{ref.conversation_id}/activities"
use_thread_reply = self.config.reply_in_thread and bool(ref.activity_id)
url = f"{base_url}/{ref.activity_id}" if use_thread_reply else base_url
headers = {
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
}
payload = {
"type": "message",
"text": msg.content or " ",
}
if use_thread_reply:
payload["replyToId"] = ref.activity_id
try:
resp = await self._http.post(url, headers=headers, json=payload)
resp.raise_for_status()
logger.info("MSTeams message sent to {}", ref.conversation_id)
except Exception as e:
logger.error("MSTeams send failed: {}", e)
raise
async def _handle_activity(self, activity: dict[str, Any]) -> None:
"""Handle inbound Teams/Bot Framework activity."""
if activity.get("type") != "message":
return
conversation = activity.get("conversation") or {}
from_user = activity.get("from") or {}
recipient = activity.get("recipient") or {}
channel_data = activity.get("channelData") or {}
sender_id = str(from_user.get("aadObjectId") or from_user.get("id") or "").strip()
conversation_id = str(conversation.get("id") or "").strip()
service_url = str(activity.get("serviceUrl") or "").strip()
activity_id = str(activity.get("id") or "").strip()
conversation_type = str(conversation.get("conversationType") or "").strip()
if not sender_id or not conversation_id or not service_url:
return
if recipient.get("id") and from_user.get("id") == recipient.get("id"):
return
# DM-only MVP: ignore group/channel traffic for now
if conversation_type and conversation_type not in ("personal", ""):
logger.debug("MSTeams ignoring non-DM conversation {}", conversation_type)
return
text = self._sanitize_inbound_text(activity)
if not text:
text = self.config.mention_only_response.strip()
if not text:
logger.debug("MSTeams ignoring empty message after Teams text sanitization")
return
self._conversation_refs[conversation_id] = ConversationRef(
service_url=service_url,
conversation_id=conversation_id,
bot_id=str(recipient.get("id") or "") or None,
activity_id=activity_id or None,
conversation_type=conversation_type or None,
tenant_id=str((channel_data.get("tenant") or {}).get("id") or "") or None,
)
self._save_refs()
await self._handle_message(
sender_id=sender_id,
chat_id=conversation_id,
content=text,
metadata={
"msteams": {
"activity_id": activity_id,
"conversation_id": conversation_id,
"conversation_type": conversation_type or "personal",
"from_name": from_user.get("name"),
}
},
)
def _sanitize_inbound_text(self, activity: dict[str, Any]) -> str:
"""Extract the user-authored text from a Teams activity."""
text = str(activity.get("text") or "")
text = self._strip_possible_bot_mention(text)
channel_data = activity.get("channelData") or {}
reply_to_id = str(activity.get("replyToId") or "").strip()
normalized_preview = html.unescape(text).replace("&rsquo", "").strip()
normalized_preview = normalized_preview.replace("\r\n", "\n").replace("\r", "\n")
preview_lines = [line.strip() for line in normalized_preview.split("\n")]
while preview_lines and not preview_lines[0]:
preview_lines.pop(0)
first_line = preview_lines[0] if preview_lines else ""
looks_like_quote_wrapper = first_line.lower().startswith("replying to ") or first_line.startswith("FWDIOC-BOT")
if reply_to_id or channel_data.get("messageType") == "reply" or looks_like_quote_wrapper:
text = self._normalize_teams_reply_quote(text)
return text.strip()
def _strip_possible_bot_mention(self, text: str) -> str:
"""Remove simple Teams mention markup from message text."""
cleaned = re.sub(r"<at\b[^>]*>.*?</at>", " ", text, flags=re.IGNORECASE | re.DOTALL)
cleaned = re.sub(r"[^\S\r\n]+", " ", cleaned)
cleaned = re.sub(r"(?:\r?\n){3,}", "\n\n", cleaned)
return cleaned.strip()
def _normalize_teams_reply_quote(self, text: str) -> str:
"""Normalize Teams quoted replies into a compact structured form."""
cleaned = html.unescape(text).replace("&rsquo", "").strip()
if not cleaned:
return ""
normalized_newlines = cleaned.replace("\r\n", "\n").replace("\r", "\n")
lines = [line.strip() for line in normalized_newlines.split("\n")]
while lines and not lines[0]:
lines.pop(0)
# Observed native Teams reply wrapper:
# Replying to Bob Smith
# actual reply text
if len(lines) >= 2 and lines[0].lower().startswith("replying to "):
quoted = lines[0][len("replying to ") :].strip(" :")
reply = "\n".join(lines[1:]).strip()
return self._format_reply_with_quote(quoted, reply)
# Observed FWDIOC relay wrapper where the quoted content is surfaced after a
# synthetic "FWDIOC-BOT" header, sometimes with a blank line separating quote
# and reply, and sometimes as a compact line-based fallback shape.
if lines and lines[0].strip().startswith("FWDIOC-BOT"):
body = normalized_newlines.split("\n", 1)[1] if "\n" in normalized_newlines else ""
body = body.lstrip()
parts = re.split(r"\n\s*\n", body, maxsplit=1)
if len(parts) == 2:
quoted = re.sub(r"\s+", " ", parts[0]).strip()
reply = re.sub(r"\s+", " ", parts[1]).strip()
if quoted or reply:
return self._format_reply_with_quote(quoted, reply)
body_lines = [line.strip() for line in body.split("\n") if line.strip()]
if body_lines:
quoted = " ".join(body_lines[:-1]).strip()
reply = body_lines[-1].strip()
if quoted and reply:
return self._format_reply_with_quote(quoted, reply)
# Observed compact fallback where the relay flattens everything into one line
# and appends the literal reply text marker at the end.
compact = re.sub(r"\s+", " ", normalized_newlines).strip()
if compact.startswith("FWDIOC-BOT "):
compact = compact[len("FWDIOC-BOT ") :].strip()
marker = " Reply with quote test"
if compact.endswith(marker):
quoted = compact[: -len(marker)].strip()
reply = marker.strip()
return self._format_reply_with_quote(quoted, reply)
return cleaned
def _format_reply_with_quote(self, quoted: str, reply: str) -> str:
"""Format a quoted reply for the model without Teams wrapper noise."""
quoted = quoted.strip()
reply = reply.strip()
if quoted and reply:
return f"User is replying to: {quoted}\nUser reply: {reply}"
if reply:
return reply
return quoted
async def _validate_inbound_auth(self, auth_header: str, activity: dict[str, Any]) -> None:
"""Validate inbound Bot Framework bearer token."""
if not MSTEAMS_AVAILABLE:
raise RuntimeError("PyJWT not installed. Run: pip install nanobot-ai[msteams]")
if not auth_header.lower().startswith("bearer "):
raise ValueError("missing bearer token")
token = auth_header.split(" ", 1)[1].strip()
if not token:
raise ValueError("empty bearer token")
header = jwt.get_unverified_header(token)
kid = str(header.get("kid") or "").strip()
if not kid:
raise ValueError("missing token kid")
jwks = await self._get_botframework_jwks()
keys = jwks.get("keys") or []
jwk = next((key for key in keys if key.get("kid") == kid), None)
if not jwk:
raise ValueError(f"signing key not found for kid={kid}")
public_key = jwt.algorithms.RSAAlgorithm.from_jwk(json.dumps(jwk))
claims = jwt.decode(
token,
key=public_key,
algorithms=["RS256"],
audience=self.config.app_id,
issuer="https://api.botframework.com",
options={
"require": ["exp", "nbf", "iss", "aud"],
},
)
claim_service_url = str(
claims.get("serviceurl") or claims.get("serviceUrl") or "",
).strip()
activity_service_url = str(activity.get("serviceUrl") or "").strip()
if claim_service_url and activity_service_url and claim_service_url != activity_service_url:
raise ValueError("serviceUrl claim mismatch")
async def _get_botframework_openid_config(self) -> dict[str, Any]:
"""Fetch and cache Bot Framework OpenID configuration."""
import time
now = time.time()
if self._botframework_openid_config and now < self._botframework_openid_config_expires_at:
return self._botframework_openid_config
if not self._http:
raise RuntimeError("MSTeams HTTP client not initialized")
resp = await self._http.get(self._botframework_openid_config_url)
resp.raise_for_status()
self._botframework_openid_config = resp.json()
self._botframework_openid_config_expires_at = now + 3600
return self._botframework_openid_config
async def _get_botframework_jwks(self) -> dict[str, Any]:
"""Fetch and cache Bot Framework JWKS."""
import time
now = time.time()
if self._botframework_jwks and now < self._botframework_jwks_expires_at:
return self._botframework_jwks
if not self._http:
raise RuntimeError("MSTeams HTTP client not initialized")
openid_config = await self._get_botframework_openid_config()
jwks_uri = str(openid_config.get("jwks_uri") or "").strip()
if not jwks_uri:
raise RuntimeError("Bot Framework OpenID config missing jwks_uri")
resp = await self._http.get(jwks_uri)
resp.raise_for_status()
self._botframework_jwks = resp.json()
self._botframework_jwks_expires_at = now + 3600
return self._botframework_jwks
def _load_refs(self) -> dict[str, ConversationRef]:
"""Load stored conversation references."""
if not self._refs_path.exists():
return {}
try:
data = json.loads(self._refs_path.read_text(encoding="utf-8"))
out: dict[str, ConversationRef] = {}
for key, value in data.items():
out[key] = ConversationRef(**value)
return out
except Exception as e:
logger.warning("Failed to load MSTeams conversation refs: {}", e)
return {}
def _save_refs(self) -> None:
"""Persist conversation references."""
try:
data = {
key: {
"service_url": ref.service_url,
"conversation_id": ref.conversation_id,
"bot_id": ref.bot_id,
"activity_id": ref.activity_id,
"conversation_type": ref.conversation_type,
"tenant_id": ref.tenant_id,
}
for key, ref in self._conversation_refs.items()
}
self._refs_path.write_text(json.dumps(data, indent=2), encoding="utf-8")
except Exception as e:
logger.warning("Failed to save MSTeams conversation refs: {}", e)
async def _get_access_token(self) -> str:
"""Fetch an access token for Bot Framework / Azure Bot auth."""
import time
now = time.time()
if self._token and now < self._token_expires_at - 60:
return self._token
if not self._http:
raise RuntimeError("MSTeams HTTP client not initialized")
tenant = (self.config.tenant_id or "").strip() or "botframework.com"
token_url = f"https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
data = {
"grant_type": "client_credentials",
"client_id": self.config.app_id,
"client_secret": self.config.app_password,
"scope": "https://api.botframework.com/.default",
}
resp = await self._http.post(token_url, data=data)
resp.raise_for_status()
payload = resp.json()
self._token = payload["access_token"]
self._token_expires_at = now + int(payload.get("expires_in", 3600))
return self._token
+112 -82
View File
@@ -242,43 +242,46 @@ class QQChannel(BaseChannel):
async def send(self, msg: OutboundMessage) -> None: async def send(self, msg: OutboundMessage) -> None:
"""Send attachments first, then text.""" """Send attachments first, then text."""
if not self._client: try:
logger.warning("QQ client not initialized") if not self._client:
return logger.warning("QQ client not initialized")
return
msg_id = msg.metadata.get("message_id") msg_id = msg.metadata.get("message_id")
chat_type = self._chat_type_cache.get(msg.chat_id, "c2c") chat_type = self._chat_type_cache.get(msg.chat_id, "c2c")
is_group = chat_type == "group" is_group = chat_type == "group"
# 1) Send media # 1) Send media
for media_ref in msg.media or []: for media_ref in msg.media or []:
ok = await self._send_media( ok = await self._send_media(
chat_id=msg.chat_id, chat_id=msg.chat_id,
media_ref=media_ref, media_ref=media_ref,
msg_id=msg_id, msg_id=msg_id,
is_group=is_group, is_group=is_group,
)
if not ok:
filename = (
os.path.basename(urlparse(media_ref).path)
or os.path.basename(media_ref)
or "file"
) )
if not ok:
filename = (
os.path.basename(urlparse(media_ref).path)
or os.path.basename(media_ref)
or "file"
)
await self._send_text_only(
chat_id=msg.chat_id,
is_group=is_group,
msg_id=msg_id,
content=f"[Attachment send failed: {filename}]",
)
# 2) Send text
if msg.content and msg.content.strip():
await self._send_text_only( await self._send_text_only(
chat_id=msg.chat_id, chat_id=msg.chat_id,
is_group=is_group, is_group=is_group,
msg_id=msg_id, msg_id=msg_id,
content=f"[Attachment send failed: {filename}]", content=msg.content.strip(),
) )
except Exception:
# 2) Send text logger.exception("Error sending QQ message to chat_id={}", msg.chat_id)
if msg.content and msg.content.strip():
await self._send_text_only(
chat_id=msg.chat_id,
is_group=is_group,
msg_id=msg_id,
content=msg.content.strip(),
)
async def _send_text_only( async def _send_text_only(
self, self,
@@ -438,15 +441,26 @@ class QQChannel(BaseChannel):
endpoint = "/v2/users/{openid}/files" endpoint = "/v2/users/{openid}/files"
id_key = "openid" id_key = "openid"
payload = { payload: dict[str, Any] = {
id_key: chat_id, id_key: chat_id,
"file_type": file_type, "file_type": file_type,
"file_data": file_data, "file_data": file_data,
"file_name": file_name,
"srv_send_msg": srv_send_msg, "srv_send_msg": srv_send_msg,
} }
# Only pass file_name for non-image types (file_type=4).
# Passing file_name for images causes QQ client to render them as
# file attachments instead of inline images.
if file_type != QQ_FILE_TYPE_IMAGE and file_name:
payload["file_name"] = file_name
route = Route("POST", endpoint, **{id_key: chat_id}) route = Route("POST", endpoint, **{id_key: chat_id})
return await self._client.api._http.request(route, json=payload) result = await self._client.api._http.request(route, json=payload)
# Extract only the file_info field to avoid extra fields (file_uuid, ttl, etc.)
# that may confuse QQ client when sending the media object.
if isinstance(result, dict) and "file_info" in result:
return {"file_info": result["file_info"]}
return result
# --------------------------- # ---------------------------
# Inbound (receive) # Inbound (receive)
@@ -454,58 +468,68 @@ class QQChannel(BaseChannel):
async def _on_message(self, data: C2CMessage | GroupMessage, is_group: bool = False) -> None: async def _on_message(self, data: C2CMessage | GroupMessage, is_group: bool = False) -> None:
"""Parse inbound message, download attachments, and publish to the bus.""" """Parse inbound message, download attachments, and publish to the bus."""
if data.id in self._processed_ids: try:
return if data.id in self._processed_ids:
self._processed_ids.append(data.id) return
self._processed_ids.append(data.id)
if is_group: if is_group:
chat_id = data.group_openid chat_id = data.group_openid
user_id = data.author.member_openid user_id = data.author.member_openid
self._chat_type_cache[chat_id] = "group" self._chat_type_cache[chat_id] = "group"
else: else:
chat_id = str( chat_id = str(
getattr(data.author, "id", None) or getattr(data.author, "user_openid", "unknown") getattr(data.author, "id", None)
) or getattr(data.author, "user_openid", "unknown")
user_id = chat_id
self._chat_type_cache[chat_id] = "c2c"
content = (data.content or "").strip()
# the data used by tests don't contain attachments property
# so we use getattr with a default of [] to avoid AttributeError in tests
attachments = getattr(data, "attachments", None) or []
media_paths, recv_lines, att_meta = await self._handle_attachments(attachments)
# Compose content that always contains actionable saved paths
if recv_lines:
tag = "[Image]" if any(_is_image_name(Path(p).name) for p in media_paths) else "[File]"
file_block = "Received files:\n" + "\n".join(recv_lines)
content = f"{content}\n\n{file_block}".strip() if content else f"{tag}\n{file_block}"
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: user_id = chat_id
logger.debug("QQ ack message failed for chat_id={}", chat_id) self._chat_type_cache[chat_id] = "c2c"
await self._handle_message( content = (data.content or "").strip()
sender_id=user_id,
chat_id=chat_id, # the data used by tests don't contain attachments property
content=content, # so we use getattr with a default of [] to avoid AttributeError in tests
media=media_paths if media_paths else None, attachments = getattr(data, "attachments", None) or []
metadata={ media_paths, recv_lines, att_meta = await self._handle_attachments(attachments)
"message_id": data.id,
"attachments": att_meta, # Compose content that always contains actionable saved paths
}, if recv_lines:
) tag = (
"[Image]"
if any(_is_image_name(Path(p).name) for p in media_paths)
else "[File]"
)
file_block = "Received files:\n" + "\n".join(recv_lines)
content = (
f"{content}\n\n{file_block}".strip() if content else f"{tag}\n{file_block}"
)
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,
content=content,
media=media_paths if media_paths else None,
metadata={
"message_id": data.id,
"attachments": att_meta,
},
)
except Exception:
logger.exception("Error handling QQ inbound message id={}", getattr(data, "id", "?"))
async def _handle_attachments( async def _handle_attachments(
self, self,
@@ -520,7 +544,9 @@ class QQChannel(BaseChannel):
return media_paths, recv_lines, att_meta return media_paths, recv_lines, att_meta
for att in attachments: for att in attachments:
url, filename, ctype = att.url, att.filename, att.content_type url = getattr(att, "url", None) or ""
filename = getattr(att, "filename", None) or ""
ctype = getattr(att, "content_type", None) or ""
logger.info("Downloading file from QQ: {}", filename or url) logger.info("Downloading file from QQ: {}", filename or url)
local_path = await self._download_to_media_dir_chunked(url, filename_hint=filename) local_path = await self._download_to_media_dir_chunked(url, filename_hint=filename)
@@ -555,6 +581,10 @@ class QQChannel(BaseChannel):
Enforces a max download size and writes to a .part temp file Enforces a max download size and writes to a .part temp file
that is atomically renamed on success. that is atomically renamed on success.
""" """
# Handle protocol-relative URLs (e.g. "//multimedia.nt.qq.com/...")
if url.startswith("//"):
url = f"https:{url}"
if not self._http: if not self._http:
self._http = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) self._http = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120))
+457
View File
@@ -0,0 +1,457 @@
"""WebSocket server channel: nanobot acts as a WebSocket server and serves connected clients."""
from __future__ import annotations
import asyncio
import email.utils
import hmac
import http
import json
import secrets
import ssl
import time
import uuid
from typing import Any, Self
from urllib.parse import parse_qs, urlparse
from loguru import logger
from pydantic import Field, field_validator, model_validator
from websockets.asyncio.server import ServerConnection, serve
from websockets.datastructures import Headers
from websockets.exceptions import ConnectionClosed
from websockets.http11 import Request as WsRequest, Response
from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel
from nanobot.config.schema import Base
def _strip_trailing_slash(path: str) -> str:
if len(path) > 1 and path.endswith("/"):
return path.rstrip("/")
return path or "/"
def _normalize_config_path(path: str) -> str:
return _strip_trailing_slash(path)
class WebSocketConfig(Base):
"""WebSocket server channel configuration.
Clients connect with URLs like ``ws://{host}:{port}{path}?client_id=...&token=...``.
- ``client_id``: Used for ``allow_from`` authorization; if omitted, a value is generated and logged.
- ``token``: If non-empty, the ``token`` query param may match this static secret; short-lived tokens
from ``token_issue_path`` are also accepted.
- ``token_issue_path``: If non-empty, **GET** (HTTP/1.1) to this path returns JSON
``{"token": "...", "expires_in": <seconds>}``; use ``?token=...`` when opening the WebSocket.
Must differ from ``path`` (the WS upgrade path). If the client runs in the **same process** as
nanobot and shares the asyncio loop, use a thread or async HTTP client for GET—do not call
blocking ``urllib`` or synchronous ``httpx`` from inside a coroutine.
- ``token_issue_secret``: If non-empty, token requests must send ``Authorization: Bearer <secret>`` or
``X-Nanobot-Auth: <secret>``.
- ``websocket_requires_token``: If True, the handshake must include a valid token (static or issued and not expired).
- Each connection has its own session: a unique ``chat_id`` maps to the agent session internally.
- ``media`` field in outbound messages contains local filesystem paths; remote clients need a
shared filesystem or an HTTP file server to access these files.
"""
enabled: bool = False
host: str = "127.0.0.1"
port: int = 8765
path: str = "/"
token: str = ""
token_issue_path: str = ""
token_issue_secret: str = ""
token_ttl_s: int = Field(default=300, ge=30, le=86_400)
websocket_requires_token: bool = True
allow_from: list[str] = Field(default_factory=lambda: ["*"])
streaming: bool = True
max_message_bytes: int = Field(default=1_048_576, ge=1024, le=16_777_216)
ping_interval_s: float = Field(default=20.0, ge=5.0, le=300.0)
ping_timeout_s: float = Field(default=20.0, ge=5.0, le=300.0)
ssl_certfile: str = ""
ssl_keyfile: str = ""
@field_validator("path")
@classmethod
def path_must_start_with_slash(cls, value: str) -> str:
if not value.startswith("/"):
raise ValueError('path must start with "/"')
return _normalize_config_path(value)
@field_validator("token_issue_path")
@classmethod
def token_issue_path_format(cls, value: str) -> str:
value = value.strip()
if not value:
return ""
if not value.startswith("/"):
raise ValueError('token_issue_path must start with "/"')
return _normalize_config_path(value)
@model_validator(mode="after")
def token_issue_path_differs_from_ws_path(self) -> Self:
if not self.token_issue_path:
return self
if _normalize_config_path(self.token_issue_path) == _normalize_config_path(self.path):
raise ValueError("token_issue_path must differ from path (the WebSocket upgrade path)")
return self
def _http_json_response(data: dict[str, Any], *, status: int = 200) -> Response:
body = json.dumps(data, ensure_ascii=False).encode("utf-8")
headers = Headers(
[
("Date", email.utils.formatdate(usegmt=True)),
("Connection", "close"),
("Content-Length", str(len(body))),
("Content-Type", "application/json; charset=utf-8"),
]
)
reason = http.HTTPStatus(status).phrase
return Response(status, reason, headers, body)
def _parse_request_path(path_with_query: str) -> tuple[str, dict[str, list[str]]]:
"""Parse normalized path and query parameters in one pass."""
parsed = urlparse("ws://x" + path_with_query)
path = _strip_trailing_slash(parsed.path or "/")
return path, parse_qs(parsed.query)
def _normalize_http_path(path_with_query: str) -> str:
"""Return the path component (no query string), with trailing slash normalized (root stays ``/``)."""
return _parse_request_path(path_with_query)[0]
def _parse_query(path_with_query: str) -> dict[str, list[str]]:
return _parse_request_path(path_with_query)[1]
def _query_first(query: dict[str, list[str]], key: str) -> str | None:
"""Return the first value for *key*, or None."""
values = query.get(key)
return values[0] if values else None
def _parse_inbound_payload(raw: str) -> str | None:
"""Parse a client frame into text; return None for empty or unrecognized content."""
text = raw.strip()
if not text:
return None
if text.startswith("{"):
try:
data = json.loads(text)
except json.JSONDecodeError:
return text
if isinstance(data, dict):
for key in ("content", "text", "message"):
value = data.get(key)
if isinstance(value, str) and value.strip():
return value
return None
return None
return text
def _issue_route_secret_matches(headers: Any, configured_secret: str) -> bool:
"""Return True if the token-issue HTTP request carries credentials matching ``token_issue_secret``."""
if not configured_secret:
return True
authorization = headers.get("Authorization") or headers.get("authorization")
if authorization and authorization.lower().startswith("bearer "):
supplied = authorization[7:].strip()
return hmac.compare_digest(supplied, configured_secret)
header_token = headers.get("X-Nanobot-Auth") or headers.get("x-nanobot-auth")
if not header_token:
return False
return hmac.compare_digest(header_token.strip(), configured_secret)
class WebSocketChannel(BaseChannel):
"""Run a local WebSocket server; forward text/JSON messages to the message bus."""
name = "websocket"
display_name = "WebSocket"
def __init__(self, config: Any, bus: MessageBus):
if isinstance(config, dict):
config = WebSocketConfig.model_validate(config)
super().__init__(config, bus)
self.config: WebSocketConfig = config
self._connections: dict[str, Any] = {}
self._issued_tokens: dict[str, float] = {}
self._stop_event: asyncio.Event | None = None
self._server_task: asyncio.Task[None] | None = None
@classmethod
def default_config(cls) -> dict[str, Any]:
return WebSocketConfig().model_dump(by_alias=True)
def _expected_path(self) -> str:
return _normalize_config_path(self.config.path)
def _build_ssl_context(self) -> ssl.SSLContext | None:
cert = self.config.ssl_certfile.strip()
key = self.config.ssl_keyfile.strip()
if not cert and not key:
return None
if not cert or not key:
raise ValueError(
"websocket: ssl_certfile and ssl_keyfile must both be set for WSS, or both left empty"
)
ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
ctx.minimum_version = ssl.TLSVersion.TLSv1_2
ctx.load_cert_chain(certfile=cert, keyfile=key)
return ctx
_MAX_ISSUED_TOKENS = 10_000
def _purge_expired_issued_tokens(self) -> None:
now = time.monotonic()
for token_key, expiry in list(self._issued_tokens.items()):
if now > expiry:
self._issued_tokens.pop(token_key, None)
def _take_issued_token_if_valid(self, token_value: str | None) -> bool:
"""Validate and consume one issued token (single use per connection attempt).
Uses single-step pop to minimize the window between lookup and removal;
safe under asyncio's single-threaded cooperative model.
"""
if not token_value:
return False
self._purge_expired_issued_tokens()
expiry = self._issued_tokens.pop(token_value, None)
if expiry is None:
return False
if time.monotonic() > expiry:
return False
return True
def _handle_token_issue_http(self, connection: Any, request: Any) -> Any:
secret = self.config.token_issue_secret.strip()
if secret:
if not _issue_route_secret_matches(request.headers, secret):
return connection.respond(401, "Unauthorized")
else:
logger.warning(
"websocket: token_issue_path is set but token_issue_secret is empty; "
"any client can obtain connection tokens — set token_issue_secret for production."
)
self._purge_expired_issued_tokens()
if len(self._issued_tokens) >= self._MAX_ISSUED_TOKENS:
logger.error(
"websocket: too many outstanding issued tokens ({}), rejecting issuance",
len(self._issued_tokens),
)
return _http_json_response({"error": "too many outstanding tokens"}, status=429)
token_value = f"nbwt_{secrets.token_urlsafe(32)}"
self._issued_tokens[token_value] = time.monotonic() + float(self.config.token_ttl_s)
return _http_json_response(
{"token": token_value, "expires_in": self.config.token_ttl_s}
)
def _authorize_websocket_handshake(self, connection: Any, query: dict[str, list[str]]) -> Any:
supplied = _query_first(query, "token")
static_token = self.config.token.strip()
if static_token:
if supplied and hmac.compare_digest(supplied, static_token):
return None
if supplied and self._take_issued_token_if_valid(supplied):
return None
return connection.respond(401, "Unauthorized")
if self.config.websocket_requires_token:
if supplied and self._take_issued_token_if_valid(supplied):
return None
return connection.respond(401, "Unauthorized")
if supplied:
self._take_issued_token_if_valid(supplied)
return None
async def start(self) -> None:
self._running = True
self._stop_event = asyncio.Event()
ssl_context = self._build_ssl_context()
scheme = "wss" if ssl_context else "ws"
async def process_request(
connection: ServerConnection,
request: WsRequest,
) -> Any:
got, _ = _parse_request_path(request.path)
if self.config.token_issue_path:
issue_expected = _normalize_config_path(self.config.token_issue_path)
if got == issue_expected:
return self._handle_token_issue_http(connection, request)
expected_ws = self._expected_path()
if got != expected_ws:
return connection.respond(404, "Not Found")
# Early reject before WebSocket upgrade to avoid unnecessary overhead;
# _handle_message() performs a second check as defense-in-depth.
query = _parse_query(request.path)
client_id = _query_first(query, "client_id") or ""
if len(client_id) > 128:
client_id = client_id[:128]
if not self.is_allowed(client_id):
return connection.respond(403, "Forbidden")
return self._authorize_websocket_handshake(connection, query)
async def handler(connection: ServerConnection) -> None:
await self._connection_loop(connection)
logger.info(
"WebSocket server listening on {}://{}:{}{}",
scheme,
self.config.host,
self.config.port,
self.config.path,
)
if self.config.token_issue_path:
logger.info(
"WebSocket token issue route: {}://{}:{}{}",
scheme,
self.config.host,
self.config.port,
_normalize_config_path(self.config.token_issue_path),
)
async def runner() -> None:
async with serve(
handler,
self.config.host,
self.config.port,
process_request=process_request,
max_size=self.config.max_message_bytes,
ping_interval=self.config.ping_interval_s,
ping_timeout=self.config.ping_timeout_s,
ssl=ssl_context,
):
assert self._stop_event is not None
await self._stop_event.wait()
self._server_task = asyncio.create_task(runner())
await self._server_task
async def _connection_loop(self, connection: Any) -> None:
request = connection.request
path_part = request.path if request else "/"
_, query = _parse_request_path(path_part)
client_id_raw = _query_first(query, "client_id")
client_id = client_id_raw.strip() if client_id_raw else ""
if not client_id:
client_id = f"anon-{uuid.uuid4().hex[:12]}"
elif len(client_id) > 128:
logger.warning("websocket: client_id too long ({} chars), truncating", len(client_id))
client_id = client_id[:128]
chat_id = str(uuid.uuid4())
try:
await connection.send(
json.dumps(
{
"event": "ready",
"chat_id": chat_id,
"client_id": client_id,
},
ensure_ascii=False,
)
)
# Register only after ready is successfully sent to avoid out-of-order sends
self._connections[chat_id] = connection
async for raw in connection:
if isinstance(raw, bytes):
try:
raw = raw.decode("utf-8")
except UnicodeDecodeError:
logger.warning("websocket: ignoring non-utf8 binary frame")
continue
content = _parse_inbound_payload(raw)
if content is None:
continue
await self._handle_message(
sender_id=client_id,
chat_id=chat_id,
content=content,
metadata={"remote": getattr(connection, "remote_address", None)},
)
except Exception as e:
logger.debug("websocket connection ended: {}", e)
finally:
self._connections.pop(chat_id, None)
async def stop(self) -> None:
if not self._running:
return
self._running = False
if self._stop_event:
self._stop_event.set()
if self._server_task:
try:
await self._server_task
except Exception as e:
logger.warning("websocket: server task error during shutdown: {}", e)
self._server_task = None
self._connections.clear()
self._issued_tokens.clear()
async def _safe_send(self, chat_id: str, raw: str, *, label: str = "") -> None:
"""Send a raw frame, cleaning up dead connections on ConnectionClosed."""
connection = self._connections.get(chat_id)
if connection is None:
return
try:
await connection.send(raw)
except ConnectionClosed:
self._connections.pop(chat_id, None)
logger.warning("websocket{}connection gone for chat_id={}", label, chat_id)
except Exception as e:
logger.error("websocket{}send failed: {}", label, e)
raise
async def send(self, msg: OutboundMessage) -> None:
connection = self._connections.get(msg.chat_id)
if connection is None:
logger.warning("websocket: no active connection for chat_id={}", msg.chat_id)
return
payload: dict[str, Any] = {
"event": "message",
"text": msg.content,
}
if msg.media:
payload["media"] = msg.media
if msg.reply_to:
payload["reply_to"] = msg.reply_to
raw = json.dumps(payload, ensure_ascii=False)
await self._safe_send(msg.chat_id, raw, label=" ")
async def send_delta(
self,
chat_id: str,
delta: str,
metadata: dict[str, Any] | None = None,
) -> None:
if self._connections.get(chat_id) is None:
return
meta = metadata or {}
if meta.get("_stream_end"):
body: dict[str, Any] = {"event": "stream_end"}
else:
body = {
"event": "delta",
"text": delta,
}
if meta.get("_stream_id") is not None:
body["stream_id"] = meta["_stream_id"]
raw = json.dumps(body, ensure_ascii=False)
await self._safe_send(chat_id, raw, label=" stream ")
+195 -26
View File
@@ -1,9 +1,13 @@
"""WeCom (Enterprise WeChat) channel implementation using wecom_aibot_sdk.""" """WeCom (Enterprise WeChat) channel implementation using wecom_aibot_sdk."""
import asyncio import asyncio
import base64
import hashlib
import importlib.util import importlib.util
import os import os
import re
from collections import OrderedDict from collections import OrderedDict
from pathlib import Path
from typing import Any from typing import Any
from loguru import logger from loguru import logger
@@ -17,6 +21,37 @@ from pydantic import Field
WECOM_AVAILABLE = importlib.util.find_spec("wecom_aibot_sdk") is not None WECOM_AVAILABLE = importlib.util.find_spec("wecom_aibot_sdk") is not None
# Upload safety limits (matching QQ channel defaults)
WECOM_UPLOAD_MAX_BYTES = 1024 * 1024 * 200 # 200MB
# Replace unsafe characters with "_", keep Chinese and common safe punctuation.
_SAFE_NAME_RE = re.compile(r"[^\w.\-()\[\]()【】\u4e00-\u9fff]+", re.UNICODE)
def _sanitize_filename(name: str) -> str:
"""Sanitize filename to avoid traversal and problematic chars."""
name = (name or "").strip()
name = Path(name).name
name = _SAFE_NAME_RE.sub("_", name).strip("._ ")
return name
_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"}
_VIDEO_EXTS = {".mp4", ".avi", ".mov"}
_AUDIO_EXTS = {".amr", ".mp3", ".wav", ".ogg"}
def _guess_wecom_media_type(filename: str) -> str:
"""Classify file extension as WeCom media_type string."""
ext = Path(filename).suffix.lower()
if ext in _IMAGE_EXTS:
return "image"
if ext in _VIDEO_EXTS:
return "video"
if ext in _AUDIO_EXTS:
return "voice"
return "file"
class WecomConfig(Base): class WecomConfig(Base):
"""WeCom (Enterprise WeChat) AI Bot channel configuration.""" """WeCom (Enterprise WeChat) AI Bot channel configuration."""
@@ -217,6 +252,7 @@ class WecomChannel(BaseChannel):
chat_id = body.get("chatid", sender_id) chat_id = body.get("chatid", sender_id)
content_parts = [] content_parts = []
media_paths: list[str] = []
if msg_type == "text": if msg_type == "text":
text = body.get("text", {}).get("content", "") text = body.get("text", {}).get("content", "")
@@ -232,7 +268,8 @@ class WecomChannel(BaseChannel):
file_path = await self._download_and_save_media(file_url, aes_key, "image") file_path = await self._download_and_save_media(file_url, aes_key, "image")
if file_path: if file_path:
filename = os.path.basename(file_path) filename = os.path.basename(file_path)
content_parts.append(f"[image: {filename}]\n[Image: source: {file_path}]") content_parts.append(f"[image: {filename}]")
media_paths.append(file_path)
else: else:
content_parts.append("[image: download failed]") content_parts.append("[image: download failed]")
else: else:
@@ -256,7 +293,8 @@ class WecomChannel(BaseChannel):
if file_url and aes_key: if file_url and aes_key:
file_path = await self._download_and_save_media(file_url, aes_key, "file", file_name) file_path = await self._download_and_save_media(file_url, aes_key, "file", file_name)
if file_path: if file_path:
content_parts.append(f"[file: {file_name}]\n[File: source: {file_path}]") content_parts.append(f"[file: {file_name}]")
media_paths.append(file_path)
else: else:
content_parts.append(f"[file: {file_name}: download failed]") content_parts.append(f"[file: {file_name}: download failed]")
else: else:
@@ -286,12 +324,11 @@ class WecomChannel(BaseChannel):
self._chat_frames[chat_id] = frame self._chat_frames[chat_id] = frame
# Forward to message bus # Forward to message bus
# Note: media paths are included in content for broader model compatibility
await self._handle_message( await self._handle_message(
sender_id=sender_id, sender_id=sender_id,
chat_id=chat_id, chat_id=chat_id,
content=content, content=content,
media=None, media=media_paths or None,
metadata={ metadata={
"message_id": msg_id, "message_id": msg_id,
"msg_type": msg_type, "msg_type": msg_type,
@@ -322,13 +359,21 @@ class WecomChannel(BaseChannel):
logger.warning("Failed to download media from WeCom") logger.warning("Failed to download media from WeCom")
return None return None
if len(data) > WECOM_UPLOAD_MAX_BYTES:
logger.warning(
"WeCom inbound media too large: {} bytes (max {})",
len(data),
WECOM_UPLOAD_MAX_BYTES,
)
return None
media_dir = get_media_dir("wecom") media_dir = get_media_dir("wecom")
if not filename: if not filename:
filename = fname or f"{media_type}_{hash(file_url) % 100000}" filename = fname or f"{media_type}_{hash(file_url) % 100000}"
filename = os.path.basename(filename) filename = _sanitize_filename(filename)
file_path = media_dir / filename file_path = media_dir / filename
file_path.write_bytes(data) await asyncio.to_thread(file_path.write_bytes, data)
logger.debug("Downloaded {} to {}", media_type, file_path) logger.debug("Downloaded {} to {}", media_type, file_path)
return str(file_path) return str(file_path)
@@ -336,6 +381,100 @@ class WecomChannel(BaseChannel):
logger.error("Error downloading media: {}", e) logger.error("Error downloading media: {}", e)
return None return None
async def _upload_media_ws(
self, client: Any, file_path: str,
) -> "tuple[str, str] | tuple[None, None]":
"""Upload a local file to WeCom via WebSocket 3-step protocol (base64).
Uses the WeCom WebSocket upload commands directly via
``client._ws_manager.send_reply()``:
``aibot_upload_media_init`` → upload_id
``aibot_upload_media_chunk`` × N (≤512 KB raw per chunk, base64)
``aibot_upload_media_finish`` → media_id
Returns (media_id, media_type) on success, (None, None) on failure.
"""
from wecom_aibot_sdk.utils import generate_req_id as _gen_req_id
try:
fname = os.path.basename(file_path)
media_type = _guess_wecom_media_type(fname)
# Read file size and data in a thread to avoid blocking the event loop
def _read_file():
file_size = os.path.getsize(file_path)
if file_size > WECOM_UPLOAD_MAX_BYTES:
raise ValueError(
f"File too large: {file_size} bytes (max {WECOM_UPLOAD_MAX_BYTES})"
)
with open(file_path, "rb") as f:
return file_size, f.read()
file_size, data = await asyncio.to_thread(_read_file)
# MD5 is used for file integrity only, not cryptographic security
md5_hash = hashlib.md5(data).hexdigest()
CHUNK_SIZE = 512 * 1024 # 512 KB raw (before base64)
mv = memoryview(data)
chunk_list = [bytes(mv[i : i + CHUNK_SIZE]) for i in range(0, file_size, CHUNK_SIZE)]
n_chunks = len(chunk_list)
del mv, data
# Step 1: init
req_id = _gen_req_id("upload_init")
resp = await client._ws_manager.send_reply(req_id, {
"type": media_type,
"filename": fname,
"total_size": file_size,
"total_chunks": n_chunks,
"md5": md5_hash,
}, "aibot_upload_media_init")
if resp.errcode != 0:
logger.warning("WeCom upload init failed ({}): {}", resp.errcode, resp.errmsg)
return None, None
upload_id = resp.body.get("upload_id") if resp.body else None
if not upload_id:
logger.warning("WeCom upload init: no upload_id in response")
return None, None
# Step 2: send chunks
for i, chunk in enumerate(chunk_list):
req_id = _gen_req_id("upload_chunk")
resp = await client._ws_manager.send_reply(req_id, {
"upload_id": upload_id,
"chunk_index": i,
"base64_data": base64.b64encode(chunk).decode(),
}, "aibot_upload_media_chunk")
if resp.errcode != 0:
logger.warning("WeCom upload chunk {} failed ({}): {}", i, resp.errcode, resp.errmsg)
return None, None
# Step 3: finish
req_id = _gen_req_id("upload_finish")
resp = await client._ws_manager.send_reply(req_id, {
"upload_id": upload_id,
}, "aibot_upload_media_finish")
if resp.errcode != 0:
logger.warning("WeCom upload finish failed ({}): {}", resp.errcode, resp.errmsg)
return None, None
media_id = resp.body.get("media_id") if resp.body else None
if not media_id:
logger.warning("WeCom upload finish: no media_id in response body={}", resp.body)
return None, None
suffix = "..." if len(media_id) > 16 else ""
logger.debug("WeCom uploaded {} ({}) → media_id={}", fname, media_type, media_id[:16] + suffix)
return media_id, media_type
except ValueError as e:
logger.warning("WeCom upload skipped for {}: {}", file_path, e)
return None, None
except Exception as e:
logger.error("WeCom _upload_media_ws error for {}: {}", file_path, e)
return None, None
async def send(self, msg: OutboundMessage) -> None: async def send(self, msg: OutboundMessage) -> None:
"""Send a message through WeCom.""" """Send a message through WeCom."""
if not self._client: if not self._client:
@@ -343,29 +482,59 @@ class WecomChannel(BaseChannel):
return return
try: try:
content = msg.content.strip() content = (msg.content or "").strip()
if not content: is_progress = bool(msg.metadata.get("_progress"))
return
# Get the stored frame for this chat # Get the stored frame for this chat
frame = self._chat_frames.get(msg.chat_id) frame = self._chat_frames.get(msg.chat_id)
if not frame:
logger.warning("No frame found for chat {}, cannot reply", msg.chat_id) # Send media files via WebSocket upload
for file_path in msg.media or []:
if not os.path.isfile(file_path):
logger.warning("WeCom media file not found: {}", file_path)
continue
media_id, media_type = await self._upload_media_ws(self._client, file_path)
if media_id:
if frame:
await self._client.reply(frame, {
"msgtype": media_type,
media_type: {"media_id": media_id},
})
else:
await self._client.send_message(msg.chat_id, {
"msgtype": media_type,
media_type: {"media_id": media_id},
})
logger.debug("WeCom sent {}{}", media_type, msg.chat_id)
else:
content += f"\n[file upload failed: {os.path.basename(file_path)}]"
if not content:
return return
# Use streaming reply for better UX if frame:
stream_id = self._generate_req_id("stream") # Both progress and final messages must use reply_stream (cmd="aibot_respond_msg").
# The plain reply() uses cmd="reply" which does not support "text" msgtype
# and causes errcode=40008 from WeCom API.
stream_id = self._generate_req_id("stream")
await self._client.reply_stream(
frame,
stream_id,
content,
finish=not is_progress,
)
logger.debug(
"WeCom {} sent to {}",
"progress" if is_progress else "message",
msg.chat_id,
)
else:
# No frame (e.g. cron push): proactive send only supports markdown
await self._client.send_message(msg.chat_id, {
"msgtype": "markdown",
"markdown": {"content": content},
})
logger.info("WeCom proactive send to {}", msg.chat_id)
# Send as streaming message with finish=True except Exception:
await self._client.reply_stream( logger.exception("Error sending WeCom message to chat_id={}", msg.chat_id)
frame,
stream_id,
content,
finish=True,
)
logger.debug("WeCom message sent to {}", msg.chat_id)
except Exception as e:
logger.error("Error sending WeCom message: {}", e)
raise
+3
View File
@@ -590,6 +590,7 @@ def serve(
mcp_servers=runtime_config.tools.mcp_servers, mcp_servers=runtime_config.tools.mcp_servers,
channels_config=runtime_config.channels, channels_config=runtime_config.channels,
timezone=runtime_config.agents.defaults.timezone, timezone=runtime_config.agents.defaults.timezone,
session_ttl_minutes=runtime_config.agents.defaults.session_ttl_minutes,
) )
model_name = runtime_config.agents.defaults.model model_name = runtime_config.agents.defaults.model
@@ -681,6 +682,7 @@ def gateway(
mcp_servers=config.tools.mcp_servers, mcp_servers=config.tools.mcp_servers,
channels_config=config.channels, channels_config=config.channels,
timezone=config.agents.defaults.timezone, timezone=config.agents.defaults.timezone,
session_ttl_minutes=config.agents.defaults.session_ttl_minutes,
) )
# Set cron callback (needs agent) # Set cron callback (needs agent)
@@ -912,6 +914,7 @@ def agent(
mcp_servers=config.tools.mcp_servers, mcp_servers=config.tools.mcp_servers,
channels_config=config.channels, channels_config=config.channels,
timezone=config.agents.defaults.timezone, timezone=config.agents.defaults.timezone,
session_ttl_minutes=config.agents.defaults.session_ttl_minutes,
) )
restart_notice = consume_restart_notice_from_env() restart_notice = consume_restart_notice_from_env()
if restart_notice and should_show_cli_restart_notice(restart_notice, session_id): if restart_notice and should_show_cli_restart_notice(restart_notice, session_id):
+1
View File
@@ -76,6 +76,7 @@ class AgentDefaults(Base):
provider_retry_mode: Literal["standard", "persistent"] = "standard" provider_retry_mode: Literal["standard", "persistent"] = "standard"
reasoning_effort: str | None = None # low / medium / high - enables LLM thinking mode reasoning_effort: str | None = None # low / medium / high - enables LLM thinking mode
timezone: str = "UTC" # IANA timezone, e.g. "Asia/Shanghai", "America/New_York" timezone: str = "UTC" # IANA timezone, e.g. "Asia/Shanghai", "America/New_York"
session_ttl_minutes: int = Field(default=0, ge=0) # Auto /new after idle (0 = disabled)
dream: DreamConfig = Field(default_factory=DreamConfig) dream: DreamConfig = Field(default_factory=DreamConfig)
+1
View File
@@ -81,6 +81,7 @@ class Nanobot:
restrict_to_workspace=config.tools.restrict_to_workspace, restrict_to_workspace=config.tools.restrict_to_workspace,
mcp_servers=config.tools.mcp_servers, mcp_servers=config.tools.mcp_servers,
timezone=defaults.timezone, timezone=defaults.timezone,
session_ttl_minutes=defaults.session_ttl_minutes,
) )
return cls(loop) return cls(loop)
+3
View File
@@ -155,6 +155,7 @@ class SessionManager:
messages = [] messages = []
metadata = {} metadata = {}
created_at = None created_at = None
updated_at = None
last_consolidated = 0 last_consolidated = 0
with open(path, encoding="utf-8") as f: with open(path, encoding="utf-8") as f:
@@ -168,6 +169,7 @@ class SessionManager:
if data.get("_type") == "metadata": if data.get("_type") == "metadata":
metadata = data.get("metadata", {}) metadata = data.get("metadata", {})
created_at = datetime.fromisoformat(data["created_at"]) if data.get("created_at") else None created_at = datetime.fromisoformat(data["created_at"]) if data.get("created_at") else None
updated_at = datetime.fromisoformat(data["updated_at"]) if data.get("updated_at") else None
last_consolidated = data.get("last_consolidated", 0) last_consolidated = data.get("last_consolidated", 0)
else: else:
messages.append(data) messages.append(data)
@@ -176,6 +178,7 @@ class SessionManager:
key=key, key=key,
messages=messages, messages=messages,
created_at=created_at or datetime.now(), created_at=created_at or datetime.now(),
updated_at=updated_at or datetime.now(),
metadata=metadata, metadata=metadata,
last_consolidated=last_consolidated last_consolidated=last_consolidated
) )
+18 -8
View File
@@ -1,13 +1,23 @@
Compare conversation history against current memory files. Compare conversation history against current memory files. Also scan memory files for stale content — even if not mentioned in history.
Output one line per finding:
[FILE] atomic fact or change description
Files: USER (identity, preferences, habits), SOUL (bot behavior, tone), MEMORY (knowledge, project context, tool patterns) Output one line per finding:
[FILE] atomic fact (not already in memory)
[FILE-REMOVE] reason for removal
Files: USER (identity, preferences), SOUL (bot behavior, tone), MEMORY (knowledge, project context)
Rules: Rules:
- Only new or conflicting information — skip duplicates and ephemera - Atomic facts: "has a cat named Luna" not "discussed pet care"
- Prefer atomic facts: "has a cat named Luna" not "discussed pet care"
- Corrections: [USER] location is Tokyo, not Osaka - Corrections: [USER] location is Tokyo, not Osaka
- Also capture confirmed approaches: if the user validated a non-obvious choice, note it - Capture confirmed approaches the user validated
If nothing needs updating: [SKIP] no new information Staleness — flag for [FILE-REMOVE]:
- Time-sensitive data older than 14 days: weather, daily status, one-time meetings, passed events
- Completed one-time tasks: triage, one-time reviews, finished research, resolved incidents
- Resolved tracking: merged/closed PRs, fixed issues, completed migrations
- Detailed incident info after 14 days — reduce to one-line summary
- Superseded: approaches replaced by newer solutions, deprecated dependencies
Do not add: current weather, transient status, temporary errors, conversational filler.
[SKIP] if nothing needs updating.
+18 -7
View File
@@ -1,13 +1,24 @@
Update memory files based on the analysis below. Update memory files based on the analysis below.
- [FILE] entries: add the described content to the appropriate file
- [FILE-REMOVE] entries: delete the corresponding content from memory files
## Quality standards ## File paths (relative to workspace root)
- Every line must carry standalone value — no filler - SOUL.md
- Concise bullet points under clear headers - USER.md
- Remove outdated or contradicted information - memory/MEMORY.md
## Editing Do NOT guess paths.
- File contents provided below — edit directly, no read_file needed
## Editing rules
- Edit directly — file contents provided below, no read_file needed
- Use exact text as old_text, include surrounding blank lines for unique match
- Batch changes to the same file into one edit_file call - Batch changes to the same file into one edit_file call
- For deletions: section header + all bullets as old_text, new_text empty
- Surgical edits only — never rewrite entire files - Surgical edits only — never rewrite entire files
- Do NOT overwrite correct entries — only add, update, or remove
- If nothing to update, stop without calling tools - If nothing to update, stop without calling tools
## Quality
- Every line must carry standalone value
- Concise bullets under clear headers
- When reducing (not deleting): keep essential facts, drop verbose details
- If uncertain whether to delete, keep but add "(verify currency)"
+4 -1
View File
@@ -15,9 +15,12 @@ from loguru import logger
def strip_think(text: str) -> str: def strip_think(text: str) -> str:
"""Remove <think>…</think> blocks and any unclosed trailing <think> tag.""" """Remove thinking blocks and any unclosed trailing tag."""
text = re.sub(r"<think>[\s\S]*?</think>", "", text) text = re.sub(r"<think>[\s\S]*?</think>", "", text)
text = re.sub(r"<think>[\s\S]*$", "", text) text = re.sub(r"<think>[\s\S]*$", "", text)
# Gemma 4 and similar models use <thought>...</thought> blocks
text = re.sub(r"<thought>[\s\S]*?</thought>", "", text)
text = re.sub(r"<thought>[\s\S]*$", "", text)
return text.strip() return text.strip()
+34 -23
View File
@@ -2,6 +2,8 @@
from __future__ import annotations from __future__ import annotations
import re
from nanobot.utils.path import abbreviate_path from nanobot.utils.path import abbreviate_path
# Registry: tool_name -> (key_args, template, is_path, is_command) # Registry: tool_name -> (key_args, template, is_path, is_command)
@@ -17,27 +19,37 @@ _TOOL_FORMATS: dict[str, tuple[list[str], str, bool, bool]] = {
"list_dir": (["path"], "ls {}", True, False), "list_dir": (["path"], "ls {}", True, False),
} }
# Matches file paths embedded in shell commands (Windows drive, ~/, or absolute after space)
_PATH_IN_CMD_RE = re.compile(
r"(?:[A-Za-z]:[/\\]|~/|(?<=\s)/)[^\s;&|<>\"']+"
)
def format_tool_hints(tool_calls: list) -> str: def format_tool_hints(tool_calls: list) -> str:
"""Format tool calls as concise hints with smart abbreviation.""" """Format tool calls as concise hints with smart abbreviation."""
if not tool_calls: if not tool_calls:
return "" return ""
hints = [] formatted = []
for name, count, example_tc in _group_consecutive(tool_calls): for tc in tool_calls:
fmt = _TOOL_FORMATS.get(name) fmt = _TOOL_FORMATS.get(tc.name)
if fmt: if fmt:
hint = _fmt_known(example_tc, fmt) formatted.append(_fmt_known(tc, fmt))
elif name.startswith("mcp_"): elif tc.name.startswith("mcp_"):
hint = _fmt_mcp(example_tc) formatted.append(_fmt_mcp(tc))
else: else:
hint = _fmt_fallback(example_tc) formatted.append(_fmt_fallback(tc))
if count > 1: hints = []
hint = f"{hint} \u00d7 {count}" for hint in formatted:
hints.append(hint) if hints and hints[-1][0] == hint:
hints[-1] = (hint, hints[-1][1] + 1)
else:
hints.append((hint, 1))
return ", ".join(hints) return ", ".join(
f"{h} \u00d7 {c}" if c > 1 else h for h, c in hints
)
def _get_args(tc) -> dict: def _get_args(tc) -> dict:
@@ -51,17 +63,6 @@ def _get_args(tc) -> dict:
return {} return {}
def _group_consecutive(calls: list) -> list[tuple[str, int, object]]:
"""Group consecutive calls to the same tool: [(name, count, first), ...]."""
groups: list[tuple[str, int, object]] = []
for tc in calls:
if groups and groups[-1][0] == tc.name:
groups[-1] = (groups[-1][0], groups[-1][1] + 1, groups[-1][2])
else:
groups.append((tc.name, 1, tc))
return groups
def _extract_arg(tc, key_args: list[str]) -> str | None: def _extract_arg(tc, key_args: list[str]) -> str | None:
"""Extract the first available value from preferred key names.""" """Extract the first available value from preferred key names."""
args = _get_args(tc) args = _get_args(tc)
@@ -85,10 +86,20 @@ def _fmt_known(tc, fmt: tuple) -> str:
if fmt[2]: # is_path if fmt[2]: # is_path
val = abbreviate_path(val) val = abbreviate_path(val)
elif fmt[3]: # is_command elif fmt[3]: # is_command
val = val[:40] + "\u2026" if len(val) > 40 else val val = _abbreviate_command(val)
return fmt[1].format(val) return fmt[1].format(val)
def _abbreviate_command(cmd: str, max_len: int = 40) -> str:
"""Abbreviate paths in a command string, then truncate."""
abbreviated = _PATH_IN_CMD_RE.sub(
lambda m: abbreviate_path(m.group(), max_len=25), cmd
)
if len(abbreviated) <= max_len:
return abbreviated
return abbreviated[:max_len - 1] + "\u2026"
def _fmt_mcp(tc) -> str: def _fmt_mcp(tc) -> str:
"""Format MCP tool as server::tool.""" """Format MCP tool as server::tool."""
name = tc.name name = tc.name
+4
View File
@@ -63,6 +63,10 @@ weixin = [
"qrcode[pil]>=8.0", "qrcode[pil]>=8.0",
"pycryptodome>=3.20.0", "pycryptodome>=3.20.0",
] ]
msteams = [
"PyJWT>=2.0,<3.0",
"cryptography>=41.0",
]
matrix = [ matrix = [
"matrix-nio[e2e]>=0.25.2", "matrix-nio[e2e]>=0.25.2",
+931
View File
@@ -0,0 +1,931 @@
"""Tests for auto compact (idle TTL) feature."""
import asyncio
from datetime import datetime, timedelta
from unittest.mock import AsyncMock, MagicMock
from pathlib import Path
import pytest
from nanobot.agent.loop import AgentLoop
from nanobot.bus.events import InboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.config.schema import AgentDefaults
from nanobot.command import CommandContext
from nanobot.providers.base import LLMResponse
def _make_loop(tmp_path: Path, session_ttl_minutes: int = 15) -> AgentLoop:
"""Create a minimal AgentLoop for testing."""
bus = MessageBus()
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.estimate_prompt_tokens.return_value = (10_000, "test")
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
provider.generation.max_tokens = 4096
loop = AgentLoop(
bus=bus,
provider=provider,
workspace=tmp_path,
model="test-model",
context_window_tokens=128_000,
session_ttl_minutes=session_ttl_minutes,
)
loop.tools.get_definitions = MagicMock(return_value=[])
return loop
class TestSessionTTLConfig:
"""Test session TTL configuration."""
def test_default_ttl_is_zero(self):
"""Default TTL should be 0 (disabled)."""
defaults = AgentDefaults()
assert defaults.session_ttl_minutes == 0
def test_custom_ttl(self):
"""Custom TTL should be stored correctly."""
defaults = AgentDefaults(session_ttl_minutes=30)
assert defaults.session_ttl_minutes == 30
class TestAgentLoopTTLParam:
"""Test that AutoCompact receives and stores session_ttl_minutes."""
def test_loop_stores_ttl(self, tmp_path):
"""AutoCompact should store the TTL value."""
loop = _make_loop(tmp_path, session_ttl_minutes=25)
assert loop.auto_compact._ttl == 25
def test_loop_default_ttl_zero(self, tmp_path):
"""AutoCompact default TTL should be 0 (disabled)."""
loop = _make_loop(tmp_path, session_ttl_minutes=0)
assert loop.auto_compact._ttl == 0
class TestAutoCompact:
"""Test the _archive method."""
@pytest.mark.asyncio
async def test_is_expired_boundary(self, tmp_path):
"""Exactly at TTL boundary should be expired (>= not >)."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
ts = datetime.now() - timedelta(minutes=15)
assert loop.auto_compact._is_expired(ts) is True
ts2 = datetime.now() - timedelta(minutes=14, seconds=59)
assert loop.auto_compact._is_expired(ts2) is False
await loop.close_mcp()
@pytest.mark.asyncio
async def test_is_expired_string_timestamp(self, tmp_path):
"""_is_expired should parse ISO string timestamps."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
ts = (datetime.now() - timedelta(minutes=20)).isoformat()
assert loop.auto_compact._is_expired(ts) is True
assert loop.auto_compact._is_expired(None) is False
assert loop.auto_compact._is_expired("") is False
await loop.close_mcp()
@pytest.mark.asyncio
async def test_check_expired_only_archives_expired_sessions(self, tmp_path):
"""With multiple sessions, only the expired one should be archived."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
# Expired session
s1 = loop.sessions.get_or_create("cli:expired")
s1.add_message("user", "old")
s1.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(s1)
# Active session
s2 = loop.sessions.get_or_create("cli:active")
s2.add_message("user", "recent")
loop.sessions.save(s2)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
loop.auto_compact.check_expired(loop._schedule_background)
await asyncio.sleep(0.1)
active_after = loop.sessions.get_or_create("cli:active")
assert len(active_after.messages) == 1
assert active_after.messages[0]["content"] == "recent"
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_archives_and_clears(self, tmp_path):
"""_archive should archive un-consolidated messages and clear session."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
for i in range(4):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
loop.sessions.save(session)
archived_messages = []
async def _fake_archive(messages):
archived_messages.extend(messages)
return True
loop.consolidator.archive = _fake_archive
await loop.auto_compact._archive("cli:test")
assert len(archived_messages) == 8
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_stores_summary(self, tmp_path):
"""_archive should store the summary in _summaries."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "hello")
session.add_message("assistant", "hi there")
loop.sessions.save(session)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "User said hello.",
}
await loop.auto_compact._archive("cli:test")
entry = loop.auto_compact._summaries.get("cli:test")
assert entry is not None
assert entry[0] == "User said hello."
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_empty_session(self, tmp_path):
"""_archive on empty session should not archive."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
archive_called = False
async def _fake_archive(messages):
nonlocal archive_called
archive_called = True
return True
loop.consolidator.archive = _fake_archive
await loop.auto_compact._archive("cli:test")
assert not archive_called
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_respects_last_consolidated(self, tmp_path):
"""_archive should only archive un-consolidated messages."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
for i in range(10):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
session.last_consolidated = 18
loop.sessions.save(session)
archived_count = 0
async def _fake_archive(messages):
nonlocal archived_count
archived_count = len(messages)
return True
loop.consolidator.archive = _fake_archive
await loop.auto_compact._archive("cli:test")
assert archived_count == 2
await loop.close_mcp()
class TestAutoCompactIdleDetection:
"""Test idle detection triggers auto-new in _process_message."""
@pytest.mark.asyncio
async def test_no_auto_compact_when_ttl_disabled(self, tmp_path):
"""No auto-new should happen when TTL is 0 (disabled)."""
loop = _make_loop(tmp_path, session_ttl_minutes=0)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.updated_at = datetime.now() - timedelta(minutes=30)
loop.sessions.save(session)
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="new msg")
await loop._process_message(msg)
session_after = loop.sessions.get_or_create("cli:test")
assert any(m["content"] == "old message" for m in session_after.messages)
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_triggers_on_idle(self, tmp_path):
"""Proactive auto-new archives expired session; _process_message reloads it."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archived_messages = []
async def _fake_archive(messages):
archived_messages.extend(messages)
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "Summary.",
}
# Simulate proactive archive completing before message arrives
await loop.auto_compact._archive("cli:test")
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="new msg")
await loop._process_message(msg)
session_after = loop.sessions.get_or_create("cli:test")
assert not any(m["content"] == "old message" for m in session_after.messages)
assert any(m["content"] == "new msg" for m in session_after.messages)
await loop.close_mcp()
@pytest.mark.asyncio
async def test_no_auto_compact_when_active(self, tmp_path):
"""No auto-new should happen when session is recently active."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "recent message")
loop.sessions.save(session)
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="new msg")
await loop._process_message(msg)
session_after = loop.sessions.get_or_create("cli:test")
assert any(m["content"] == "recent message" for m in session_after.messages)
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_does_not_affect_priority_commands(self, tmp_path):
"""Priority commands (/stop, /restart) bypass _process_message entirely via run()."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
# Priority commands are dispatched in run() before _process_message is called.
# Simulate that path directly via dispatch_priority.
raw = "/stop"
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content=raw)
ctx = CommandContext(msg=msg, session=session, key="cli:test", raw=raw, loop=loop)
result = await loop.commands.dispatch_priority(ctx)
assert result is not None
assert "stopped" in result.content.lower() or "no active task" in result.content.lower()
# Session should be untouched since priority commands skip _process_message
session_after = loop.sessions.get_or_create("cli:test")
assert any(m["content"] == "old message" for m in session_after.messages)
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_with_slash_new(self, tmp_path):
"""Auto-new fires before /new dispatches; session is cleared twice but idempotent."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
for i in range(4):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(msg)
assert response is not None
assert "new session started" in response.content.lower()
session_after = loop.sessions.get_or_create("cli:test")
# Session is empty (auto-new archived and cleared, /new cleared again)
assert len(session_after.messages) == 0
await loop.close_mcp()
class TestAutoCompactSystemMessages:
"""Test that auto-new also works for system messages."""
@pytest.mark.asyncio
async def test_auto_compact_triggers_for_system_messages(self, tmp_path):
"""Proactive auto-new archives expired session; system messages reload it."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message from subagent context")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "Summary.",
}
# Simulate proactive archive completing before system message arrives
await loop.auto_compact._archive("cli:test")
msg = InboundMessage(
channel="system", sender_id="subagent", chat_id="cli:test",
content="subagent result",
)
await loop._process_message(msg)
session_after = loop.sessions.get_or_create("cli:test")
assert not any(
m["content"] == "old message from subagent context"
for m in session_after.messages
)
await loop.close_mcp()
class TestAutoCompactEdgeCases:
"""Edge cases for auto session new."""
@pytest.mark.asyncio
async def test_auto_compact_with_nothing_summary(self, tmp_path):
"""Auto-new should not inject when archive produces '(nothing)'."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "thanks")
session.add_message("assistant", "you're welcome")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="(nothing)", tool_calls=[])
)
await loop.auto_compact._archive("cli:test")
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
# "(nothing)" summary should not be stored
assert "cli:test" not in loop.auto_compact._summaries
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_archive_failure_still_clears(self, tmp_path):
"""Auto-new should clear session even if LLM archive fails (raw_archive fallback)."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "important data")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
loop.provider.chat_with_retry = AsyncMock(side_effect=Exception("API down"))
# Should not raise
await loop.auto_compact._archive("cli:test")
session_after = loop.sessions.get_or_create("cli:test")
# Session should be cleared (archive falls back to raw dump)
assert len(session_after.messages) == 0
await loop.close_mcp()
@pytest.mark.asyncio
async def test_auto_compact_preserves_runtime_checkpoint_before_check(self, tmp_path):
"""Runtime checkpoint is restored; proactive archive handles the expired session."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.metadata[AgentLoop._RUNTIME_CHECKPOINT_KEY] = {
"assistant_message": {"role": "assistant", "content": "interrupted response"},
"completed_tool_results": [],
"pending_tool_calls": [],
}
session.add_message("user", "previous message")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archived_messages = []
async def _fake_archive(messages):
archived_messages.extend(messages)
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "Summary.",
}
# Simulate proactive archive completing before message arrives
await loop.auto_compact._archive("cli:test")
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="continue")
await loop._process_message(msg)
# The checkpoint-restored message should have been archived by proactive path
assert len(archived_messages) >= 1
await loop.close_mcp()
class TestAutoCompactIntegration:
"""End-to-end test of auto session new feature."""
@pytest.mark.asyncio
async def test_full_lifecycle(self, tmp_path):
"""
Full lifecycle: messages -> idle -> auto-new -> archive -> clear -> summary injected as runtime context.
"""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
# Phase 1: User has a conversation
session.add_message("user", "I'm learning English, teach me past tense")
session.add_message("assistant", "Past tense is used for actions completed in the past...")
session.add_message("user", "Give me an example")
session.add_message("assistant", '"I walked to the store yesterday."')
loop.sessions.save(session)
# Phase 2: Time passes (simulate idle)
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
# Phase 3: User returns with a new message
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(
content="User is learning English past tense. Example: 'I walked to the store yesterday.'",
tool_calls=[],
)
)
msg = InboundMessage(
channel="cli", sender_id="user", chat_id="test",
content="Let's continue, teach me present perfect",
)
response = await loop._process_message(msg)
# Phase 4: Verify
session_after = loop.sessions.get_or_create("cli:test")
# Old messages should be gone
assert not any(
"past tense is used" in str(m.get("content", "")) for m in session_after.messages
)
# Summary should NOT be persisted in session (ephemeral, one-shot)
assert not any(
"[Resumed Session]" in str(m.get("content", "")) for m in session_after.messages
)
# Runtime context end marker should NOT be persisted
assert not any(
"[/Runtime Context]" in str(m.get("content", "")) for m in session_after.messages
)
# Pending summary should be consumed (one-shot)
assert "cli:test" not in loop.auto_compact._summaries
# The new message should be processed (response exists)
assert response is not None
await loop.close_mcp()
@pytest.mark.asyncio
async def test_multi_paragraph_user_message_preserved(self, tmp_path):
"""Multi-paragraph user messages must be fully preserved after auto-new."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "Summary.",
}
# Simulate proactive archive completing before message arrives
await loop.auto_compact._archive("cli:test")
msg = InboundMessage(
channel="cli", sender_id="user", chat_id="test",
content="Paragraph one\n\nParagraph two\n\nParagraph three",
)
await loop._process_message(msg)
session_after = loop.sessions.get_or_create("cli:test")
user_msgs = [m for m in session_after.messages if m.get("role") == "user"]
assert len(user_msgs) >= 1
# All three paragraphs must be preserved
persisted = user_msgs[-1]["content"]
assert "Paragraph one" in persisted
assert "Paragraph two" in persisted
assert "Paragraph three" in persisted
# No runtime context markers in persisted message
assert "[Runtime Context" not in persisted
assert "[/Runtime Context]" not in persisted
await loop.close_mcp()
class TestProactiveAutoCompact:
"""Test proactive auto-new on idle ticks (TimeoutError path in run loop)."""
@staticmethod
async def _run_check_expired(loop):
"""Helper: run check_expired via callback and wait for background tasks."""
loop.auto_compact.check_expired(loop._schedule_background)
await asyncio.sleep(0.1)
@pytest.mark.asyncio
async def test_no_check_when_ttl_disabled(self, tmp_path):
"""check_expired should be a no-op when TTL is 0."""
loop = _make_loop(tmp_path, session_ttl_minutes=0)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.updated_at = datetime.now() - timedelta(minutes=30)
loop.sessions.save(session)
await self._run_check_expired(loop)
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 1
await loop.close_mcp()
@pytest.mark.asyncio
async def test_proactive_archive_on_idle_tick(self, tmp_path):
"""Expired session should be archived during idle tick."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.add_message("assistant", "old response")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archived_messages = []
async def _fake_archive(messages):
archived_messages.extend(messages)
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "User chatted about old things.",
}
await self._run_check_expired(loop)
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
assert len(archived_messages) == 2
entry = loop.auto_compact._summaries.get("cli:test")
assert entry is not None
assert entry[0] == "User chatted about old things."
await loop.close_mcp()
@pytest.mark.asyncio
async def test_no_proactive_archive_when_active(self, tmp_path):
"""Recently active session should NOT be archived on idle tick."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "recent message")
loop.sessions.save(session)
await self._run_check_expired(loop)
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 1
await loop.close_mcp()
@pytest.mark.asyncio
async def test_no_duplicate_archive(self, tmp_path):
"""Should not archive the same session twice if already in progress."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
started = asyncio.Event()
block_forever = asyncio.Event()
async def _slow_archive(messages):
nonlocal archive_count
archive_count += 1
started.set()
await block_forever.wait()
return True
loop.consolidator.archive = _slow_archive
# First call starts archiving via callback
loop.auto_compact.check_expired(loop._schedule_background)
await started.wait()
assert archive_count == 1
# Second call should skip (key is in _archiving)
loop.auto_compact.check_expired(loop._schedule_background)
await asyncio.sleep(0.05)
assert archive_count == 1
# Clean up
block_forever.set()
await asyncio.sleep(0.1)
await loop.close_mcp()
@pytest.mark.asyncio
async def test_proactive_archive_error_does_not_block(self, tmp_path):
"""Proactive archive failure should be caught and not block future ticks."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _failing_archive(messages):
raise RuntimeError("LLM down")
loop.consolidator.archive = _failing_archive
# Should not raise
await self._run_check_expired(loop)
# Key should be removed from _archiving (finally block)
assert "cli:test" not in loop.auto_compact._archiving
await loop.close_mcp()
@pytest.mark.asyncio
async def test_proactive_archive_skips_empty_sessions(self, tmp_path):
"""Proactive archive should not call LLM for sessions with no un-consolidated messages."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_called = False
async def _fake_archive(messages):
nonlocal archive_called
archive_called = True
return True
loop.consolidator.archive = _fake_archive
await self._run_check_expired(loop)
assert not archive_called
await loop.close_mcp()
@pytest.mark.asyncio
async def test_no_reschedule_after_successful_archive(self, tmp_path):
"""Already-archived session should NOT be re-scheduled on subsequent ticks."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "old message")
session.add_message("assistant", "old response")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "Summary.",
}
# First tick: archives the session
await self._run_check_expired(loop)
assert archive_count == 1
# Second tick: should NOT re-schedule (updated_at is fresh after clear)
await self._run_check_expired(loop)
assert archive_count == 1 # Still 1, not re-scheduled
await loop.close_mcp()
@pytest.mark.asyncio
async def test_empty_skip_refreshes_updated_at_prevents_reschedule(self, tmp_path):
"""Empty session skip refreshes updated_at, preventing immediate re-scheduling."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return True
loop.consolidator.archive = _fake_archive
# First tick: skips (no messages), refreshes updated_at
await self._run_check_expired(loop)
assert archive_count == 0
# Second tick: should NOT re-schedule because updated_at is fresh
await self._run_check_expired(loop)
assert archive_count == 0
await loop.close_mcp()
@pytest.mark.asyncio
async def test_session_can_be_compacted_again_after_new_messages(self, tmp_path):
"""After successful compact + user sends new messages + idle again, should compact again."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "first conversation")
session.add_message("assistant", "first response")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
archive_count = 0
async def _fake_archive(messages):
nonlocal archive_count
archive_count += 1
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "Summary.",
}
# First compact cycle
await loop.auto_compact._archive("cli:test")
assert archive_count == 1
# User returns, sends new messages
msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="second topic")
await loop._process_message(msg)
# Simulate idle again
loop.sessions.invalidate("cli:test")
session2 = loop.sessions.get_or_create("cli:test")
session2.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session2)
# Second compact cycle should succeed
await loop.auto_compact._archive("cli:test")
assert archive_count == 2
await loop.close_mcp()
class TestSummaryPersistence:
"""Test that summary survives restart via session metadata."""
@pytest.mark.asyncio
async def test_summary_persisted_in_session_metadata(self, tmp_path):
"""After archive, _last_summary should be in session metadata."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "hello")
session.add_message("assistant", "hi there")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "User said hello.",
}
await loop.auto_compact._archive("cli:test")
# Summary should be persisted in session metadata
session_after = loop.sessions.get_or_create("cli:test")
meta = session_after.metadata.get("_last_summary")
assert meta is not None
assert meta["text"] == "User said hello."
assert "last_active" in meta
await loop.close_mcp()
@pytest.mark.asyncio
async def test_summary_recovered_after_restart(self, tmp_path):
"""Summary should be recovered from metadata when _summaries is empty (simulates restart)."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "hello")
session.add_message("assistant", "hi there")
last_active = datetime.now() - timedelta(minutes=20)
session.updated_at = last_active
loop.sessions.save(session)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "User said hello.",
}
# Archive
await loop.auto_compact._archive("cli:test")
# Simulate restart: clear in-memory state
loop.auto_compact._summaries.clear()
loop.sessions.invalidate("cli:test")
# prepare_session should recover summary from metadata
reloaded = loop.sessions.get_or_create("cli:test")
_, summary = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert summary is not None
assert "User said hello." in summary
assert "Inactive for" in summary
# Metadata should be cleaned up after consumption
assert "_last_summary" not in reloaded.metadata
await loop.close_mcp()
@pytest.mark.asyncio
async def test_metadata_cleanup_no_leak(self, tmp_path):
"""_last_summary should be removed from metadata after being consumed."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "hello")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "Summary.",
}
await loop.auto_compact._archive("cli:test")
# Clear in-memory to force metadata path
loop.auto_compact._summaries.clear()
loop.sessions.invalidate("cli:test")
reloaded = loop.sessions.get_or_create("cli:test")
# First call: consumes from metadata
_, summary = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert summary is not None
# Second call: no summary (already consumed)
_, summary2 = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert summary2 is None
assert "_last_summary" not in reloaded.metadata
await loop.close_mcp()
@pytest.mark.asyncio
async def test_metadata_cleanup_on_inmemory_path(self, tmp_path):
"""In-memory _summaries path should also clean up _last_summary from metadata."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
session.add_message("user", "hello")
session.updated_at = datetime.now() - timedelta(minutes=20)
loop.sessions.save(session)
async def _fake_archive(messages):
return True
loop.consolidator.archive = _fake_archive
loop.consolidator.get_last_history_entry = lambda: {
"cursor": 1, "timestamp": "2026-01-01 00:00", "content": "Summary.",
}
await loop.auto_compact._archive("cli:test")
# Both _summaries and metadata have the summary
assert "cli:test" in loop.auto_compact._summaries
loop.sessions.invalidate("cli:test")
reloaded = loop.sessions.get_or_create("cli:test")
assert "_last_summary" in reloaded.metadata
# In-memory path is taken (no restart)
_, summary = loop.auto_compact.prepare_session(reloaded, "cli:test")
assert summary is not None
# Metadata should also be cleaned up
assert "_last_summary" not in reloaded.metadata
await loop.close_mcp()
+3 -3
View File
@@ -278,7 +278,7 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path):
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
content, tools_used, messages = await loop._run_agent_loop( content, tools_used, messages, _, _ = await loop._run_agent_loop(
[{"role": "user", "content": "hi"}] [{"role": "user", "content": "hi"}]
) )
@@ -302,7 +302,7 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path):
) )
loop.tools.get_definitions = MagicMock(return_value=[]) loop.tools.get_definitions = MagicMock(return_value=[])
content, _, _ = await loop._run_agent_loop( content, _, _, _, _ = await loop._run_agent_loop(
[{"role": "user", "content": "hi"}] [{"role": "user", "content": "hi"}]
) )
@@ -344,7 +344,7 @@ async def test_agent_loop_no_hooks_backward_compat(tmp_path):
loop.tools.execute = AsyncMock(return_value="ok") loop.tools.execute = AsyncMock(return_value="ok")
loop.max_iterations = 2 loop.max_iterations = 2
content, tools_used, _ = await loop._run_agent_loop([]) content, tools_used, _, _, _ = await loop._run_agent_loop([])
assert content == ( assert content == (
"I reached the maximum number of tool call iterations (2) " "I reached the maximum number of tool call iterations (2) "
"without completing the task. You can try breaking the task into smaller steps." "without completing the task. You can try breaking the task into smaller steps."
+412 -3
View File
@@ -798,7 +798,7 @@ async def test_loop_max_iterations_message_stays_stable(tmp_path):
loop.tools.execute = AsyncMock(return_value="ok") loop.tools.execute = AsyncMock(return_value="ok")
loop.max_iterations = 2 loop.max_iterations = 2
final_content, _, _ = await loop._run_agent_loop([]) final_content, _, _, _, _ = await loop._run_agent_loop([])
assert final_content == ( assert final_content == (
"I reached the maximum number of tool call iterations (2) " "I reached the maximum number of tool call iterations (2) "
@@ -825,7 +825,7 @@ async def test_loop_stream_filter_handles_think_only_prefix_without_crashing(tmp
async def on_stream_end(*, resuming: bool = False) -> None: async def on_stream_end(*, resuming: bool = False) -> None:
endings.append(resuming) endings.append(resuming)
final_content, _, _ = await loop._run_agent_loop( final_content, _, _, _, _ = await loop._run_agent_loop(
[], [],
on_stream=on_stream, on_stream=on_stream,
on_stream_end=on_stream_end, on_stream_end=on_stream_end,
@@ -849,7 +849,7 @@ async def test_loop_retries_think_only_final_response(tmp_path):
loop.provider.chat_with_retry = chat_with_retry loop.provider.chat_with_retry = chat_with_retry
final_content, _, _ = await loop._run_agent_loop([]) final_content, _, _, _, _ = await loop._run_agent_loop([])
assert final_content == "Recovered answer" assert final_content == "Recovered answer"
assert call_count["n"] == 2 assert call_count["n"] == 2
@@ -999,3 +999,412 @@ async def test_runner_passes_cached_tokens_to_hook_context():
assert len(captured_usage) == 1 assert len(captured_usage) == 1
assert captured_usage[0]["cached_tokens"] == 150 assert captured_usage[0]["cached_tokens"] == 150
# ── Mid-turn injection tests ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_drain_injections_returns_empty_when_no_callback():
"""No injection_callback → empty list."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner
provider = MagicMock()
runner = AgentRunner(provider)
tools = MagicMock()
tools.get_definitions.return_value = []
spec = AgentRunSpec(
initial_messages=[], tools=tools, model="m",
max_iterations=1, max_tool_result_chars=1000,
injection_callback=None,
)
result = await runner._drain_injections(spec)
assert result == []
@pytest.mark.asyncio
async def test_drain_injections_extracts_content_from_inbound_messages():
"""Should extract .content from InboundMessage objects."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner, _MAX_INJECTIONS_PER_TURN
from nanobot.bus.events import InboundMessage
provider = MagicMock()
runner = AgentRunner(provider)
tools = MagicMock()
tools.get_definitions.return_value = []
msgs = [
InboundMessage(channel="cli", sender_id="u", chat_id="c", content="hello"),
InboundMessage(channel="cli", sender_id="u", chat_id="c", content="world"),
]
async def cb():
return msgs
spec = AgentRunSpec(
initial_messages=[], tools=tools, model="m",
max_iterations=1, max_tool_result_chars=1000,
injection_callback=cb,
)
result = await runner._drain_injections(spec)
assert result == ["hello", "world"]
@pytest.mark.asyncio
async def test_drain_injections_caps_at_max_and_logs_warning():
"""When more than _MAX_INJECTIONS_PER_TURN items, only the last N are kept."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner, _MAX_INJECTIONS_PER_TURN
from nanobot.bus.events import InboundMessage
provider = MagicMock()
runner = AgentRunner(provider)
tools = MagicMock()
tools.get_definitions.return_value = []
msgs = [
InboundMessage(channel="cli", sender_id="u", chat_id="c", content=f"msg{i}")
for i in range(_MAX_INJECTIONS_PER_TURN + 3)
]
async def cb():
return msgs
spec = AgentRunSpec(
initial_messages=[], tools=tools, model="m",
max_iterations=1, max_tool_result_chars=1000,
injection_callback=cb,
)
result = await runner._drain_injections(spec)
assert len(result) == _MAX_INJECTIONS_PER_TURN
# Should keep the LAST _MAX_INJECTIONS_PER_TURN items
assert result[0] == "msg3"
assert result[-1] == f"msg{_MAX_INJECTIONS_PER_TURN + 2}"
@pytest.mark.asyncio
async def test_drain_injections_skips_empty_content():
"""Messages with blank content should be filtered out."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner
from nanobot.bus.events import InboundMessage
provider = MagicMock()
runner = AgentRunner(provider)
tools = MagicMock()
tools.get_definitions.return_value = []
msgs = [
InboundMessage(channel="cli", sender_id="u", chat_id="c", content=""),
InboundMessage(channel="cli", sender_id="u", chat_id="c", content=" "),
InboundMessage(channel="cli", sender_id="u", chat_id="c", content="valid"),
]
async def cb():
return msgs
spec = AgentRunSpec(
initial_messages=[], tools=tools, model="m",
max_iterations=1, max_tool_result_chars=1000,
injection_callback=cb,
)
result = await runner._drain_injections(spec)
assert result == ["valid"]
@pytest.mark.asyncio
async def test_drain_injections_handles_callback_exception():
"""If the callback raises, return empty list (error is logged)."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner
provider = MagicMock()
runner = AgentRunner(provider)
tools = MagicMock()
tools.get_definitions.return_value = []
async def cb():
raise RuntimeError("boom")
spec = AgentRunSpec(
initial_messages=[], tools=tools, model="m",
max_iterations=1, max_tool_result_chars=1000,
injection_callback=cb,
)
result = await runner._drain_injections(spec)
assert result == []
@pytest.mark.asyncio
async def test_checkpoint1_injects_after_tool_execution():
"""Follow-up messages are injected after tool execution, before next LLM call."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner
from nanobot.bus.events import InboundMessage
provider = MagicMock()
call_count = {"n": 0}
captured_messages = []
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
captured_messages.append(list(messages))
if call_count["n"] == 1:
return LLMResponse(
content="using tool",
tool_calls=[ToolCallRequest(id="c1", name="read_file", arguments={"path": "x"})],
usage={},
)
return LLMResponse(content="final answer", tool_calls=[], usage={})
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
tools.execute = AsyncMock(return_value="file content")
injection_queue = asyncio.Queue()
async def inject_cb():
items = []
while not injection_queue.empty():
items.append(await injection_queue.get())
return items
# Put a follow-up message in the queue before the run starts
await injection_queue.put(
InboundMessage(channel="cli", sender_id="u", chat_id="c", content="follow-up question")
)
runner = AgentRunner(provider)
result = await runner.run(AgentRunSpec(
initial_messages=[{"role": "user", "content": "hello"}],
tools=tools,
model="test-model",
max_iterations=5,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
injection_callback=inject_cb,
))
assert result.had_injections is True
assert result.final_content == "final answer"
# The second call should have the injected user message
assert call_count["n"] == 2
last_messages = captured_messages[-1]
injected = [m for m in last_messages if m.get("role") == "user" and m.get("content") == "follow-up question"]
assert len(injected) == 1
@pytest.mark.asyncio
async def test_checkpoint2_injects_after_final_response_with_resuming_stream():
"""After final response, if injections exist, stream_end should get resuming=True."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.bus.events import InboundMessage
provider = MagicMock()
call_count = {"n": 0}
stream_end_calls = []
class TrackingHook(AgentHook):
def wants_streaming(self) -> bool:
return True
async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None:
stream_end_calls.append(resuming)
def finalize_content(self, context: AgentHookContext, content: str | None) -> str | None:
return content
async def chat_stream_with_retry(*, messages, on_content_delta=None, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
return LLMResponse(content="first answer", tool_calls=[], usage={})
return LLMResponse(content="second answer", tool_calls=[], usage={})
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
injection_queue = asyncio.Queue()
async def inject_cb():
items = []
while not injection_queue.empty():
items.append(await injection_queue.get())
return items
# Inject a follow-up that arrives during the first response
await injection_queue.put(
InboundMessage(channel="cli", sender_id="u", chat_id="c", content="quick follow-up")
)
runner = AgentRunner(provider)
result = await runner.run(AgentRunSpec(
initial_messages=[{"role": "user", "content": "hello"}],
tools=tools,
model="test-model",
max_iterations=5,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=TrackingHook(),
injection_callback=inject_cb,
))
assert result.had_injections is True
assert result.final_content == "second answer"
assert call_count["n"] == 2
# First stream_end should have resuming=True (because injections found)
assert stream_end_calls[0] is True
# Second (final) stream_end should have resuming=False
assert stream_end_calls[-1] is False
@pytest.mark.asyncio
async def test_injection_cycles_capped_at_max():
"""Injection cycles should be capped at _MAX_INJECTION_CYCLES."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner, _MAX_INJECTION_CYCLES
from nanobot.bus.events import InboundMessage
provider = MagicMock()
call_count = {"n": 0}
async def chat_with_retry(*, messages, **kwargs):
call_count["n"] += 1
return LLMResponse(content=f"answer-{call_count['n']}", tool_calls=[], usage={})
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
drain_count = {"n": 0}
async def inject_cb():
drain_count["n"] += 1
# Only inject for the first _MAX_INJECTION_CYCLES drains
if drain_count["n"] <= _MAX_INJECTION_CYCLES:
return [InboundMessage(channel="cli", sender_id="u", chat_id="c", content=f"msg-{drain_count['n']}")]
return []
runner = AgentRunner(provider)
result = await runner.run(AgentRunSpec(
initial_messages=[{"role": "user", "content": "start"}],
tools=tools,
model="test-model",
max_iterations=20,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
injection_callback=inject_cb,
))
assert result.had_injections is True
# Should be capped: _MAX_INJECTION_CYCLES injection rounds + 1 final round
assert call_count["n"] == _MAX_INJECTION_CYCLES + 1
@pytest.mark.asyncio
async def test_no_injections_flag_is_false_by_default():
"""had_injections should be False when no injection callback or no messages."""
from nanobot.agent.runner import AgentRunSpec, AgentRunner
provider = MagicMock()
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
runner = AgentRunner(provider)
result = await runner.run(AgentRunSpec(
initial_messages=[{"role": "user", "content": "hi"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
assert result.had_injections is False
@pytest.mark.asyncio
async def test_pending_queue_cleanup_on_dispatch(tmp_path):
"""_pending_queues should be cleaned up after _dispatch completes."""
loop = _make_loop(tmp_path)
async def chat_with_retry(**kwargs):
return LLMResponse(content="done", tool_calls=[], usage={})
loop.provider.chat_with_retry = chat_with_retry
from nanobot.bus.events import InboundMessage
msg = InboundMessage(channel="cli", sender_id="u", chat_id="c", content="hello")
# The queue should not exist before dispatch
assert msg.session_key not in loop._pending_queues
await loop._dispatch(msg)
# The queue should be cleaned up after dispatch
assert msg.session_key not in loop._pending_queues
@pytest.mark.asyncio
async def test_followup_routed_to_pending_queue(tmp_path):
"""When a session has an active dispatch, follow-up messages go to pending queue."""
from nanobot.bus.events import InboundMessage
loop = _make_loop(tmp_path)
# Simulate an active dispatch by manually adding a pending queue
pending = asyncio.Queue(maxsize=20)
loop._pending_queues["cli:c"] = pending
msg = InboundMessage(channel="cli", sender_id="u", chat_id="c", content="follow-up")
# Directly test the routing logic from run() — if session_key is in
# _pending_queues, the message should be put into the queue.
assert msg.session_key in loop._pending_queues
loop._pending_queues[msg.session_key].put_nowait(msg)
assert not pending.empty()
queued_msg = pending.get_nowait()
assert queued_msg.content == "follow-up"
@pytest.mark.asyncio
async def test_dispatch_republishes_leftover_queue_messages(tmp_path):
"""Messages left in the pending queue after _dispatch are re-published to the bus.
This tests the finally-block cleanup that prevents message loss when
the runner exits early (e.g., max_iterations, tool_error) with messages
still in the queue.
"""
from nanobot.bus.events import InboundMessage
loop = _make_loop(tmp_path)
bus = loop.bus
# Simulate a completed dispatch by manually registering a queue
# with leftover messages, then running the cleanup logic directly.
pending = asyncio.Queue(maxsize=20)
session_key = "cli:c"
loop._pending_queues[session_key] = pending
pending.put_nowait(InboundMessage(channel="cli", sender_id="u", chat_id="c", content="leftover-1"))
pending.put_nowait(InboundMessage(channel="cli", sender_id="u", chat_id="c", content="leftover-2"))
# Execute the cleanup logic from the finally block
queue = loop._pending_queues.pop(session_key, None)
assert queue is not None
leftover = 0
while True:
try:
item = queue.get_nowait()
except asyncio.QueueEmpty:
break
await bus.publish_inbound(item)
leftover += 1
assert leftover == 2
# Verify the messages are now on the bus
msgs = []
while not bus.inbound.empty():
msgs.append(await asyncio.wait_for(bus.consume_inbound(), timeout=0.5))
contents = [m.content for m in msgs]
assert "leftover-1" in contents
assert "leftover-2" in contents
+45 -7
View File
@@ -52,6 +52,37 @@ class TestToolHintKnownTools:
assert result.startswith("$ ") assert result.startswith("$ ")
assert len(result) <= 50 # reasonable limit assert len(result) <= 50 # reasonable limit
def test_exec_abbreviates_paths_in_command(self):
"""Windows paths in exec commands should be folded, not blindly truncated."""
cmd = "cd D:\\Documents\\GitHub\\nanobot\\.worktree\\tomain\\nanobot && git diff origin/main...pr-2706 --name-only 2>&1"
result = _hint([_tc("exec", {"command": cmd})])
assert "\u2026/" in result # path should be folded with …/
assert "worktree" not in result # middle segments should be collapsed
def test_exec_abbreviates_linux_paths(self):
"""Unix absolute paths in exec commands should be folded."""
cmd = "cd /home/user/projects/nanobot/.worktree/tomain && make build"
result = _hint([_tc("exec", {"command": cmd})])
assert "\u2026/" in result
assert "projects" not in result
def test_exec_abbreviates_home_paths(self):
"""~/ paths in exec commands should be folded."""
cmd = "cd ~/projects/nanobot/workspace && pytest tests/"
result = _hint([_tc("exec", {"command": cmd})])
assert "\u2026/" in result
def test_exec_short_command_unchanged(self):
result = _hint([_tc("exec", {"command": "npm install typescript"})])
assert result == "$ npm install typescript"
def test_exec_chained_commands_truncated_not_mid_path(self):
"""Long chained commands should truncate preserving abbreviated paths."""
cmd = "cd D:\\Documents\\GitHub\\project && npm run build && npm test"
result = _hint([_tc("exec", {"command": cmd})])
assert "\u2026/" in result # path folded
assert "npm" in result # chained command still visible
def test_web_search(self): def test_web_search(self):
result = _hint([_tc("web_search", {"query": "Claude 4 vs GPT-4"})]) result = _hint([_tc("web_search", {"query": "Claude 4 vs GPT-4"})])
assert result == 'search "Claude 4 vs GPT-4"' assert result == 'search "Claude 4 vs GPT-4"'
@@ -105,22 +136,30 @@ class TestToolHintFolding:
result = _hint(calls) result = _hint(calls)
assert "\u00d7" not in result assert "\u00d7" not in result
def test_two_consecutive_same_folded(self): def test_two_consecutive_different_args_not_folded(self):
calls = [ calls = [
_tc("grep", {"pattern": "*.py"}), _tc("grep", {"pattern": "*.py"}),
_tc("grep", {"pattern": "*.ts"}), _tc("grep", {"pattern": "*.ts"}),
] ]
result = _hint(calls) result = _hint(calls)
assert "\u00d7" not in result
def test_two_consecutive_same_args_folded(self):
calls = [
_tc("grep", {"pattern": "TODO"}),
_tc("grep", {"pattern": "TODO"}),
]
result = _hint(calls)
assert "\u00d7 2" in result assert "\u00d7 2" in result
def test_three_consecutive_same_folded(self): def test_three_consecutive_different_args_not_folded(self):
calls = [ calls = [
_tc("read_file", {"path": "a.py"}), _tc("read_file", {"path": "a.py"}),
_tc("read_file", {"path": "b.py"}), _tc("read_file", {"path": "b.py"}),
_tc("read_file", {"path": "c.py"}), _tc("read_file", {"path": "c.py"}),
] ]
result = _hint(calls) result = _hint(calls)
assert "\u00d7 3" in result assert "\u00d7" not in result
def test_different_tools_not_folded(self): def test_different_tools_not_folded(self):
calls = [ calls = [
@@ -187,7 +226,7 @@ class TestToolHintMixedFolding:
"""G4: Mixed folding groups with interleaved same-tool segments.""" """G4: Mixed folding groups with interleaved same-tool segments."""
def test_read_read_grep_grep_read(self): def test_read_read_grep_grep_read(self):
"""read×2, grep×2, read — should produce two separate groups.""" """All different args — each hint listed separately."""
calls = [ calls = [
_tc("read_file", {"path": "a.py"}), _tc("read_file", {"path": "a.py"}),
_tc("read_file", {"path": "b.py"}), _tc("read_file", {"path": "b.py"}),
@@ -196,7 +235,6 @@ class TestToolHintMixedFolding:
_tc("read_file", {"path": "c.py"}), _tc("read_file", {"path": "c.py"}),
] ]
result = _hint(calls) result = _hint(calls)
assert "\u00d7 2" in result assert "\u00d7" not in result
# Should have 3 groups: read×2, grep×2, read
parts = result.split(", ") parts = result.split(", ")
assert len(parts) == 3 assert len(parts) == 5
+48
View File
@@ -0,0 +1,48 @@
"""Tests for Feishu/Lark domain configuration."""
from unittest.mock import MagicMock
import pytest
from nanobot.bus.queue import MessageBus
from nanobot.channels.feishu import FeishuChannel, FeishuConfig
def _make_channel(domain: str = "feishu") -> FeishuChannel:
config = FeishuConfig(
enabled=True,
app_id="cli_test",
app_secret="secret",
allow_from=["*"],
domain=domain,
)
ch = FeishuChannel(config, MessageBus())
ch._client = MagicMock()
ch._loop = None
return ch
class TestFeishuConfigDomain:
def test_domain_default_is_feishu(self):
config = FeishuConfig()
assert config.domain == "feishu"
def test_domain_accepts_lark(self):
config = FeishuConfig(domain="lark")
assert config.domain == "lark"
def test_domain_accepts_feishu(self):
config = FeishuConfig(domain="feishu")
assert config.domain == "feishu"
def test_default_config_includes_domain(self):
default_cfg = FeishuChannel.default_config()
assert "domain" in default_cfg
assert default_cfg["domain"] == "feishu"
def test_channel_persists_domain_from_config(self):
ch = _make_channel(domain="lark")
assert ch.config.domain == "lark"
def test_channel_persists_feishu_domain_from_config(self):
ch = _make_channel(domain="feishu")
assert ch.config.domain == "feishu"
+190
View File
@@ -5,6 +5,7 @@ from unittest.mock import MagicMock
import pytest import pytest
from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.channels.feishu import FeishuChannel, FeishuConfig, _FeishuStreamBuf from nanobot.channels.feishu import FeishuChannel, FeishuConfig, _FeishuStreamBuf
@@ -203,6 +204,55 @@ class TestSendDelta:
ch._client.cardkit.v1.card_element.content.assert_not_called() ch._client.cardkit.v1.card_element.content.assert_not_called()
ch._client.im.v1.message.create.assert_called_once() ch._client.im.v1.message.create.assert_called_once()
@pytest.mark.asyncio
async def test_stream_end_resuming_keeps_buffer(self):
"""_resuming=True flushes text to card but keeps the buffer for the next segment."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Partial answer", card_id="card_1", sequence=2, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True, "_resuming": True})
assert "oc_chat1" in ch._stream_bufs
buf = ch._stream_bufs["oc_chat1"]
assert buf.card_id == "card_1"
assert buf.sequence == 3
ch._client.cardkit.v1.card_element.content.assert_called_once()
ch._client.cardkit.v1.card.settings.assert_not_called()
@pytest.mark.asyncio
async def test_stream_end_resuming_then_final_end(self):
"""Full multi-segment flow: resuming mid-turn, then final end closes the card."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Seg1", card_id="card_1", sequence=1, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
ch._client.cardkit.v1.card.settings.return_value = _mock_content_response()
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True, "_resuming": True})
assert "oc_chat1" in ch._stream_bufs
ch._stream_bufs["oc_chat1"].text += " Seg2"
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True})
assert "oc_chat1" not in ch._stream_bufs
ch._client.cardkit.v1.card.settings.assert_called_once()
@pytest.mark.asyncio
async def test_stream_end_resuming_no_card_is_noop(self):
"""_resuming with no card_id (card creation failed) is a safe no-op."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="text", card_id=None, sequence=0, last_edit=0.0,
)
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True, "_resuming": True})
assert "oc_chat1" in ch._stream_bufs
ch._client.cardkit.v1.card_element.content.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_stream_end_without_buf_is_noop(self): async def test_stream_end_without_buf_is_noop(self):
ch = _make_channel() ch = _make_channel()
@@ -239,6 +289,146 @@ class TestSendDelta:
assert buf.sequence == 7 assert buf.sequence == 7
class TestToolHintInlineStreaming:
"""Tool hint messages should be inlined into active streaming cards."""
@pytest.mark.asyncio
async def test_tool_hint_inlined_when_stream_active(self):
"""With an active streaming buffer, tool hint appends to the card."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Partial answer", card_id="card_1", sequence=2, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
msg = OutboundMessage(
channel="feishu", chat_id="oc_chat1",
content='web_fetch("https://example.com")',
metadata={"_tool_hint": True},
)
await ch.send(msg)
buf = ch._stream_bufs["oc_chat1"]
assert '🔧 web_fetch("https://example.com")' in buf.text
assert buf.sequence == 3
ch._client.cardkit.v1.card_element.content.assert_called_once()
ch._client.im.v1.message.create.assert_not_called()
@pytest.mark.asyncio
async def test_tool_hint_preserved_on_next_delta(self):
"""When new delta arrives, the tool hint is kept as permanent content and delta appends after it."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Partial answer\n\n🔧 web_fetch(\"url\")\n\n",
card_id="card_1", sequence=3, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
await ch.send_delta("oc_chat1", " continued")
buf = ch._stream_bufs["oc_chat1"]
assert "Partial answer" in buf.text
assert "🔧 web_fetch" in buf.text
assert buf.text.endswith(" continued")
@pytest.mark.asyncio
async def test_tool_hint_fallback_when_no_stream(self):
"""Without an active buffer, tool hint falls back to a standalone card."""
ch = _make_channel()
ch._client.im.v1.message.create.return_value = _mock_send_response("om_hint")
msg = OutboundMessage(
channel="feishu", chat_id="oc_chat1",
content='read_file("path")',
metadata={"_tool_hint": True},
)
await ch.send(msg)
assert "oc_chat1" not in ch._stream_bufs
ch._client.im.v1.message.create.assert_called_once()
@pytest.mark.asyncio
async def test_consecutive_tool_hints_append(self):
"""When multiple tool hints arrive consecutively, each appends to the card."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Partial answer", card_id="card_1", sequence=2, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
msg1 = OutboundMessage(
channel="feishu", chat_id="oc_chat1",
content='$ cd /project', metadata={"_tool_hint": True},
)
await ch.send(msg1)
msg2 = OutboundMessage(
channel="feishu", chat_id="oc_chat1",
content='$ git status', metadata={"_tool_hint": True},
)
await ch.send(msg2)
buf = ch._stream_bufs["oc_chat1"]
assert "$ cd /project" in buf.text
assert "$ git status" in buf.text
assert buf.text.startswith("Partial answer")
assert "🔧 $ cd /project" in buf.text
assert "🔧 $ git status" in buf.text
@pytest.mark.asyncio
async def test_tool_hint_preserved_on_resuming_flush(self):
"""When _resuming flushes the buffer, tool hint is kept as permanent content."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Partial answer\n\n🔧 $ cd /project\n\n",
card_id="card_1", sequence=2, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True, "_resuming": True})
buf = ch._stream_bufs["oc_chat1"]
assert "Partial answer" in buf.text
assert "🔧 $ cd /project" in buf.text
@pytest.mark.asyncio
async def test_tool_hint_preserved_on_final_stream_end(self):
"""When final _stream_end closes the card, tool hint is kept in the final text."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Final content\n\n🔧 web_fetch(\"url\")\n\n",
card_id="card_1", sequence=3, last_edit=0.0,
)
ch._client.cardkit.v1.card_element.content.return_value = _mock_content_response()
ch._client.cardkit.v1.card.settings.return_value = _mock_content_response()
await ch.send_delta("oc_chat1", "", metadata={"_stream_end": True})
assert "oc_chat1" not in ch._stream_bufs
update_call = ch._client.cardkit.v1.card_element.content.call_args[0][0]
assert "🔧" in update_call.body.content
@pytest.mark.asyncio
async def test_empty_tool_hint_is_noop(self):
"""Empty or whitespace-only tool hint content is silently ignored."""
ch = _make_channel()
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
text="Partial answer", card_id="card_1", sequence=2, last_edit=0.0,
)
for content in ("", " ", "\t\n"):
msg = OutboundMessage(
channel="feishu", chat_id="oc_chat1",
content=content, metadata={"_tool_hint": True},
)
await ch.send(msg)
buf = ch._stream_bufs["oc_chat1"]
assert buf.text == "Partial answer"
assert buf.sequence == 2
ch._client.cardkit.v1.card_element.content.assert_not_called()
class TestSendMessageReturnsId: class TestSendMessageReturnsId:
def test_returns_message_id_on_success(self): def test_returns_message_id_on_success(self):
ch = _make_channel() ch = _make_channel()
+304
View File
@@ -0,0 +1,304 @@
"""Tests for QQ channel media support: helpers, send, inbound, and upload."""
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
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.events import OutboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.channels.qq import (
QQ_FILE_TYPE_FILE,
QQ_FILE_TYPE_IMAGE,
QQChannel,
QQConfig,
_guess_send_file_type,
_is_image_name,
_sanitize_filename,
)
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 _FakeHttp:
"""Fake _http for _post_base64file tests."""
def __init__(self, return_value: dict | None = None) -> None:
self.return_value = return_value or {}
self.calls: list[tuple] = []
async def request(self, route, **kwargs):
self.calls.append((route, kwargs))
return self.return_value
class _FakeClient:
def __init__(self, http_return: dict | None = None) -> None:
self.api = _FakeApi()
self.api._http = _FakeHttp(http_return)
# ── Helper function tests (pure, no async) ──────────────────────────
def test_sanitize_filename_strips_path_traversal() -> None:
assert _sanitize_filename("../../etc/passwd") == "passwd"
def test_sanitize_filename_keeps_chinese_chars() -> None:
assert _sanitize_filename("文件(1.jpg") == "文件(1.jpg"
def test_sanitize_filename_strips_unsafe_chars() -> None:
result = _sanitize_filename('file<>:"|?*.txt')
# All unsafe chars replaced with "_", but * is replaced too
assert result.startswith("file")
assert result.endswith(".txt")
assert "<" not in result
assert ">" not in result
assert '"' not in result
assert "|" not in result
assert "?" not in result
def test_sanitize_filename_empty_input() -> None:
assert _sanitize_filename("") == ""
def test_is_image_name_with_known_extensions() -> None:
for ext in (".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp", ".tif", ".tiff", ".ico", ".svg"):
assert _is_image_name(f"photo{ext}") is True
def test_is_image_name_with_unknown_extension() -> None:
for ext in (".pdf", ".txt", ".mp3", ".mp4"):
assert _is_image_name(f"doc{ext}") is False
def test_guess_send_file_type_image() -> None:
assert _guess_send_file_type("photo.png") == QQ_FILE_TYPE_IMAGE
assert _guess_send_file_type("pic.jpg") == QQ_FILE_TYPE_IMAGE
def test_guess_send_file_type_file() -> None:
assert _guess_send_file_type("doc.pdf") == QQ_FILE_TYPE_FILE
def test_guess_send_file_type_by_mime() -> None:
# A filename with no known extension but whose mime type is image/*
assert _guess_send_file_type("photo.xyz_image_test") == QQ_FILE_TYPE_FILE
# ── send() exception handling ───────────────────────────────────────
@pytest.mark.asyncio
async def test_send_exception_caught_not_raised() -> None:
"""Exceptions inside send() must not propagate."""
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
channel._client = _FakeClient()
with patch.object(channel, "_send_text_only", new_callable=AsyncMock, side_effect=RuntimeError("boom")):
await channel.send(
OutboundMessage(channel="qq", chat_id="user1", content="hello")
)
# No exception raised — test passes if we get here.
@pytest.mark.asyncio
async def test_send_media_then_text() -> None:
"""Media is sent before text when both are present."""
import tempfile
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
channel._client = _FakeClient()
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(b"\x89PNG\r\n")
tmp = f.name
try:
with patch.object(channel, "_post_base64file", new_callable=AsyncMock, return_value={"file_info": "1"}) as mock_upload:
await channel.send(
OutboundMessage(
channel="qq",
chat_id="user1",
content="text after image",
media=[tmp],
metadata={"message_id": "m1"},
)
)
assert mock_upload.called
# Text should have been sent via c2c (default chat type)
text_calls = [c for c in channel._client.api.c2c_calls if c.get("msg_type") == 0]
assert len(text_calls) >= 1
assert text_calls[-1]["content"] == "text after image"
finally:
import os
os.unlink(tmp)
@pytest.mark.asyncio
async def test_send_media_failure_falls_back_to_text() -> None:
"""When _send_media returns False, a failure notice is appended."""
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
channel._client = _FakeClient()
with patch.object(channel, "_send_media", new_callable=AsyncMock, return_value=False):
await channel.send(
OutboundMessage(
channel="qq",
chat_id="user1",
content="hello",
media=["https://example.com/bad.png"],
metadata={"message_id": "m1"},
)
)
# Should have the failure text among the c2c calls
failure_calls = [c for c in channel._client.api.c2c_calls if "Attachment send failed" in c.get("content", "")]
assert len(failure_calls) == 1
assert "bad.png" in failure_calls[0]["content"]
# ── _on_message() exception handling ────────────────────────────────
@pytest.mark.asyncio
async def test_on_message_exception_caught_not_raised() -> None:
"""Missing required attributes should not crash _on_message."""
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
channel._client = _FakeClient()
# Construct a message-like object that lacks 'author' — triggers AttributeError
bad_data = SimpleNamespace(id="x1", content="hi")
# Should not raise
await channel._on_message(bad_data, is_group=False)
@pytest.mark.asyncio
async def test_on_message_with_attachments() -> None:
"""Messages with attachments produce media_paths and formatted content."""
import tempfile
channel = QQChannel(QQConfig(app_id="app", secret="secret", allow_from=["*"]), MessageBus())
channel._client = _FakeClient()
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(b"\x89PNG\r\n")
saved_path = f.name
att = SimpleNamespace(url="", filename="screenshot.png", content_type="image/png")
# Patch _download_to_media_dir_chunked to return the temp file path
async def fake_download(url, filename_hint=""):
return saved_path
try:
with patch.object(channel, "_download_to_media_dir_chunked", side_effect=fake_download):
data = SimpleNamespace(
id="att1",
content="look at this",
author=SimpleNamespace(user_openid="u1"),
attachments=[att],
)
await channel._on_message(data, is_group=False)
msg = await channel.bus.consume_inbound()
assert "look at this" in msg.content
assert "screenshot.png" in msg.content
assert "Received files:" in msg.content
assert len(msg.media) == 1
assert msg.media[0] == saved_path
finally:
import os
os.unlink(saved_path)
# ── _post_base64file() ─────────────────────────────────────────────
@pytest.mark.asyncio
async def test_post_base64file_omits_file_name_for_images() -> None:
"""file_type=1 (image) → payload must not contain file_name."""
channel = QQChannel(QQConfig(app_id="app", secret="secret"), MessageBus())
channel._client = _FakeClient(http_return={"file_info": "img_abc"})
await channel._post_base64file(
chat_id="user1",
is_group=False,
file_type=QQ_FILE_TYPE_IMAGE,
file_data="ZmFrZQ==",
file_name="photo.png",
)
http = channel._client.api._http
assert len(http.calls) == 1
payload = http.calls[0][1]["json"]
assert "file_name" not in payload
assert payload["file_type"] == QQ_FILE_TYPE_IMAGE
@pytest.mark.asyncio
async def test_post_base64file_includes_file_name_for_files() -> None:
"""file_type=4 (file) → payload must contain file_name."""
channel = QQChannel(QQConfig(app_id="app", secret="secret"), MessageBus())
channel._client = _FakeClient(http_return={"file_info": "file_abc"})
await channel._post_base64file(
chat_id="user1",
is_group=False,
file_type=QQ_FILE_TYPE_FILE,
file_data="ZmFrZQ==",
file_name="report.pdf",
)
http = channel._client.api._http
assert len(http.calls) == 1
payload = http.calls[0][1]["json"]
assert payload["file_name"] == "report.pdf"
assert payload["file_type"] == QQ_FILE_TYPE_FILE
@pytest.mark.asyncio
async def test_post_base64file_filters_response_to_file_info() -> None:
"""Response with file_info + extra fields must be filtered to only file_info."""
channel = QQChannel(QQConfig(app_id="app", secret="secret"), MessageBus())
channel._client = _FakeClient(http_return={
"file_info": "fi_123",
"file_uuid": "uuid_xxx",
"ttl": 3600,
})
result = await channel._post_base64file(
chat_id="user1",
is_group=False,
file_type=QQ_FILE_TYPE_FILE,
file_data="ZmFrZQ==",
file_name="doc.pdf",
)
assert result == {"file_info": "fi_123"}
assert "file_uuid" not in result
assert "ttl" not in result
+598
View File
@@ -0,0 +1,598 @@
"""Unit and lightweight integration tests for the WebSocket channel."""
import asyncio
import functools
import json
import time
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
import websockets
from websockets.exceptions import ConnectionClosed
from websockets.frames import Close
from nanobot.bus.events import OutboundMessage
from nanobot.channels.websocket import (
WebSocketChannel,
WebSocketConfig,
_issue_route_secret_matches,
_normalize_config_path,
_normalize_http_path,
_parse_inbound_payload,
_parse_query,
_parse_request_path,
)
# -- Shared helpers (aligned with test_websocket_integration.py) ---------------
_PORT = 29876
def _ch(bus: Any, **kw: Any) -> WebSocketChannel:
cfg: dict[str, Any] = {
"enabled": True,
"allowFrom": ["*"],
"host": "127.0.0.1",
"port": _PORT,
"path": "/ws",
"websocketRequiresToken": False,
}
cfg.update(kw)
return WebSocketChannel(cfg, bus)
@pytest.fixture()
def bus() -> MagicMock:
b = MagicMock()
b.publish_inbound = AsyncMock()
return b
async def _http_get(url: str, headers: dict[str, str] | None = None) -> httpx.Response:
"""Run GET in a thread to avoid blocking the asyncio loop shared with websockets."""
return await asyncio.to_thread(
functools.partial(httpx.get, url, headers=headers or {}, timeout=5.0)
)
def test_normalize_http_path_strips_trailing_slash_except_root() -> None:
assert _normalize_http_path("/chat/") == "/chat"
assert _normalize_http_path("/chat?x=1") == "/chat"
assert _normalize_http_path("/") == "/"
def test_parse_request_path_matches_normalize_and_query() -> None:
path, query = _parse_request_path("/ws/?token=secret&client_id=u1")
assert path == _normalize_http_path("/ws/?token=secret&client_id=u1")
assert query == _parse_query("/ws/?token=secret&client_id=u1")
def test_normalize_config_path_matches_request() -> None:
assert _normalize_config_path("/ws/") == "/ws"
assert _normalize_config_path("/") == "/"
def test_parse_query_extracts_token_and_client_id() -> None:
query = _parse_query("/?token=secret&client_id=u1")
assert query.get("token") == ["secret"]
assert query.get("client_id") == ["u1"]
@pytest.mark.parametrize(
("raw", "expected"),
[
("plain", "plain"),
('{"content": "hi"}', "hi"),
('{"text": "there"}', "there"),
('{"message": "x"}', "x"),
(" ", None),
("{}", None),
],
)
def test_parse_inbound_payload(raw: str, expected: str | None) -> None:
assert _parse_inbound_payload(raw) == expected
def test_parse_inbound_invalid_json_falls_back_to_raw_string() -> None:
assert _parse_inbound_payload("{not json") == "{not json"
@pytest.mark.parametrize(
("raw", "expected"),
[
('{"content": ""}', None), # empty string content
('{"content": 123}', None), # non-string content
('{"content": " "}', None), # whitespace-only content
('["hello"]', '["hello"]'), # JSON array: not a dict, treated as plain text
('{"unknown_key": "val"}', None), # unrecognized key
('{"content": null}', None), # null content
],
)
def test_parse_inbound_payload_edge_cases(raw: str, expected: str | None) -> None:
assert _parse_inbound_payload(raw) == expected
def test_web_socket_config_path_must_start_with_slash() -> None:
with pytest.raises(ValueError, match='path must start with "/"'):
WebSocketConfig(path="bad")
def test_ssl_context_requires_both_cert_and_key_files() -> None:
bus = MagicMock()
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"], "sslCertfile": "/tmp/c.pem", "sslKeyfile": ""},
bus,
)
with pytest.raises(ValueError, match="ssl_certfile and ssl_keyfile"):
channel._build_ssl_context()
def test_default_config_includes_safe_bind_and_streaming() -> None:
defaults = WebSocketChannel.default_config()
assert defaults["enabled"] is False
assert defaults["host"] == "127.0.0.1"
assert defaults["streaming"] is True
assert defaults["allowFrom"] == ["*"]
assert defaults.get("tokenIssuePath", "") == ""
def test_token_issue_path_must_differ_from_websocket_path() -> None:
with pytest.raises(ValueError, match="token_issue_path must differ"):
WebSocketConfig(path="/ws", token_issue_path="/ws")
def test_issue_route_secret_matches_bearer_and_header() -> None:
from websockets.datastructures import Headers
secret = "my-secret"
bearer_headers = Headers([("Authorization", "Bearer my-secret")])
assert _issue_route_secret_matches(bearer_headers, secret) is True
x_headers = Headers([("X-Nanobot-Auth", "my-secret")])
assert _issue_route_secret_matches(x_headers, secret) is True
wrong = Headers([("Authorization", "Bearer other")])
assert _issue_route_secret_matches(wrong, secret) is False
def test_issue_route_secret_matches_empty_secret() -> None:
from websockets.datastructures import Headers
# Empty secret always returns True regardless of headers
assert _issue_route_secret_matches(Headers([]), "") is True
assert _issue_route_secret_matches(Headers([("Authorization", "Bearer anything")]), "") is True
@pytest.mark.asyncio
async def test_send_delivers_json_message_with_media_and_reply() -> None:
bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
mock_ws = AsyncMock()
channel._connections["chat-1"] = mock_ws
msg = OutboundMessage(
channel="websocket",
chat_id="chat-1",
content="hello",
reply_to="m1",
media=["/tmp/a.png"],
)
await channel.send(msg)
mock_ws.send.assert_awaited_once()
payload = json.loads(mock_ws.send.call_args[0][0])
assert payload["event"] == "message"
assert payload["text"] == "hello"
assert payload["reply_to"] == "m1"
assert payload["media"] == ["/tmp/a.png"]
@pytest.mark.asyncio
async def test_send_missing_connection_is_noop_without_error() -> None:
bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
msg = OutboundMessage(channel="websocket", chat_id="missing", content="x")
await channel.send(msg)
@pytest.mark.asyncio
async def test_send_removes_connection_on_connection_closed() -> None:
bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
mock_ws = AsyncMock()
mock_ws.send.side_effect = ConnectionClosed(Close(1006, ""), Close(1006, ""), True)
channel._connections["chat-1"] = mock_ws
msg = OutboundMessage(channel="websocket", chat_id="chat-1", content="hello")
await channel.send(msg)
assert "chat-1" not in channel._connections
@pytest.mark.asyncio
async def test_send_delta_removes_connection_on_connection_closed() -> None:
bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus)
mock_ws = AsyncMock()
mock_ws.send.side_effect = ConnectionClosed(Close(1006, ""), Close(1006, ""), True)
channel._connections["chat-1"] = mock_ws
await channel.send_delta("chat-1", "chunk", {"_stream_delta": True, "_stream_id": "s1"})
assert "chat-1" not in channel._connections
@pytest.mark.asyncio
async def test_send_delta_emits_delta_and_stream_end() -> None:
bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus)
mock_ws = AsyncMock()
channel._connections["chat-1"] = mock_ws
await channel.send_delta("chat-1", "part", {"_stream_delta": True, "_stream_id": "sid"})
await channel.send_delta("chat-1", "", {"_stream_end": True, "_stream_id": "sid"})
assert mock_ws.send.await_count == 2
first = json.loads(mock_ws.send.call_args_list[0][0][0])
second = json.loads(mock_ws.send.call_args_list[1][0][0])
assert first["event"] == "delta"
assert first["text"] == "part"
assert first["stream_id"] == "sid"
assert second["event"] == "stream_end"
assert second["stream_id"] == "sid"
@pytest.mark.asyncio
async def test_send_non_connection_closed_exception_is_raised() -> None:
bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
mock_ws = AsyncMock()
mock_ws.send.side_effect = RuntimeError("unexpected")
channel._connections["chat-1"] = mock_ws
msg = OutboundMessage(channel="websocket", chat_id="chat-1", content="hello")
with pytest.raises(RuntimeError, match="unexpected"):
await channel.send(msg)
@pytest.mark.asyncio
async def test_send_delta_missing_connection_is_noop() -> None:
bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"], "streaming": True}, bus)
# No exception, no error — just a no-op
await channel.send_delta("nonexistent", "chunk", {"_stream_delta": True, "_stream_id": "s1"})
@pytest.mark.asyncio
async def test_stop_is_idempotent() -> None:
bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus)
# stop() before start() should not raise
await channel.stop()
await channel.stop()
@pytest.mark.asyncio
async def test_end_to_end_client_receives_ready_and_agent_sees_inbound(bus: MagicMock) -> None:
port = 29876
channel = _ch(bus, port=port)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=tester") as client:
ready_raw = await client.recv()
ready = json.loads(ready_raw)
assert ready["event"] == "ready"
assert ready["client_id"] == "tester"
chat_id = ready["chat_id"]
await client.send(json.dumps({"content": "ping from client"}))
await asyncio.sleep(0.08)
bus.publish_inbound.assert_awaited()
inbound = bus.publish_inbound.call_args[0][0]
assert inbound.channel == "websocket"
assert inbound.sender_id == "tester"
assert inbound.chat_id == chat_id
assert inbound.content == "ping from client"
await client.send("plain text frame")
await asyncio.sleep(0.08)
assert bus.publish_inbound.await_count >= 2
second = [c[0][0] for c in bus.publish_inbound.call_args_list][-1]
assert second.content == "plain text frame"
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_token_rejects_handshake_when_mismatch(bus: MagicMock) -> None:
port = 29877
channel = _ch(bus, port=port, path="/", token="secret")
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
with pytest.raises(websockets.exceptions.InvalidStatus) as excinfo:
async with websockets.connect(f"ws://127.0.0.1:{port}/?token=wrong"):
pass
assert excinfo.value.response.status_code == 401
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_wrong_path_returns_404(bus: MagicMock) -> None:
port = 29878
channel = _ch(bus, port=port)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
with pytest.raises(websockets.exceptions.InvalidStatus) as excinfo:
async with websockets.connect(f"ws://127.0.0.1:{port}/other"):
pass
assert excinfo.value.response.status_code == 404
finally:
await channel.stop()
await server_task
def test_registry_discovers_websocket_channel() -> None:
from nanobot.channels.registry import load_channel_class
cls = load_channel_class("websocket")
assert cls.name == "websocket"
@pytest.mark.asyncio
async def test_http_route_issues_token_then_websocket_requires_it(bus: MagicMock) -> None:
port = 29879
channel = _ch(
bus, port=port,
tokenIssuePath="/auth/token",
tokenIssueSecret="route-secret",
websocketRequiresToken=True,
)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
deny = await _http_get(f"http://127.0.0.1:{port}/auth/token")
assert deny.status_code == 401
issue = await _http_get(
f"http://127.0.0.1:{port}/auth/token",
headers={"Authorization": "Bearer route-secret"},
)
assert issue.status_code == 200
token = issue.json()["token"]
assert token.startswith("nbwt_")
with pytest.raises(websockets.exceptions.InvalidStatus) as missing_token:
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=x"):
pass
assert missing_token.value.response.status_code == 401
uri = f"ws://127.0.0.1:{port}/ws?token={token}&client_id=caller"
async with websockets.connect(uri) as client:
ready = json.loads(await client.recv())
assert ready["event"] == "ready"
assert ready["client_id"] == "caller"
with pytest.raises(websockets.exceptions.InvalidStatus) as reuse:
async with websockets.connect(uri):
pass
assert reuse.value.response.status_code == 401
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_end_to_end_server_pushes_streaming_deltas_to_client(bus: MagicMock) -> None:
port = 29880
channel = _ch(bus, port=port, streaming=True)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=stream-tester") as client:
ready_raw = await client.recv()
ready = json.loads(ready_raw)
chat_id = ready["chat_id"]
# Server pushes deltas directly
await channel.send_delta(
chat_id, "Hello ", {"_stream_delta": True, "_stream_id": "s1"}
)
await channel.send_delta(
chat_id, "world", {"_stream_delta": True, "_stream_id": "s1"}
)
await channel.send_delta(
chat_id, "", {"_stream_end": True, "_stream_id": "s1"}
)
delta1 = json.loads(await client.recv())
assert delta1["event"] == "delta"
assert delta1["text"] == "Hello "
assert delta1["stream_id"] == "s1"
delta2 = json.loads(await client.recv())
assert delta2["event"] == "delta"
assert delta2["text"] == "world"
assert delta2["stream_id"] == "s1"
end = json.loads(await client.recv())
assert end["event"] == "stream_end"
assert end["stream_id"] == "s1"
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_token_issue_rejects_when_at_capacity(bus: MagicMock) -> None:
port = 29881
channel = _ch(bus, port=port, tokenIssuePath="/auth/token", tokenIssueSecret="s")
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
# Fill issued tokens to capacity
channel._issued_tokens = {
f"nbwt_fill_{i}": time.monotonic() + 300 for i in range(channel._MAX_ISSUED_TOKENS)
}
resp = await _http_get(
f"http://127.0.0.1:{port}/auth/token",
headers={"Authorization": "Bearer s"},
)
assert resp.status_code == 429
data = resp.json()
assert "error" in data
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_allow_from_rejects_unauthorized_client_id(bus: MagicMock) -> None:
port = 29882
channel = _ch(bus, port=port, allowFrom=["alice", "bob"])
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=eve"):
pass
assert exc_info.value.response.status_code == 403
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_client_id_truncation(bus: MagicMock) -> None:
port = 29883
channel = _ch(bus, port=port)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
long_id = "x" * 200
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id={long_id}") as client:
ready = json.loads(await client.recv())
assert ready["client_id"] == "x" * 128
assert len(ready["client_id"]) == 128
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_non_utf8_binary_frame_ignored(bus: MagicMock) -> None:
port = 29884
channel = _ch(bus, port=port)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=bin-test") as client:
await client.recv() # consume ready
# Send non-UTF-8 bytes
await client.send(b"\xff\xfe\xfd")
await asyncio.sleep(0.05)
# publish_inbound should NOT have been called
bus.publish_inbound.assert_not_awaited()
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_static_token_accepts_issued_token_as_fallback(bus: MagicMock) -> None:
port = 29885
channel = _ch(
bus, port=port,
token="static-secret",
tokenIssuePath="/auth/token",
tokenIssueSecret="route-secret",
)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
# Get an issued token
resp = await _http_get(
f"http://127.0.0.1:{port}/auth/token",
headers={"Authorization": "Bearer route-secret"},
)
assert resp.status_code == 200
issued_token = resp.json()["token"]
# Connect using issued token (not the static one)
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?token={issued_token}&client_id=caller") as client:
ready = json.loads(await client.recv())
assert ready["event"] == "ready"
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_allow_from_empty_list_denies_all(bus: MagicMock) -> None:
port = 29886
channel = _ch(bus, port=port, allowFrom=[])
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=anyone"):
pass
assert exc_info.value.response.status_code == 403
finally:
await channel.stop()
await server_task
@pytest.mark.asyncio
async def test_websocket_requires_token_without_issue_path(bus: MagicMock) -> None:
"""When websocket_requires_token is True but no token or issue path configured, all connections are rejected."""
port = 29887
channel = _ch(bus, port=port, websocketRequiresToken=True)
server_task = asyncio.create_task(channel.start())
await asyncio.sleep(0.3)
try:
# No token at all → 401
with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=u"):
pass
assert exc_info.value.response.status_code == 401
# Wrong token → 401
with pytest.raises(websockets.exceptions.InvalidStatus) as exc_info:
async with websockets.connect(f"ws://127.0.0.1:{port}/ws?client_id=u&token=wrong"):
pass
assert exc_info.value.response.status_code == 401
finally:
await channel.stop()
await server_task
@@ -0,0 +1,477 @@
"""Integration tests for the WebSocket channel using WsTestClient.
Complements the unit/lightweight tests in test_websocket_channel.py by covering
multi-client scenarios, edge cases, and realistic usage patterns.
"""
from __future__ import annotations
import asyncio
import json
from unittest.mock import AsyncMock, MagicMock
import pytest
import websockets
from nanobot.channels.websocket import WebSocketChannel
from nanobot.bus.events import OutboundMessage
from ws_test_client import WsTestClient, issue_token, issue_token_ok
def _ch(bus: Any, port: int, **kw: Any) -> WebSocketChannel:
cfg: dict[str, Any] = {
"enabled": True,
"allowFrom": ["*"],
"host": "127.0.0.1",
"port": port,
"path": "/",
"websocketRequiresToken": False,
}
cfg.update(kw)
return WebSocketChannel(cfg, bus)
@pytest.fixture()
def bus() -> MagicMock:
b = MagicMock()
b.publish_inbound = AsyncMock()
return b
# -- Connection basics ----------------------------------------------------
@pytest.mark.asyncio
async def test_ready_event_fields(bus: MagicMock) -> None:
ch = _ch(bus, 29901)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29901/", client_id="c1") as c:
r = await c.recv_ready()
assert r.event == "ready"
assert len(r.chat_id) == 36
assert r.client_id == "c1"
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_anonymous_client_gets_generated_id(bus: MagicMock) -> None:
ch = _ch(bus, 29902)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29902/", client_id="") as c:
r = await c.recv_ready()
assert r.client_id.startswith("anon-")
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_each_connection_unique_chat_id(bus: MagicMock) -> None:
ch = _ch(bus, 29903)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29903/", client_id="a") as c1:
async with WsTestClient("ws://127.0.0.1:29903/", client_id="b") as c2:
assert (await c1.recv_ready()).chat_id != (await c2.recv_ready()).chat_id
finally:
await ch.stop(); await t
# -- Inbound messages (client -> server) ----------------------------------
@pytest.mark.asyncio
async def test_plain_text(bus: MagicMock) -> None:
ch = _ch(bus, 29904)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29904/", client_id="p") as c:
await c.recv_ready()
await c.send_text("hello world")
await asyncio.sleep(0.1)
inbound = bus.publish_inbound.call_args[0][0]
assert inbound.content == "hello world"
assert inbound.sender_id == "p"
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_json_content_field(bus: MagicMock) -> None:
ch = _ch(bus, 29905)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29905/", client_id="j") as c:
await c.recv_ready()
await c.send_json({"content": "structured"})
await asyncio.sleep(0.1)
assert bus.publish_inbound.call_args[0][0].content == "structured"
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_json_text_and_message_fields(bus: MagicMock) -> None:
ch = _ch(bus, 29906)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29906/", client_id="x") as c:
await c.recv_ready()
await c.send_json({"text": "via text"})
await asyncio.sleep(0.1)
assert bus.publish_inbound.call_args[0][0].content == "via text"
await c.send_json({"message": "via message"})
await asyncio.sleep(0.1)
assert bus.publish_inbound.call_args[0][0].content == "via message"
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_empty_payload_ignored(bus: MagicMock) -> None:
ch = _ch(bus, 29907)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29907/", client_id="e") as c:
await c.recv_ready()
await c.send_text(" ")
await c.send_json({})
await asyncio.sleep(0.1)
bus.publish_inbound.assert_not_awaited()
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_messages_preserve_order(bus: MagicMock) -> None:
ch = _ch(bus, 29908)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29908/", client_id="o") as c:
await c.recv_ready()
for i in range(5):
await c.send_text(f"msg-{i}")
await asyncio.sleep(0.2)
contents = [call[0][0].content for call in bus.publish_inbound.call_args_list]
assert contents == [f"msg-{i}" for i in range(5)]
finally:
await ch.stop(); await t
# -- Outbound messages (server -> client) ---------------------------------
@pytest.mark.asyncio
async def test_server_send_message(bus: MagicMock) -> None:
ch = _ch(bus, 29909)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29909/", client_id="r") as c:
ready = await c.recv_ready()
await ch.send(OutboundMessage(
channel="websocket", chat_id=ready.chat_id, content="reply",
))
msg = await c.recv_message()
assert msg.text == "reply"
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_server_send_with_media_and_reply(bus: MagicMock) -> None:
ch = _ch(bus, 29910)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29910/", client_id="m") as c:
ready = await c.recv_ready()
await ch.send(OutboundMessage(
channel="websocket", chat_id=ready.chat_id, content="img",
media=["/tmp/a.png"], reply_to="m1",
))
msg = await c.recv_message()
assert msg.text == "img"
assert msg.media == ["/tmp/a.png"]
assert msg.reply_to == "m1"
finally:
await ch.stop(); await t
# -- Streaming ------------------------------------------------------------
@pytest.mark.asyncio
async def test_streaming_deltas_and_end(bus: MagicMock) -> None:
ch = _ch(bus, 29911, streaming=True)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29911/", client_id="s") as c:
cid = (await c.recv_ready()).chat_id
for part in ("Hello", " ", "world", "!"):
await ch.send_delta(cid, part, {"_stream_delta": True, "_stream_id": "s1"})
await ch.send_delta(cid, "", {"_stream_end": True, "_stream_id": "s1"})
msgs = await c.collect_stream()
deltas = [m for m in msgs if m.event == "delta"]
assert "".join(d.text for d in deltas) == "Hello world!"
ends = [m for m in msgs if m.event == "stream_end"]
assert len(ends) == 1
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_interleaved_streams(bus: MagicMock) -> None:
ch = _ch(bus, 29912, streaming=True)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29912/", client_id="i") as c:
cid = (await c.recv_ready()).chat_id
await ch.send_delta(cid, "A1", {"_stream_delta": True, "_stream_id": "sa"})
await ch.send_delta(cid, "B1", {"_stream_delta": True, "_stream_id": "sb"})
await ch.send_delta(cid, "A2", {"_stream_delta": True, "_stream_id": "sa"})
await ch.send_delta(cid, "", {"_stream_end": True, "_stream_id": "sa"})
await ch.send_delta(cid, "B2", {"_stream_delta": True, "_stream_id": "sb"})
await ch.send_delta(cid, "", {"_stream_end": True, "_stream_id": "sb"})
msgs = await c.recv_n(6)
sa = "".join(m.text for m in msgs if m.event == "delta" and m.stream_id == "sa")
sb = "".join(m.text for m in msgs if m.event == "delta" and m.stream_id == "sb")
assert sa == "A1A2"
assert sb == "B1B2"
finally:
await ch.stop(); await t
# -- Multi-client ---------------------------------------------------------
@pytest.mark.asyncio
async def test_independent_sessions(bus: MagicMock) -> None:
ch = _ch(bus, 29913)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29913/", client_id="u1") as c1:
async with WsTestClient("ws://127.0.0.1:29913/", client_id="u2") as c2:
r1, r2 = await c1.recv_ready(), await c2.recv_ready()
await ch.send(OutboundMessage(
channel="websocket", chat_id=r1.chat_id, content="for-u1",
))
assert (await c1.recv_message()).text == "for-u1"
await ch.send(OutboundMessage(
channel="websocket", chat_id=r2.chat_id, content="for-u2",
))
assert (await c2.recv_message()).text == "for-u2"
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_disconnected_client_cleanup(bus: MagicMock) -> None:
ch = _ch(bus, 29914)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29914/", client_id="tmp") as c:
chat_id = (await c.recv_ready()).chat_id
# disconnected
await ch.send(OutboundMessage(
channel="websocket", chat_id=chat_id, content="orphan",
))
assert chat_id not in ch._connections
finally:
await ch.stop(); await t
# -- Authentication -------------------------------------------------------
@pytest.mark.asyncio
async def test_static_token_accepted(bus: MagicMock) -> None:
ch = _ch(bus, 29915, token="secret")
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29915/", client_id="a", token="secret") as c:
assert (await c.recv_ready()).client_id == "a"
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_static_token_rejected(bus: MagicMock) -> None:
ch = _ch(bus, 29916, token="correct")
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
async with WsTestClient("ws://127.0.0.1:29916/", client_id="b", token="wrong"):
pass
assert exc.value.response.status_code == 401
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_token_issue_full_flow(bus: MagicMock) -> None:
ch = _ch(bus, 29917, path="/ws",
tokenIssuePath="/auth/token", tokenIssueSecret="s",
websocketRequiresToken=True)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
# no secret -> 401
_, status = await issue_token(port=29917, issue_path="/auth/token")
assert status == 401
# with secret -> token
token = await issue_token_ok(port=29917, issue_path="/auth/token", secret="s")
# no token -> 401
with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
async with WsTestClient("ws://127.0.0.1:29917/ws", client_id="x"):
pass
assert exc.value.response.status_code == 401
# valid token -> ok
async with WsTestClient("ws://127.0.0.1:29917/ws", client_id="ok", token=token) as c:
assert (await c.recv_ready()).client_id == "ok"
# reuse -> 401
with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
async with WsTestClient("ws://127.0.0.1:29917/ws", client_id="r", token=token):
pass
assert exc.value.response.status_code == 401
finally:
await ch.stop(); await t
# -- Path routing ---------------------------------------------------------
@pytest.mark.asyncio
async def test_custom_path(bus: MagicMock) -> None:
ch = _ch(bus, 29918, path="/my-chat")
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29918/my-chat", client_id="p") as c:
assert (await c.recv_ready()).event == "ready"
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_wrong_path_404(bus: MagicMock) -> None:
ch = _ch(bus, 29919, path="/ws")
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
with pytest.raises(websockets.exceptions.InvalidStatus) as exc:
async with WsTestClient("ws://127.0.0.1:29919/wrong", client_id="x"):
pass
assert exc.value.response.status_code == 404
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_trailing_slash_normalized(bus: MagicMock) -> None:
ch = _ch(bus, 29920, path="/ws")
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29920/ws/", client_id="s") as c:
assert (await c.recv_ready()).event == "ready"
finally:
await ch.stop(); await t
# -- Edge cases -----------------------------------------------------------
@pytest.mark.asyncio
async def test_large_message(bus: MagicMock) -> None:
ch = _ch(bus, 29921)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29921/", client_id="big") as c:
await c.recv_ready()
big = "x" * 100_000
await c.send_text(big)
await asyncio.sleep(0.2)
assert bus.publish_inbound.call_args[0][0].content == big
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_unicode_roundtrip(bus: MagicMock) -> None:
ch = _ch(bus, 29922)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29922/", client_id="u") as c:
ready = await c.recv_ready()
text = "你好世界 🌍 日本語テスト"
await c.send_text(text)
await asyncio.sleep(0.1)
assert bus.publish_inbound.call_args[0][0].content == text
await ch.send(OutboundMessage(
channel="websocket", chat_id=ready.chat_id, content=text,
))
assert (await c.recv_message()).text == text
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_rapid_fire(bus: MagicMock) -> None:
ch = _ch(bus, 29923)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29923/", client_id="r") as c:
ready = await c.recv_ready()
for i in range(50):
await c.send_text(f"in-{i}")
await asyncio.sleep(0.5)
assert bus.publish_inbound.await_count == 50
for i in range(50):
await ch.send(OutboundMessage(
channel="websocket", chat_id=ready.chat_id, content=f"out-{i}",
))
received = [(await c.recv_message()).text for _ in range(50)]
assert received == [f"out-{i}" for i in range(50)]
finally:
await ch.stop(); await t
@pytest.mark.asyncio
async def test_invalid_json_as_plain_text(bus: MagicMock) -> None:
ch = _ch(bus, 29924)
t = asyncio.create_task(ch.start())
await asyncio.sleep(0.3)
try:
async with WsTestClient("ws://127.0.0.1:29924/", client_id="j") as c:
await c.recv_ready()
await c.send_text("{broken json")
await asyncio.sleep(0.1)
assert bus.publish_inbound.call_args[0][0].content == "{broken json"
finally:
await ch.stop(); await t
+584
View File
@@ -0,0 +1,584 @@
"""Tests for WeCom channel: helpers, download, upload, send, and message processing."""
import os
import tempfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
try:
import importlib.util
WECOM_AVAILABLE = importlib.util.find_spec("wecom_aibot_sdk") is not None
except ImportError:
WECOM_AVAILABLE = False
if not WECOM_AVAILABLE:
pytest.skip("WeCom dependencies not installed (wecom_aibot_sdk)", allow_module_level=True)
from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus
from nanobot.channels.wecom import (
WecomChannel,
WecomConfig,
_guess_wecom_media_type,
_sanitize_filename,
)
# Try to import the real response class; fall back to a stub if unavailable.
try:
from wecom_aibot_sdk.utils import WsResponse
_RealWsResponse = WsResponse
except ImportError:
_RealWsResponse = None
class _FakeResponse:
"""Minimal stand-in for wecom_aibot_sdk WsResponse."""
def __init__(self, errcode: int = 0, body: dict | None = None, errmsg: str = "ok"):
self.errcode = errcode
self.errmsg = errmsg
self.body = body or {}
class _FakeWsManager:
"""Tracks send_reply calls and returns configurable responses."""
def __init__(self, responses: list[_FakeResponse] | None = None):
self.responses = responses or []
self.calls: list[tuple[str, dict, str]] = []
self._idx = 0
async def send_reply(self, req_id: str, data: dict, cmd: str) -> _FakeResponse:
self.calls.append((req_id, data, cmd))
if self._idx < len(self.responses):
resp = self.responses[self._idx]
self._idx += 1
return resp
return _FakeResponse()
class _FakeFrame:
"""Minimal frame object with a body dict."""
def __init__(self, body: dict | None = None):
self.body = body or {}
class _FakeWeComClient:
"""Fake WeCom client with mock methods."""
def __init__(self, ws_responses: list[_FakeResponse] | None = None):
self._ws_manager = _FakeWsManager(ws_responses)
self.download_file = AsyncMock(return_value=(None, None))
self.reply = AsyncMock()
self.reply_stream = AsyncMock()
self.send_message = AsyncMock()
self.reply_welcome = AsyncMock()
# ── Helper function tests (pure, no async) ──────────────────────────
def test_sanitize_filename_strips_path_traversal() -> None:
assert _sanitize_filename("../../etc/passwd") == "passwd"
def test_sanitize_filename_keeps_chinese_chars() -> None:
assert _sanitize_filename("文件(1.jpg") == "文件(1.jpg"
def test_sanitize_filename_empty_input() -> None:
assert _sanitize_filename("") == ""
def test_guess_wecom_media_type_image() -> None:
for ext in (".jpg", ".jpeg", ".png", ".gif", ".webp", ".bmp"):
assert _guess_wecom_media_type(f"photo{ext}") == "image"
def test_guess_wecom_media_type_video() -> None:
for ext in (".mp4", ".avi", ".mov"):
assert _guess_wecom_media_type(f"video{ext}") == "video"
def test_guess_wecom_media_type_voice() -> None:
for ext in (".amr", ".mp3", ".wav", ".ogg"):
assert _guess_wecom_media_type(f"audio{ext}") == "voice"
def test_guess_wecom_media_type_file_fallback() -> None:
for ext in (".pdf", ".doc", ".xlsx", ".zip"):
assert _guess_wecom_media_type(f"doc{ext}") == "file"
def test_guess_wecom_media_type_case_insensitive() -> None:
assert _guess_wecom_media_type("photo.PNG") == "image"
assert _guess_wecom_media_type("photo.Jpg") == "image"
# ── _download_and_save_media() ──────────────────────────────────────
@pytest.mark.asyncio
async def test_download_and_save_success() -> None:
"""Successful download writes file and returns sanitized path."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
fake_data = b"\x89PNG\r\nfake image"
client.download_file.return_value = (fake_data, "raw_photo.png")
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(tempfile.gettempdir())):
path = await channel._download_and_save_media("https://example.com/img.png", "aes_key", "image", "photo.png")
assert path is not None
assert os.path.isfile(path)
assert os.path.basename(path) == "photo.png"
# Cleanup
os.unlink(path)
@pytest.mark.asyncio
async def test_download_and_save_oversized_rejected() -> None:
"""Data exceeding 200MB is rejected → returns None."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
big_data = b"\x00" * (200 * 1024 * 1024 + 1) # 200MB + 1 byte
client.download_file.return_value = (big_data, "big.bin")
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(tempfile.gettempdir())):
result = await channel._download_and_save_media("https://example.com/big.bin", "key", "file", "big.bin")
assert result is None
@pytest.mark.asyncio
async def test_download_and_save_failure() -> None:
"""SDK returns None data → returns None."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
client.download_file.return_value = (None, None)
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(tempfile.gettempdir())):
result = await channel._download_and_save_media("https://example.com/fail.png", "key", "image")
assert result is None
# ── _upload_media_ws() ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_upload_media_ws_success() -> None:
"""Happy path: init → chunk → finish → returns (media_id, media_type)."""
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(b"\x89PNG\r\n")
tmp = f.name
try:
responses = [
_FakeResponse(errcode=0, body={"upload_id": "up_1"}),
_FakeResponse(errcode=0, body={}),
_FakeResponse(errcode=0, body={"media_id": "media_abc"}),
]
client = _FakeWeComClient(responses)
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
channel._client = client
with patch("wecom_aibot_sdk.utils.generate_req_id", side_effect=lambda x: f"req_{x}"):
media_id, media_type = await channel._upload_media_ws(client, tmp)
assert media_id == "media_abc"
assert media_type == "image"
finally:
os.unlink(tmp)
@pytest.mark.asyncio
async def test_upload_media_ws_oversized_file() -> None:
"""File >200MB triggers ValueError → returns (None, None)."""
# Instead of creating a real 200MB+ file, mock os.path.getsize and open
with patch("os.path.getsize", return_value=200 * 1024 * 1024 + 1), \
patch("builtins.open", MagicMock()):
client = _FakeWeComClient()
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
channel._client = client
result = await channel._upload_media_ws(client, "/fake/large.bin")
assert result == (None, None)
@pytest.mark.asyncio
async def test_upload_media_ws_init_failure() -> None:
"""Init step returns errcode != 0 → returns (None, None)."""
with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as f:
f.write(b"hello")
tmp = f.name
try:
responses = [
_FakeResponse(errcode=50001, errmsg="invalid"),
]
client = _FakeWeComClient(responses)
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
channel._client = client
with patch("wecom_aibot_sdk.utils.generate_req_id", side_effect=lambda x: f"req_{x}"):
result = await channel._upload_media_ws(client, tmp)
assert result == (None, None)
finally:
os.unlink(tmp)
@pytest.mark.asyncio
async def test_upload_media_ws_chunk_failure() -> None:
"""Chunk step returns errcode != 0 → returns (None, None)."""
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(b"\x89PNG\r\n")
tmp = f.name
try:
responses = [
_FakeResponse(errcode=0, body={"upload_id": "up_1"}),
_FakeResponse(errcode=50002, errmsg="chunk fail"),
]
client = _FakeWeComClient(responses)
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
channel._client = client
with patch("wecom_aibot_sdk.utils.generate_req_id", side_effect=lambda x: f"req_{x}"):
result = await channel._upload_media_ws(client, tmp)
assert result == (None, None)
finally:
os.unlink(tmp)
@pytest.mark.asyncio
async def test_upload_media_ws_finish_no_media_id() -> None:
"""Finish step returns empty media_id → returns (None, None)."""
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(b"\x89PNG\r\n")
tmp = f.name
try:
responses = [
_FakeResponse(errcode=0, body={"upload_id": "up_1"}),
_FakeResponse(errcode=0, body={}),
_FakeResponse(errcode=0, body={}), # no media_id
]
client = _FakeWeComClient(responses)
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
channel._client = client
with patch("wecom_aibot_sdk.utils.generate_req_id", side_effect=lambda x: f"req_{x}"):
result = await channel._upload_media_ws(client, tmp)
assert result == (None, None)
finally:
os.unlink(tmp)
# ── send() ──────────────────────────────────────────────────────────
@pytest.mark.asyncio
async def test_send_text_with_frame() -> None:
"""When frame is stored, send uses reply_stream for final text."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
channel._generate_req_id = lambda x: f"req_{x}"
channel._chat_frames["chat1"] = _FakeFrame()
await channel.send(
OutboundMessage(channel="wecom", chat_id="chat1", content="hello")
)
client.reply_stream.assert_called_once()
call_args = client.reply_stream.call_args
assert call_args[0][2] == "hello" # content arg
@pytest.mark.asyncio
async def test_send_progress_with_frame() -> None:
"""When metadata has _progress, send uses reply_stream with finish=False."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
channel._generate_req_id = lambda x: f"req_{x}"
channel._chat_frames["chat1"] = _FakeFrame()
await channel.send(
OutboundMessage(channel="wecom", chat_id="chat1", content="thinking...", metadata={"_progress": True})
)
client.reply_stream.assert_called_once()
call_args = client.reply_stream.call_args
assert call_args[0][2] == "thinking..." # content arg
assert call_args[1]["finish"] is False
@pytest.mark.asyncio
async def test_send_proactive_without_frame() -> None:
"""Without stored frame, send uses send_message with markdown."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
await channel.send(
OutboundMessage(channel="wecom", chat_id="chat1", content="proactive msg")
)
client.send_message.assert_called_once()
call_args = client.send_message.call_args
assert call_args[0][0] == "chat1"
assert call_args[0][1]["msgtype"] == "markdown"
@pytest.mark.asyncio
async def test_send_media_then_text() -> None:
"""Media files are uploaded and sent before text content."""
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(b"\x89PNG\r\n")
tmp = f.name
try:
responses = [
_FakeResponse(errcode=0, body={"upload_id": "up_1"}),
_FakeResponse(errcode=0, body={}),
_FakeResponse(errcode=0, body={"media_id": "media_123"}),
]
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient(responses)
channel._client = client
channel._generate_req_id = lambda x: f"req_{x}"
channel._chat_frames["chat1"] = _FakeFrame()
await channel.send(
OutboundMessage(channel="wecom", chat_id="chat1", content="see image", media=[tmp])
)
# Media should have been sent via reply
media_calls = [c for c in client.reply.call_args_list if c[0][1].get("msgtype") == "image"]
assert len(media_calls) == 1
assert media_calls[0][0][1]["image"]["media_id"] == "media_123"
# Text should have been sent via reply_stream
client.reply_stream.assert_called_once()
finally:
os.unlink(tmp)
@pytest.mark.asyncio
async def test_send_media_file_not_found() -> None:
"""Non-existent media path is skipped with a warning."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
channel._generate_req_id = lambda x: f"req_{x}"
channel._chat_frames["chat1"] = _FakeFrame()
await channel.send(
OutboundMessage(channel="wecom", chat_id="chat1", content="hello", media=["/nonexistent/file.png"])
)
# reply_stream should still be called for the text part
client.reply_stream.assert_called_once()
# No media reply should happen
media_calls = [c for c in client.reply.call_args_list if c[0][1].get("msgtype") in ("image", "file", "video")]
assert len(media_calls) == 0
@pytest.mark.asyncio
async def test_send_exception_caught_not_raised() -> None:
"""Exceptions inside send() must not propagate."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["*"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
channel._generate_req_id = lambda x: f"req_{x}"
channel._chat_frames["chat1"] = _FakeFrame()
# Make reply_stream raise
client.reply_stream.side_effect = RuntimeError("boom")
await channel.send(
OutboundMessage(channel="wecom", chat_id="chat1", content="fail test")
)
# No exception — test passes if we reach here.
# ── _process_message() ──────────────────────────────────────────────
@pytest.mark.asyncio
async def test_process_text_message() -> None:
"""Text message is routed to bus with correct fields."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
frame = _FakeFrame(body={
"msgid": "msg_text_1",
"chatid": "chat1",
"chattype": "single",
"from": {"userid": "user1"},
"text": {"content": "hello wecom"},
})
await channel._process_message(frame, "text")
msg = await channel.bus.consume_inbound()
assert msg.sender_id == "user1"
assert msg.chat_id == "chat1"
assert msg.content == "hello wecom"
assert msg.metadata["msg_type"] == "text"
@pytest.mark.asyncio
async def test_process_image_message() -> None:
"""Image message: download success → media_paths non-empty."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
client = _FakeWeComClient()
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
f.write(b"\x89PNG\r\n")
saved = f.name
client.download_file.return_value = (b"\x89PNG\r\n", "photo.png")
channel._client = client
try:
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(os.path.dirname(saved))):
frame = _FakeFrame(body={
"msgid": "msg_img_1",
"chatid": "chat1",
"from": {"userid": "user1"},
"image": {"url": "https://example.com/img.png", "aeskey": "key123"},
})
await channel._process_message(frame, "image")
msg = await channel.bus.consume_inbound()
assert len(msg.media) == 1
assert msg.media[0].endswith("photo.png")
assert "[image:" in msg.content
finally:
if os.path.exists(saved):
pass # may have been overwritten; clean up if exists
# Clean up any photo.png in tempdir
p = os.path.join(os.path.dirname(saved), "photo.png")
if os.path.exists(p):
os.unlink(p)
@pytest.mark.asyncio
async def test_process_file_message() -> None:
"""File message: download success → media_paths non-empty (critical fix verification)."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
client = _FakeWeComClient()
with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as f:
f.write(b"%PDF-1.4 fake")
saved = f.name
client.download_file.return_value = (b"%PDF-1.4 fake", "report.pdf")
channel._client = client
try:
with patch("nanobot.channels.wecom.get_media_dir", return_value=Path(os.path.dirname(saved))):
frame = _FakeFrame(body={
"msgid": "msg_file_1",
"chatid": "chat1",
"from": {"userid": "user1"},
"file": {"url": "https://example.com/report.pdf", "aeskey": "key456", "name": "report.pdf"},
})
await channel._process_message(frame, "file")
msg = await channel.bus.consume_inbound()
assert len(msg.media) == 1
assert msg.media[0].endswith("report.pdf")
assert "[file: report.pdf]" in msg.content
finally:
p = os.path.join(os.path.dirname(saved), "report.pdf")
if os.path.exists(p):
os.unlink(p)
@pytest.mark.asyncio
async def test_process_voice_message() -> None:
"""Voice message: transcribed text is included in content."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
frame = _FakeFrame(body={
"msgid": "msg_voice_1",
"chatid": "chat1",
"from": {"userid": "user1"},
"voice": {"content": "transcribed text here"},
})
await channel._process_message(frame, "voice")
msg = await channel.bus.consume_inbound()
assert "transcribed text here" in msg.content
assert "[voice]" in msg.content
@pytest.mark.asyncio
async def test_process_message_deduplication() -> None:
"""Same msg_id is not processed twice."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
frame = _FakeFrame(body={
"msgid": "msg_dup_1",
"chatid": "chat1",
"from": {"userid": "user1"},
"text": {"content": "once"},
})
await channel._process_message(frame, "text")
await channel._process_message(frame, "text")
msg = await channel.bus.consume_inbound()
assert msg.content == "once"
# Second message should not appear on the bus
assert channel.bus.inbound.empty()
@pytest.mark.asyncio
async def test_process_message_empty_content_skipped() -> None:
"""Message with empty content produces no bus message."""
channel = WecomChannel(WecomConfig(bot_id="b", secret="s", allow_from=["user1"]), MessageBus())
client = _FakeWeComClient()
channel._client = client
frame = _FakeFrame(body={
"msgid": "msg_empty_1",
"chatid": "chat1",
"from": {"userid": "user1"},
"text": {"content": ""},
})
await channel._process_message(frame, "text")
assert channel.bus.inbound.empty()
+227
View File
@@ -0,0 +1,227 @@
"""Lightweight WebSocket test client for integration testing the nanobot WebSocket channel.
Provides an async ``WsTestClient`` class and token-issuance helpers that
integration tests can import and use directly::
from ws_test_client import WsTestClient
async with WsTestClient("ws://127.0.0.1:8765/", client_id="t") as c:
ready = await c.recv_ready()
await c.send_text("hello")
msg = await c.recv_message()
"""
from __future__ import annotations
import asyncio
import json
from dataclasses import dataclass, field
from typing import Any
import httpx
import websockets
from websockets.asyncio.client import ClientConnection
@dataclass
class WsMessage:
"""A parsed message received from the WebSocket server."""
event: str
raw: dict[str, Any] = field(repr=False)
@property
def text(self) -> str | None:
return self.raw.get("text")
@property
def chat_id(self) -> str | None:
return self.raw.get("chat_id")
@property
def client_id(self) -> str | None:
return self.raw.get("client_id")
@property
def media(self) -> list[str] | None:
return self.raw.get("media")
@property
def reply_to(self) -> str | None:
return self.raw.get("reply_to")
@property
def stream_id(self) -> str | None:
return self.raw.get("stream_id")
def __eq__(self, other: object) -> bool:
if not isinstance(other, WsMessage):
return NotImplemented
return self.event == other.event and self.raw == other.raw
class WsTestClient:
"""Async WebSocket test client with helper methods for common operations.
Usage::
async with WsTestClient("ws://127.0.0.1:8765/", client_id="tester") as client:
ready = await client.recv_ready()
await client.send_text("hello")
msg = await client.recv_message(timeout=5.0)
"""
def __init__(
self,
uri: str,
*,
client_id: str = "test-client",
token: str = "",
extra_headers: dict[str, str] | None = None,
) -> None:
params: list[str] = []
if client_id:
params.append(f"client_id={client_id}")
if token:
params.append(f"token={token}")
sep = "&" if "?" in uri else "?"
self._uri = uri + sep + "&".join(params) if params else uri
self._extra_headers = extra_headers
self._ws: ClientConnection | None = None
async def connect(self) -> None:
self._ws = await websockets.connect(
self._uri,
additional_headers=self._extra_headers,
)
async def close(self) -> None:
if self._ws:
await self._ws.close()
self._ws = None
async def __aenter__(self) -> WsTestClient:
await self.connect()
return self
async def __aexit__(self, *args: Any) -> None:
await self.close()
@property
def ws(self) -> ClientConnection:
assert self._ws is not None, "Client is not connected"
return self._ws
# -- Receiving --------------------------------------------------------
async def recv_raw(self, timeout: float = 10.0) -> dict[str, Any]:
"""Receive and parse one raw JSON message with timeout."""
raw = await asyncio.wait_for(self.ws.recv(), timeout=timeout)
return json.loads(raw)
async def recv(self, timeout: float = 10.0) -> WsMessage:
"""Receive one message, returning a WsMessage wrapper."""
data = await self.recv_raw(timeout)
return WsMessage(event=data.get("event", ""), raw=data)
async def recv_ready(self, timeout: float = 5.0) -> WsMessage:
"""Receive and validate the 'ready' event."""
msg = await self.recv(timeout)
assert msg.event == "ready", f"Expected 'ready' event, got '{msg.event}'"
return msg
async def recv_message(self, timeout: float = 10.0) -> WsMessage:
"""Receive and validate a 'message' event."""
msg = await self.recv(timeout)
assert msg.event == "message", f"Expected 'message' event, got '{msg.event}'"
return msg
async def recv_delta(self, timeout: float = 10.0) -> WsMessage:
"""Receive and validate a 'delta' event."""
msg = await self.recv(timeout)
assert msg.event == "delta", f"Expected 'delta' event, got '{msg.event}'"
return msg
async def recv_stream_end(self, timeout: float = 10.0) -> WsMessage:
"""Receive and validate a 'stream_end' event."""
msg = await self.recv(timeout)
assert msg.event == "stream_end", f"Expected 'stream_end' event, got '{msg.event}'"
return msg
async def collect_stream(self, timeout: float = 10.0) -> list[WsMessage]:
"""Collect all deltas and the final stream_end into a list."""
messages: list[WsMessage] = []
while True:
msg = await self.recv(timeout)
messages.append(msg)
if msg.event == "stream_end":
break
return messages
async def recv_n(self, n: int, timeout: float = 10.0) -> list[WsMessage]:
"""Receive exactly *n* messages."""
return [await self.recv(timeout) for _ in range(n)]
# -- Sending ----------------------------------------------------------
async def send_text(self, text: str) -> None:
"""Send a plain text frame."""
await self.ws.send(text)
async def send_json(self, data: dict[str, Any]) -> None:
"""Send a JSON frame."""
await self.ws.send(json.dumps(data, ensure_ascii=False))
async def send_content(self, content: str) -> None:
"""Send content in the preferred JSON format ``{"content": ...}``."""
await self.send_json({"content": content})
# -- Connection introspection -----------------------------------------
@property
def closed(self) -> bool:
return self._ws is None or self._ws.closed
# -- Token issuance helpers -----------------------------------------------
async def issue_token(
host: str = "127.0.0.1",
port: int = 8765,
issue_path: str = "/auth/token",
secret: str = "",
) -> tuple[dict[str, Any] | None, int]:
"""Request a short-lived token from the token-issue HTTP endpoint.
Returns ``(parsed_json_or_None, status_code)``.
"""
url = f"http://{host}:{port}{issue_path}"
headers: dict[str, str] = {}
if secret:
headers["Authorization"] = f"Bearer {secret}"
loop = asyncio.get_running_loop()
resp = await loop.run_in_executor(
None, lambda: httpx.get(url, headers=headers, timeout=5.0)
)
try:
data = resp.json()
except Exception:
data = None
return data, resp.status_code
async def issue_token_ok(
host: str = "127.0.0.1",
port: int = 8765,
issue_path: str = "/auth/token",
secret: str = "",
) -> str:
"""Request a token, asserting success, and return the token string."""
(data, status) = await issue_token(host, port, issue_path, secret)
assert status == 200, f"Token issue failed with status {status}"
assert data is not None
token = data["token"]
assert token.startswith("nbwt_"), f"Unexpected token format: {token}"
return token
+528
View File
@@ -0,0 +1,528 @@
import json
import pytest
# Check optional msteams dependencies before running tests
try:
from nanobot.channels import msteams
MSTEAMS_AVAILABLE = getattr(msteams, "MSTEAMS_AVAILABLE", False)
except ImportError:
MSTEAMS_AVAILABLE = False
if not MSTEAMS_AVAILABLE:
pytest.skip("MSTeams dependencies not installed (PyJWT, cryptography). Run: pip install nanobot-ai[msteams]", allow_module_level=True)
import jwt
from cryptography.hazmat.primitives.asymmetric import rsa
import nanobot.channels.msteams as msteams_module
from nanobot.bus.events import OutboundMessage
from nanobot.channels.msteams import ConversationRef, MSTeamsChannel, MSTeamsConfig
class DummyBus:
def __init__(self):
self.inbound = []
async def publish_inbound(self, msg):
self.inbound.append(msg)
class FakeResponse:
def __init__(self, payload=None, *, should_raise=False):
self._payload = payload or {}
self._should_raise = should_raise
def raise_for_status(self):
if self._should_raise:
raise RuntimeError("boom")
return None
def json(self):
return self._payload
class FakeHttpClient:
def __init__(self, payload=None, *, should_raise=False):
self.payload = payload or {"access_token": "tok", "expires_in": 3600}
self.should_raise = should_raise
self.calls = []
async def post(self, url, **kwargs):
self.calls.append((url, kwargs))
return FakeResponse(self.payload, should_raise=self.should_raise)
@pytest.fixture
def make_channel(tmp_path, monkeypatch):
monkeypatch.setattr("nanobot.channels.msteams.get_workspace_path", lambda: tmp_path)
def _make_channel(**config_overrides):
config = {
"enabled": True,
"appId": "app-id",
"appPassword": "secret",
"tenantId": "tenant-id",
"allowFrom": ["*"],
}
config.update(config_overrides)
return MSTeamsChannel(config, DummyBus())
return _make_channel
@pytest.mark.asyncio
async def test_handle_activity_personal_message_publishes_and_stores_ref(make_channel, tmp_path):
ch = make_channel()
activity = {
"type": "message",
"id": "activity-1",
"text": "Hello from Teams",
"serviceUrl": "https://smba.trafficmanager.net/amer/",
"conversation": {
"id": "conv-123",
"conversationType": "personal",
},
"from": {
"id": "29:user-id",
"aadObjectId": "aad-user-1",
"name": "Bob",
},
"recipient": {
"id": "28:bot-id",
"name": "nanobot",
},
"channelData": {
"tenant": {"id": "tenant-id"},
},
}
await ch._handle_activity(activity)
assert len(ch.bus.inbound) == 1
msg = ch.bus.inbound[0]
assert msg.channel == "msteams"
assert msg.sender_id == "aad-user-1"
assert msg.chat_id == "conv-123"
assert msg.content == "Hello from Teams"
assert msg.metadata["msteams"]["conversation_id"] == "conv-123"
assert "conv-123" in ch._conversation_refs
saved = json.loads((tmp_path / "state" / "msteams_conversations.json").read_text(encoding="utf-8"))
assert saved["conv-123"]["conversation_id"] == "conv-123"
assert saved["conv-123"]["tenant_id"] == "tenant-id"
@pytest.mark.asyncio
async def test_handle_activity_ignores_group_messages(make_channel):
ch = make_channel()
activity = {
"type": "message",
"id": "activity-2",
"text": "Hello group",
"serviceUrl": "https://smba.trafficmanager.net/amer/",
"conversation": {
"id": "conv-group",
"conversationType": "channel",
},
"from": {
"id": "29:user-id",
"aadObjectId": "aad-user-1",
"name": "Bob",
},
"recipient": {
"id": "28:bot-id",
"name": "nanobot",
},
}
await ch._handle_activity(activity)
assert ch.bus.inbound == []
assert ch._conversation_refs == {}
@pytest.mark.asyncio
async def test_handle_activity_mention_only_uses_default_response(make_channel):
ch = make_channel()
activity = {
"type": "message",
"id": "activity-3",
"text": "<at>Nanobot</at>",
"serviceUrl": "https://smba.trafficmanager.net/amer/",
"conversation": {
"id": "conv-empty",
"conversationType": "personal",
},
"from": {
"id": "29:user-id",
"aadObjectId": "aad-user-1",
"name": "Bob",
},
"recipient": {
"id": "28:bot-id",
"name": "nanobot",
},
}
await ch._handle_activity(activity)
assert len(ch.bus.inbound) == 1
assert ch.bus.inbound[0].content == "Hi — what can I help with?"
assert "conv-empty" in ch._conversation_refs
@pytest.mark.asyncio
async def test_handle_activity_mention_only_ignores_when_response_disabled(make_channel):
ch = make_channel(mentionOnlyResponse=" ")
activity = {
"type": "message",
"id": "activity-4",
"text": "<at>Nanobot</at>",
"serviceUrl": "https://smba.trafficmanager.net/amer/",
"conversation": {
"id": "conv-empty-disabled",
"conversationType": "personal",
},
"from": {
"id": "29:user-id",
"aadObjectId": "aad-user-1",
"name": "Bob",
},
"recipient": {
"id": "28:bot-id",
"name": "nanobot",
},
}
await ch._handle_activity(activity)
assert ch.bus.inbound == []
assert ch._conversation_refs == {}
def test_strip_possible_bot_mention_removes_generic_at_tags(make_channel):
ch = make_channel()
assert ch._strip_possible_bot_mention("<at>Nanobot</at> hello") == "hello"
assert ch._strip_possible_bot_mention("hi <at>Some Bot</at> there") == "hi there"
def test_sanitize_inbound_text_keeps_normal_inline_message(make_channel):
ch = make_channel()
activity = {
"text": "<at>Nanobot</at> normal inline message",
"channelData": {},
}
assert ch._sanitize_inbound_text(activity) == "normal inline message"
def test_sanitize_inbound_text_normalizes_fwdioc_wrapper_without_reply_metadata(make_channel):
ch = make_channel()
activity = {
"text": "FWDIOC-BOT \r\nQuoted prior message\r\n\r\nThis is a reply with quote test",
"channelData": {},
}
assert ch._sanitize_inbound_text(activity) == (
"User is replying to: Quoted prior message\n"
"User reply: This is a reply with quote test"
)
def test_sanitize_inbound_text_structures_reply_quote_prefix(make_channel):
ch = make_channel()
activity = {
"text": "Replying to Bob Smith\nactual reply text",
"replyToId": "parent-activity",
"channelData": {"messageType": "reply"},
}
assert ch._sanitize_inbound_text(activity) == "User is replying to: Bob Smith\nUser reply: actual reply text"
def test_sanitize_inbound_text_structures_live_fwdioc_quote_shape(make_channel):
ch = make_channel()
activity = {
"text": "FWDIOC-BOT Got it. Ill watch for the exact text reply with quote test and then inspect that turn specifically. Reply with quote test",
"replyToId": "parent-activity",
"channelData": {"messageType": "reply"},
}
assert ch._sanitize_inbound_text(activity) == (
"User is replying to: Got it. Ill watch for the exact text reply with quote test and then inspect that turn specifically.\n"
"User reply: Reply with quote test"
)
def test_sanitize_inbound_text_structures_multiline_fwdioc_quote_shape(make_channel):
ch = make_channel()
activity = {
"text": (
"FWDIOC-BOT\r\n"
"Understood — then the restart already happened, and the new Teams quote normalization should now be live. "
"Next best step: • send one more real reply-with-quote message in Teams • I&rsquo…\r\n"
"\r\n"
"This is a reply with quote"
),
"replyToId": "parent-activity",
"channelData": {"messageType": "reply"},
}
assert ch._sanitize_inbound_text(activity) == (
"User is replying to: Understood — then the restart already happened, and the new Teams quote normalization should now be live. "
"Next best step: • send one more real reply-with-quote message in Teams • I’…\n"
"User reply: This is a reply with quote"
)
def test_sanitize_inbound_text_structures_exact_live_crlf_fwdioc_shape(make_channel):
ch = make_channel()
activity = {
"text": (
"FWDIOC-BOT \r\n"
"Please send one real reply-with-quote message in Teams. That single test should be enough now: "
"• Ill check the new MSTeams sanitized inbound text ... log • and compare it to the prompt…\r\n"
"\r\n"
"This is a reply with quote test"
),
"replyToId": "parent-activity",
"channelData": {"messageType": "reply"},
}
assert ch._sanitize_inbound_text(activity) == (
"User is replying to: Please send one real reply-with-quote message in Teams. That single test should be enough now: "
"• Ill check the new MSTeams sanitized inbound text ... log • and compare it to the prompt…\n"
"User reply: This is a reply with quote test"
)
@pytest.mark.asyncio
async def test_get_access_token_uses_configured_tenant(make_channel):
ch = make_channel(tenantId="tenant-123")
fake_http = FakeHttpClient()
ch._http = fake_http
token = await ch._get_access_token()
assert token == "tok"
assert len(fake_http.calls) == 1
url, kwargs = fake_http.calls[0]
assert url == "https://login.microsoftonline.com/tenant-123/oauth2/v2.0/token"
assert kwargs["data"]["client_id"] == "app-id"
assert kwargs["data"]["client_secret"] == "secret"
assert kwargs["data"]["scope"] == "https://api.botframework.com/.default"
@pytest.mark.asyncio
async def test_send_replies_to_activity_when_reply_in_thread_enabled(make_channel):
ch = make_channel(replyInThread=True)
fake_http = FakeHttpClient()
ch._http = fake_http
ch._token = "tok"
ch._token_expires_at = 9999999999
ch._conversation_refs["conv-123"] = ConversationRef(
service_url="https://smba.trafficmanager.net/amer/",
conversation_id="conv-123",
activity_id="activity-1",
)
await ch.send(OutboundMessage(channel="msteams", chat_id="conv-123", content="Reply text"))
assert len(fake_http.calls) == 1
url, kwargs = fake_http.calls[0]
assert url == "https://smba.trafficmanager.net/amer/v3/conversations/conv-123/activities/activity-1"
assert kwargs["headers"]["Authorization"] == "Bearer tok"
assert kwargs["json"]["text"] == "Reply text"
assert kwargs["json"]["replyToId"] == "activity-1"
@pytest.mark.asyncio
async def test_send_posts_to_conversation_when_thread_reply_disabled(make_channel):
ch = make_channel(replyInThread=False)
fake_http = FakeHttpClient()
ch._http = fake_http
ch._token = "tok"
ch._token_expires_at = 9999999999
ch._conversation_refs["conv-123"] = ConversationRef(
service_url="https://smba.trafficmanager.net/amer/",
conversation_id="conv-123",
activity_id="activity-1",
)
await ch.send(OutboundMessage(channel="msteams", chat_id="conv-123", content="Reply text"))
assert len(fake_http.calls) == 1
url, kwargs = fake_http.calls[0]
assert url == "https://smba.trafficmanager.net/amer/v3/conversations/conv-123/activities"
assert kwargs["headers"]["Authorization"] == "Bearer tok"
assert kwargs["json"]["text"] == "Reply text"
assert "replyToId" not in kwargs["json"]
@pytest.mark.asyncio
async def test_send_posts_to_conversation_when_thread_reply_enabled_but_no_activity_id(make_channel):
ch = make_channel(replyInThread=True)
fake_http = FakeHttpClient()
ch._http = fake_http
ch._token = "tok"
ch._token_expires_at = 9999999999
ch._conversation_refs["conv-123"] = ConversationRef(
service_url="https://smba.trafficmanager.net/amer/",
conversation_id="conv-123",
activity_id=None,
)
await ch.send(OutboundMessage(channel="msteams", chat_id="conv-123", content="Reply text"))
assert len(fake_http.calls) == 1
url, kwargs = fake_http.calls[0]
assert url == "https://smba.trafficmanager.net/amer/v3/conversations/conv-123/activities"
assert kwargs["headers"]["Authorization"] == "Bearer tok"
assert kwargs["json"]["text"] == "Reply text"
assert "replyToId" not in kwargs["json"]
@pytest.mark.asyncio
async def test_send_raises_when_conversation_ref_missing(make_channel):
ch = make_channel()
ch._http = FakeHttpClient()
with pytest.raises(RuntimeError, match="conversation ref not found"):
await ch.send(OutboundMessage(channel="msteams", chat_id="missing", content="Reply text"))
@pytest.mark.asyncio
async def test_send_raises_delivery_failures_for_retry(make_channel):
ch = make_channel()
ch._http = FakeHttpClient(should_raise=True)
ch._token = "tok"
ch._token_expires_at = 9999999999
ch._conversation_refs["conv-123"] = ConversationRef(
service_url="https://smba.trafficmanager.net/amer/",
conversation_id="conv-123",
activity_id="activity-1",
)
with pytest.raises(RuntimeError, match="boom"):
await ch.send(OutboundMessage(channel="msteams", chat_id="conv-123", content="Reply text"))
def _make_test_rsa_jwk(kid: str = "test-kid"):
private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
public_key = private_key.public_key()
jwk = json.loads(jwt.algorithms.RSAAlgorithm.to_jwk(public_key))
jwk["kid"] = kid
jwk["use"] = "sig"
jwk["kty"] = "RSA"
jwk["alg"] = "RS256"
return private_key, jwk
@pytest.mark.asyncio
async def test_validate_inbound_auth_accepts_observed_botframework_shape(make_channel):
ch = make_channel(validateInboundAuth=True)
private_key, jwk = _make_test_rsa_jwk()
ch._botframework_jwks = {"keys": [jwk]}
ch._botframework_jwks_expires_at = 9999999999
service_url = "https://smba.trafficmanager.net/amer/tenant/"
token = jwt.encode(
{
"iss": "https://api.botframework.com",
"aud": "app-id",
"serviceurl": service_url,
"nbf": 1700000000,
"exp": 4100000000,
},
private_key,
algorithm="RS256",
headers={"kid": jwk["kid"]},
)
await ch._validate_inbound_auth(
f"Bearer {token}",
{"serviceUrl": service_url},
)
@pytest.mark.asyncio
async def test_validate_inbound_auth_rejects_service_url_mismatch(make_channel):
ch = make_channel(validateInboundAuth=True)
private_key, jwk = _make_test_rsa_jwk()
ch._botframework_jwks = {"keys": [jwk]}
ch._botframework_jwks_expires_at = 9999999999
token = jwt.encode(
{
"iss": "https://api.botframework.com",
"aud": "app-id",
"serviceurl": "https://smba.trafficmanager.net/amer/tenant-a/",
"nbf": 1700000000,
"exp": 4100000000,
},
private_key,
algorithm="RS256",
headers={"kid": jwk["kid"]},
)
with pytest.raises(ValueError, match="serviceUrl claim mismatch"):
await ch._validate_inbound_auth(
f"Bearer {token}",
{"serviceUrl": "https://smba.trafficmanager.net/amer/tenant-b/"},
)
@pytest.mark.asyncio
async def test_validate_inbound_auth_rejects_missing_bearer_token(make_channel):
ch = make_channel(validateInboundAuth=True)
with pytest.raises(ValueError, match="missing bearer token"):
await ch._validate_inbound_auth("", {"serviceUrl": "https://smba.trafficmanager.net/amer/tenant/"})
@pytest.mark.asyncio
async def test_start_logs_install_hint_when_pyjwt_missing(make_channel, monkeypatch):
ch = make_channel()
errors = []
monkeypatch.setattr(msteams_module, "MSTEAMS_AVAILABLE", False)
monkeypatch.setattr(msteams_module.logger, "error", lambda message, *args: errors.append(message.format(*args)))
await ch.start()
assert errors == ["PyJWT not installed. Run: pip install nanobot-ai[msteams]"]
def test_msteams_default_config_includes_restart_notify_fields():
cfg = MSTeamsChannel.default_config()
assert cfg["restartNotifyEnabled"] is False
assert "restartNotifyPreMessage" in cfg
assert "restartNotifyPostMessage" in cfg
def test_msteams_config_accepts_restart_notify_aliases():
cfg = MSTeamsConfig.model_validate(
{
"restartNotifyEnabled": True,
"restartNotifyPreMessage": "Restarting now.",
"restartNotifyPostMessage": "Back online.",
}
)
assert cfg.restart_notify_enabled is True
assert cfg.restart_notify_pre_message == "Restarting now."
assert cfg.restart_notify_post_message == "Back online."
+7
View File
@@ -1,10 +1,15 @@
"""Tests for exec tool environment isolation.""" """Tests for exec tool environment isolation."""
import sys
import pytest import pytest
from nanobot.agent.tools.shell import ExecTool from nanobot.agent.tools.shell import ExecTool
_UNIX_ONLY = pytest.mark.skipif(sys.platform == "win32", reason="Unix shell commands")
@_UNIX_ONLY
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_exec_does_not_leak_parent_env(monkeypatch): async def test_exec_does_not_leak_parent_env(monkeypatch):
"""Env vars from the parent process must not be visible to commands.""" """Env vars from the parent process must not be visible to commands."""
@@ -22,6 +27,7 @@ async def test_exec_has_working_path():
assert "hello" in result assert "hello" in result
@_UNIX_ONLY
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_exec_path_append(): async def test_exec_path_append():
"""The pathAppend config should be available in the command's PATH.""" """The pathAppend config should be available in the command's PATH."""
@@ -30,6 +36,7 @@ async def test_exec_path_append():
assert "/opt/custom/bin" in result assert "/opt/custom/bin" in result
@_UNIX_ONLY
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_exec_path_append_preserves_system_path(): async def test_exec_path_append_preserves_system_path():
"""pathAppend must not clobber standard system paths.""" """pathAppend must not clobber standard system paths."""
+269
View File
@@ -0,0 +1,269 @@
"""Tests for cross-platform shell execution.
Verifies that ExecTool selects the correct shell, environment, path-append
strategy, and sandbox behaviour per platform without actually running
platform-specific binaries (all subprocess calls are mocked).
"""
from unittest.mock import AsyncMock, patch
import pytest
from nanobot.agent.tools.shell import ExecTool
# ---------------------------------------------------------------------------
# _build_env
# ---------------------------------------------------------------------------
class TestBuildEnvUnix:
def test_expected_keys(self):
with patch("nanobot.agent.tools.shell._IS_WINDOWS", False):
env = ExecTool()._build_env()
assert set(env) == {"HOME", "LANG", "TERM"}
def test_home_from_environ(self, monkeypatch):
monkeypatch.setenv("HOME", "/Users/dev")
with patch("nanobot.agent.tools.shell._IS_WINDOWS", False):
env = ExecTool()._build_env()
assert env["HOME"] == "/Users/dev"
def test_secrets_excluded(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "sk-secret")
monkeypatch.setenv("NANOBOT_TOKEN", "tok-secret")
with patch("nanobot.agent.tools.shell._IS_WINDOWS", False):
env = ExecTool()._build_env()
assert "OPENAI_API_KEY" not in env
assert "NANOBOT_TOKEN" not in env
for v in env.values():
assert "secret" not in v.lower()
class TestBuildEnvWindows:
_EXPECTED_KEYS = {
"SYSTEMROOT", "COMSPEC", "USERPROFILE", "HOMEDRIVE",
"HOMEPATH", "TEMP", "TMP", "PATHEXT", "PATH",
}
def test_expected_keys(self):
with patch("nanobot.agent.tools.shell._IS_WINDOWS", True):
env = ExecTool()._build_env()
assert set(env) == self._EXPECTED_KEYS
def test_secrets_excluded(self, monkeypatch):
monkeypatch.setenv("OPENAI_API_KEY", "sk-secret")
monkeypatch.setenv("NANOBOT_TOKEN", "tok-secret")
with patch("nanobot.agent.tools.shell._IS_WINDOWS", True):
env = ExecTool()._build_env()
assert "OPENAI_API_KEY" not in env
assert "NANOBOT_TOKEN" not in env
for v in env.values():
assert "secret" not in v.lower()
def test_path_has_sensible_default(self):
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", True),
patch.dict("os.environ", {}, clear=True),
):
env = ExecTool()._build_env()
assert "system32" in env["PATH"].lower()
def test_systemroot_forwarded(self, monkeypatch):
monkeypatch.setenv("SYSTEMROOT", r"D:\Windows")
with patch("nanobot.agent.tools.shell._IS_WINDOWS", True):
env = ExecTool()._build_env()
assert env["SYSTEMROOT"] == r"D:\Windows"
# ---------------------------------------------------------------------------
# _spawn
# ---------------------------------------------------------------------------
class TestSpawnUnix:
@pytest.mark.asyncio
async def test_uses_bash(self):
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", False),
patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_exec,
):
mock_exec.return_value = AsyncMock()
await ExecTool._spawn("echo hi", "/tmp", {"HOME": "/tmp"})
args = mock_exec.call_args[0]
assert "bash" in args[0]
assert "-l" in args
assert "-c" in args
assert "echo hi" in args
class TestSpawnWindows:
@pytest.mark.asyncio
async def test_uses_comspec_from_env(self):
env = {"COMSPEC": r"C:\Windows\system32\cmd.exe", "PATH": ""}
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", True),
patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_exec,
):
mock_exec.return_value = AsyncMock()
await ExecTool._spawn("dir", r"C:\Users", env)
args = mock_exec.call_args[0]
assert "cmd.exe" in args[0]
assert "/c" in args
assert "dir" in args
@pytest.mark.asyncio
async def test_falls_back_to_default_comspec(self):
env = {"PATH": ""}
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", True),
patch.dict("os.environ", {}, clear=True),
patch("asyncio.create_subprocess_exec", new_callable=AsyncMock) as mock_exec,
):
mock_exec.return_value = AsyncMock()
await ExecTool._spawn("dir", r"C:\Users", env)
args = mock_exec.call_args[0]
assert args[0] == "cmd.exe"
# ---------------------------------------------------------------------------
# path_append
# ---------------------------------------------------------------------------
class TestPathAppendPlatform:
@pytest.mark.asyncio
async def test_unix_injects_export(self):
"""On Unix, path_append is an export statement prepended to command."""
mock_proc = AsyncMock()
mock_proc.communicate.return_value = (b"ok", b"")
mock_proc.returncode = 0
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", False),
patch.object(ExecTool, "_spawn", return_value=mock_proc) as mock_spawn,
patch.object(ExecTool, "_guard_command", return_value=None),
):
tool = ExecTool(path_append="/opt/bin")
await tool.execute(command="ls")
spawned_cmd = mock_spawn.call_args[0][0]
assert 'export PATH="$PATH:/opt/bin"' in spawned_cmd
assert spawned_cmd.endswith("ls")
@pytest.mark.asyncio
async def test_windows_modifies_env(self):
"""On Windows, path_append is appended to PATH in the env dict."""
mock_proc = AsyncMock()
mock_proc.communicate.return_value = (b"ok", b"")
mock_proc.returncode = 0
captured_env = {}
async def capture_spawn(cmd, cwd, env):
captured_env.update(env)
return mock_proc
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", True),
patch.object(ExecTool, "_spawn", side_effect=capture_spawn),
patch.object(ExecTool, "_guard_command", return_value=None),
):
tool = ExecTool(path_append=r"C:\tools\bin")
await tool.execute(command="dir")
assert captured_env["PATH"].endswith(r";C:\tools\bin")
# ---------------------------------------------------------------------------
# sandbox
# ---------------------------------------------------------------------------
class TestSandboxPlatform:
@pytest.mark.asyncio
async def test_bwrap_skipped_on_windows(self):
"""bwrap must be silently skipped on Windows, not crash."""
mock_proc = AsyncMock()
mock_proc.communicate.return_value = (b"ok", b"")
mock_proc.returncode = 0
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", True),
patch.object(ExecTool, "_spawn", return_value=mock_proc) as mock_spawn,
patch.object(ExecTool, "_guard_command", return_value=None),
):
tool = ExecTool(sandbox="bwrap")
result = await tool.execute(command="dir")
assert "ok" in result
spawned_cmd = mock_spawn.call_args[0][0]
assert "bwrap" not in spawned_cmd
@pytest.mark.asyncio
async def test_bwrap_applied_on_unix(self):
"""On Unix, sandbox wrapping should still happen normally."""
mock_proc = AsyncMock()
mock_proc.communicate.return_value = (b"sandboxed", b"")
mock_proc.returncode = 0
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", False),
patch("nanobot.agent.tools.shell.wrap_command", return_value="bwrap -- sh -c ls") as mock_wrap,
patch.object(ExecTool, "_spawn", return_value=mock_proc) as mock_spawn,
patch.object(ExecTool, "_guard_command", return_value=None),
):
tool = ExecTool(sandbox="bwrap", working_dir="/workspace")
await tool.execute(command="ls")
mock_wrap.assert_called_once()
spawned_cmd = mock_spawn.call_args[0][0]
assert "bwrap" in spawned_cmd
# ---------------------------------------------------------------------------
# end-to-end (mocked subprocess, full execute path)
# ---------------------------------------------------------------------------
class TestExecuteEndToEnd:
@pytest.mark.asyncio
async def test_windows_full_path(self):
"""Full execute() flow on Windows: env, spawn, output formatting."""
mock_proc = AsyncMock()
mock_proc.communicate.return_value = (b"hello world\r\n", b"")
mock_proc.returncode = 0
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", True),
patch.object(ExecTool, "_spawn", return_value=mock_proc),
patch.object(ExecTool, "_guard_command", return_value=None),
):
tool = ExecTool()
result = await tool.execute(command="echo hello world")
assert "hello world" in result
assert "Exit code: 0" in result
@pytest.mark.asyncio
async def test_unix_full_path(self):
"""Full execute() flow on Unix: env, spawn, output formatting."""
mock_proc = AsyncMock()
mock_proc.communicate.return_value = (b"hello world\n", b"")
mock_proc.returncode = 0
with (
patch("nanobot.agent.tools.shell._IS_WINDOWS", False),
patch.object(ExecTool, "_spawn", return_value=mock_proc),
patch.object(ExecTool, "_guard_command", return_value=None),
):
tool = ExecTool()
result = await tool.execute(command="echo hello world")
assert "hello world" in result
assert "Exit code: 0" in result
+1 -1
View File
@@ -107,7 +107,7 @@ class TestMessageToolSuppressLogic:
async def on_progress(content: str, *, tool_hint: bool = False) -> None: async def on_progress(content: str, *, tool_hint: bool = False) -> None:
progress.append((content, tool_hint)) progress.append((content, tool_hint))
final_content, _, _ = await loop._run_agent_loop([], on_progress=on_progress) final_content, _, _, _, _ = await loop._run_agent_loop([], on_progress=on_progress)
assert final_content == "Done" assert final_content == "Done"
assert progress == [ assert progress == [
+12 -2
View File
@@ -172,6 +172,15 @@ async def test_grep_files_with_matches_supports_head_limit_and_offset(tmp_path:
(tmp_path / "src" / name).write_text("needle\n", encoding="utf-8") (tmp_path / "src" / name).write_text("needle\n", encoding="utf-8")
tool = GrepTool(workspace=tmp_path, allowed_dir=tmp_path) tool = GrepTool(workspace=tmp_path, allowed_dir=tmp_path)
# Get the full (unpaginated) list to determine the expected ordering.
full_result = await tool.execute(
pattern="needle",
path="src",
head_limit=0,
)
all_files = full_result.splitlines()
result = await tool.execute( result = await tool.execute(
pattern="needle", pattern="needle",
path="src", path="src",
@@ -179,8 +188,9 @@ async def test_grep_files_with_matches_supports_head_limit_and_offset(tmp_path:
offset=1, offset=1,
) )
lines = result.splitlines() lines = [l for l in result.splitlines() if l and not l.startswith("(pagination")]
assert lines[0] == "src/b.py" assert len(lines) == 1
assert lines[0] == all_files[1]
assert "pagination: limit=1, offset=1" in result assert "pagination: limit=1, offset=1" in result
+10 -10
View File
@@ -545,18 +545,18 @@ async def test_exec_always_returns_exit_code() -> None:
assert "hello" in result assert "hello" in result
async def test_exec_head_tail_truncation() -> None: async def test_exec_head_tail_truncation(tmp_path) -> None:
"""Long output should preserve both head and tail.""" """Long output should preserve both head and tail."""
tool = ExecTool() tool = ExecTool()
# Generate output that exceeds _MAX_OUTPUT (10_000 chars) # Generate output that exceeds _MAX_OUTPUT (10_000 chars).
# Use current interpreter (PATH may not have `python`). ExecTool uses # Use current interpreter (PATH may not have ``python``). Write the
# create_subprocess_shell: POSIX needs shlex.quote; Windows uses cmd.exe # script to a file to avoid shell-quoting issues on both POSIX and Windows.
# rules, so list2cmdline is appropriate there. script_file = tmp_path / "gen.py"
script = "print('A' * 6000 + '\\n' + 'B' * 6000)" script_file.write_text(
if sys.platform == "win32": "import sys;sys.stdout.write(chr(65)*6000);sys.stdout.write(chr(10));sys.stdout.write(chr(66)*6000)",
command = subprocess.list2cmdline([sys.executable, "-c", script]) encoding="utf-8",
else: )
command = f"{shlex.quote(sys.executable)} -c {shlex.quote(script)}" command = f"{sys.executable} {script_file}"
result = await tool.execute(command=command) result = await tool.execute(command=command)
assert "chars truncated" in result assert "chars truncated" in result
# Head portion should start with As # Head portion should start with As