mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-08 13:28:43 +03:00
Compare commits
29
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c3b55ba289 | ||
|
|
bc4cc49a59 | ||
|
|
df6f9dd171 | ||
|
|
6af81bc4a3 | ||
|
|
4a33c1392b | ||
|
|
4b0fdffe39 | ||
|
|
4fe23a01c9 | ||
|
|
973b888d39 | ||
|
|
7b1ce24600 | ||
|
|
3bece171c2 | ||
|
|
8f7ce9fef7 | ||
|
|
d327c19db0 | ||
|
|
e00dca2f84 | ||
|
|
51200a954c | ||
|
|
c121547114 | ||
|
|
dcc9c057bb | ||
|
|
586d4e2411 | ||
|
|
a4bb1923ac | ||
|
|
8d6f41e484 | ||
|
|
4962867112 | ||
|
|
473637ceff | ||
|
|
c44d4f2b2b | ||
|
|
ae27d69ecb | ||
|
|
ba38d41ad1 | ||
|
|
ad4d095080 | ||
|
|
3723cd726e | ||
|
|
8f0b653a4c | ||
|
|
5857f7fdd0 | ||
|
|
7d2c62716c |
@@ -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.
|
||||||
|
|||||||
@@ -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.
|
||||||
@@ -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.
|
||||||
@@ -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
|
||||||
@@ -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
@@ -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
@@ -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
@@ -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(
|
||||||
|
|||||||
@@ -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
@@ -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},
|
||||||
|
|||||||
@@ -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
@@ -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))
|
||||||
|
|
||||||
|
|||||||
@@ -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
@@ -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
|
|
||||||
|
|||||||
@@ -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):
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
@@ -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)"
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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"
|
||||||
@@ -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()
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -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
|
||||||
@@ -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. I’ll 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. I’ll 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: "
|
||||||
|
"• I’ll 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: "
|
||||||
|
"• I’ll 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."
|
||||||
@@ -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."""
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -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 == [
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user