Compare commits

..
Author SHA1 Message Date
chengyongru be3a42ebac fix(gateway): keep event loop responsive (NAN-33)
Move blocking filesystem, persistence, subprocess, media, and DNS work off the gateway event loop while preserving existing contracts. Add bounded cancellation and responsiveness regression coverage.
2026-08-25 10:18:13 +08:00
269 changed files with 12315 additions and 16268 deletions
+5 -5
View File
@@ -146,7 +146,7 @@ Activate it with `source .venv/bin/activate` on macOS/Linux or
python -m pip install -e . python -m pip install -e .
``` ```
After that, the normal commands are identical to a stable install. `nanobot` runs the TUI After that, the normal commands are identical to a stable install. `nanobot agent` runs the TUI
from this checkout, and `nanobot webui` rebuilds stale frontend assets automatically. A later from this checkout, and `nanobot webui` rebuilds stale frontend assets automatically. A later
`git pull --ff-only` updates the Python, TUI, and WebUI source together; rerun `git pull --ff-only` updates the Python, TUI, and WebUI source together; rerun
`python -m pip install -e .` when Python dependencies change. Contributors should also read `python -m pip install -e .` when Python dependencies change. Contributors should also read
@@ -206,13 +206,13 @@ Use `nanobot gateway --background` for the same direct entry point without keepi
**Prefer to work entirely in the terminal?** **Prefer to work entirely in the terminal?**
```bash ```bash
nanobot nanobot agent
``` ```
This opens the native terminal client with the launch directory as its workspace. It shares saved conversations and the local gateway with the WebUI. The explicit `nanobot agent` form remains available for compatibility. This opens the native terminal client with the launch directory as its workspace. It shares saved conversations and the local gateway with the WebUI.
- Type `/` to discover commands, `/sessions` to switch conversations, or `@` to mention an app, MCP server, or saved session. - Type `/` to discover commands, `/sessions` to switch conversations, or `@` to mention an app, MCP server, or saved session.
- Press `Enter` to send. While nanobot is working, `Enter` sends now and `Tab` sends after the current response. Press `Shift+Enter` to add a newline (`Ctrl+J` works in terminals that cannot distinguish modified Enter keys). - Press `Enter` to send or steer, `Tab` to queue a follow-up, and `Shift+Enter` to add a newline (`Ctrl+J` works in terminals that cannot distinguish modified Enter keys).
- Use `/detach` to leave the current task running, or start with `nanobot gateway --background` when nanobot should stay online after all local clients exit. - Use `/detach` to leave the current task running, or start with `nanobot gateway --background` when nanobot should stay online after all local clients exit.
Each launch starts a new session by default. Use `--session` to resume one and `--workspace` to choose another workspace. See the [CLI reference](./docs/cli-reference.md#agent-cli) for session branching, diffs, history, shortcuts, gateway lifecycle, and compatibility options. Each launch starts a new session by default. Use `--session` to resume one and `--workspace` to choose another workspace. See the [CLI reference](./docs/cli-reference.md#agent-cli) for session branching, diffs, history, shortcuts, gateway lifecycle, and compatibility options.
@@ -220,7 +220,7 @@ Each launch starts a new session by default. Use `--session` to resume one and `
For one request and an immediate exit, use: For one request and an immediate exit, use:
```bash ```bash
nanobot -m "Hello!" nanobot agent -m "Hello!"
``` ```
The one-shot form is useful for a quick provider check, shell scripts, and local automation. If you have not configured a model yet, run `nanobot webui` and open **Settings → Models** first. The one-shot form is useful for a quick provider check, shell scripts, and local automation. If you have not configured a model yet, run `nanobot webui` and open **Settings → Models** first.
+14 -14
View File
@@ -12,8 +12,8 @@ Use this page when you know what you want to run and need the command shape. For
| Use guided setup | `nanobot onboard --wizard` | Best when you prefer prompts over hand-editing JSON | | Use guided setup | `nanobot onboard --wizard` | Best when you prefer prompts over hand-editing JSON |
| Open the browser workbench | `nanobot webui` | Prepares local WebUI settings, starts the gateway, and opens the browser | | Open the browser workbench | `nanobot webui` | Prepares local WebUI settings, starts the gateway, and opens the browser |
| Check readiness without calling a model | `nanobot status` | Summarizes config/workspace and validates the active provider/model configuration | | Check readiness without calling a model | `nanobot status` | Summarizes config/workspace and validates the active provider/model configuration |
| Send one test message | `nanobot -m "Hello!"` | First proof that install, config, provider, model, and workspace all work | | Send one test message | `nanobot agent -m "Hello!"` | First proof that install, config, provider, model, and workspace all work |
| Chat in the terminal | `nanobot` | Interactive local chat; `nanobot agent` remains an explicit alias | | Chat in the terminal | `nanobot agent` | Interactive local chat; exit with `exit`, `/exit`, `:q`, or `Ctrl+D` |
| Run the gateway directly | `nanobot gateway` | Service/ops command for WebUI, chat apps, cron, and heartbeat | | Run the gateway directly | `nanobot gateway` | Service/ops command for WebUI, chat apps, cron, and heartbeat |
| Deliver a local trigger | `nanobot trigger <id> "message"` | Created first with `/trigger <name>` in the target chat/session | | Deliver a local trigger | `nanobot trigger <id> "message"` | Created first with `/trigger <name>` in the target chat/session |
| Serve an OpenAI-compatible API | `nanobot serve` | Starts `/v1/chat/completions`, `/v1/models`, and `/health` | | Serve an OpenAI-compatible API | `nanobot serve` | Starts `/v1/chat/completions`, `/v1/models`, and `/health` |
@@ -86,15 +86,15 @@ follow the printed WebUI **Settings → Models** or `nanobot onboard --wizard` r
| Command | Description | | Command | Description |
|---|---| |---|---|
| `nanobot -m "Hello!"` | Send one message and exit | | `nanobot agent -m "Hello!"` | Send one message and exit |
| `nanobot` | Start interactive terminal chat | | `nanobot agent` | Start interactive terminal chat |
| `nanobot --session <id>` | Use a WebSocket session key; add `--classic` for another channel | | `nanobot agent --session <id>` | Use a WebSocket session key; add `--classic` for another channel |
| `nanobot --workspace <path>` | Override workspace | | `nanobot agent --workspace <path>` | Override workspace |
| `nanobot --config <path>` | Use a specific config file | | `nanobot agent --config <path>` | Use a specific config file |
| `nanobot --classic` | Use the classic Python prompt instead of the native terminal UI | | `nanobot agent --classic` | Use the classic Python prompt instead of the native terminal UI |
| `nanobot --theme auto\|dark\|light` | Auto-detect the terminal appearance or force a TUI palette | | `nanobot agent --theme auto\|dark\|light` | Auto-detect the terminal appearance or force a TUI palette |
| `nanobot --no-markdown` | Use the classic prompt and print plain text instead of Markdown | | `nanobot agent --no-markdown` | Use the classic prompt and print plain text instead of Markdown |
| `nanobot --logs` | Use the classic prompt and show runtime logs while chatting | | `nanobot agent --logs` | Use the classic prompt and show runtime logs while chatting |
Inside the native TUI, `/sessions` switches saved conversations, `/new-chat` starts another saved Inside the native TUI, `/sessions` switches saved conversations, `/new-chat` starts another saved
conversation, and `/context` explains the compacted summary and raw session suffix available to conversation, and `/context` explains the compacted summary and raw session suffix available to
@@ -127,7 +127,7 @@ Interactive mode uses nanobot's native TypeScript terminal UI. It talks to the s
The default `--theme auto` mode paints first with the terminal's default background, probes the real foreground and background colors asynchronously, and follows supported live appearance changes. Use `--theme light` or `--theme dark` when a terminal or multiplexer does not report its colors reliably. The model preset and workspace access labels above the composer can be clicked to open their selectors; arrow keys, `Enter`, and `Esc` provide the same controls without a mouse. Access changes still pass through the gateway's local-trust and active-turn policy checks. The default `--theme auto` mode paints first with the terminal's default background, probes the real foreground and background colors asynchronously, and follows supported live appearance changes. Use `--theme light` or `--theme dark` when a terminal or multiplexer does not report its colors reliably. The model preset and workspace access labels above the composer can be clicked to open their selectors; arrow keys, `Enter`, and `Esc` provide the same controls without a mouse. Access changes still pass through the gateway's local-trust and active-turn policy checks.
`Enter` sends the current message. While nanobot is working, `Enter` sends immediately, `Tab` waits until the current response is finished, and `Option+Up` on macOS (`Alt+Up` on Windows/Linux) returns the latest waiting message to the composer. Press `Shift+Enter` to add a newline; `Ctrl+J` is the universal fallback when a terminal cannot distinguish modified Enter keys. `Alt+Enter` and `Ctrl+Enter` are also accepted when distinguishable. Use `Up`/`Down` at the composer edge to recall prompts from the current saved session. Large pastes appear as a compact placeholder in the composer but are sent unchanged. Type `/` to discover nanobot commands and terminal navigation in one palette, or type `@` to complete installed apps, configured MCP servers, and saved sessions. Use the arrow keys to choose an item and `Tab` to complete it. `/sessions` opens a searchable conversation picker, `/new-chat` preserves the current conversation and starts another one, and `/branch` forks from a completed reply. `/diff` opens a read-only unified diff for the newest turn; use `Left`/`Right` to switch edits and `Esc` to close it. The core `/new` command retains its cross-channel behavior and resets the current chat. `Ctrl+C` copies a selection, stops a running turn, clears a non-empty composer, or exits when idle. Use `PageUp`/`PageDown` to scroll, `Ctrl+Home`/`Ctrl+End` to jump to the transcript edges, and `Ctrl+O` to expand or collapse long tool traces. When you leave the bottom, the TUI shows a scrollbar and a `Ctrl+End` hint until you return. The footer reports provider token/cache usage when available. Selections copy through OSC 52 when the terminal supports it. The transcript reflows when the terminal is resized, and exiting restores the previous screen. `Enter` sends the current message. While a turn is active, `Enter` steers it immediately, `Tab` queues a visible follow-up for the next turn, and `Option+Up` on macOS (`Alt+Up` on Windows/Linux) returns the latest queued message to the composer. Press `Shift+Enter` to add a newline; `Ctrl+J` is the universal fallback when a terminal cannot distinguish modified Enter keys. `Alt+Enter` and `Ctrl+Enter` are also accepted when distinguishable. Use `Up`/`Down` at the composer edge to recall prompts from the current saved session. Large pastes appear as a compact placeholder in the composer but are sent unchanged. Type `/` to discover nanobot commands and terminal navigation in one palette, or type `@` to complete installed apps, configured MCP servers, and saved sessions. Use the arrow keys to choose an item and `Tab` to complete it. `/sessions` opens a searchable conversation picker, `/new-chat` preserves the current conversation and starts another one, and `/branch` forks from a completed reply. `/diff` opens a read-only unified diff for the newest turn; use `Left`/`Right` to switch edits and `Esc` to close it. The core `/new` command retains its cross-channel behavior and resets the current chat. `Ctrl+C` copies a selection, stops a running turn, clears a non-empty composer, or exits when idle. Use `PageUp`/`PageDown` to scroll, `Ctrl+Home`/`Ctrl+End` to jump to the transcript edges, and `Ctrl+O` to expand or collapse long tool traces. When you leave the bottom, the TUI shows a scrollbar and a `Ctrl+End` hint until you return. The footer reports provider token/cache usage when available. Selections copy through OSC 52 when the terminal supports it. The transcript reflows when the terminal is resized, and exiting restores the previous screen.
Packaged releases fetch a version-matched, checksummed terminal archive for macOS (Apple Silicon and Intel), Linux (x64 and ARM64), or Windows x64 on first use. The cache keeps the executable together with its licenses, third-party notices, source offer, relinking instructions, and corresponding TUI source. Windows ARM64 currently falls back to the classic prompt because the Bun runtime disables the FFI required by OpenTUI on that platform. Set `NANOBOT_TUI_NO_DOWNLOAD=1` or pass `--classic` to keep the Python-only path. A local source install requires Bun and runs its own `tui/` source while the original checkout remains available; it never silently falls back to a release binary. Packaged releases fetch a version-matched, checksummed terminal archive for macOS (Apple Silicon and Intel), Linux (x64 and ARM64), or Windows x64 on first use. The cache keeps the executable together with its licenses, third-party notices, source offer, relinking instructions, and corresponding TUI source. Windows ARM64 currently falls back to the classic prompt because the Bun runtime disables the FFI required by OpenTUI on that platform. Set `NANOBOT_TUI_NO_DOWNLOAD=1` or pass `--classic` to keep the Python-only path. A local source install requires Bun and runs its own `tui/` source while the original checkout remains available; it never silently falls back to a release binary.
@@ -139,7 +139,7 @@ Interactive mode exits with `exit`, `quit`, `/exit`, `/quit`, `:q`, or `Ctrl+D`.
| Command | Description | | Command | Description |
|---|---| |---|---|
| `nanobot webui` | Create config/workspace if needed, enable the local WebUI channel after confirmation, start the gateway, open `http://127.0.0.1:8765`, and follow new gateway logs | | `nanobot webui` | Create config/workspace if needed, enable the local WebUI channel after confirmation, start the gateway, and open `http://127.0.0.1:8765` |
| `nanobot webui --background` | Deprecated; prints the equivalent explicit `nanobot gateway --background` command and exits | | `nanobot webui --background` | Deprecated; prints the equivalent explicit `nanobot gateway --background` command and exits |
| `nanobot webui --dev` | Start the gateway and Vite together at `http://127.0.0.1:5173`, with live frontend updates | | `nanobot webui --dev` | Start the gateway and Vite together at `http://127.0.0.1:5173`, with live frontend updates |
| `nanobot webui --no-open` | Prepare and start the WebUI without opening a browser | | `nanobot webui --no-open` | Prepare and start the WebUI without opening a browser |
@@ -344,7 +344,7 @@ remain accepted as no-op compatibility aliases.
| Command | Description | | Command | Description |
|---|---| |---|---|
| `nanobot provider login openai-codex --set-main` | Authenticate Codex and select its current default model | | `nanobot provider login openai-codex --set-main` | Authenticate Codex and select its current default model |
| `nanobot provider login xai-grok --set-main` | Authenticate an eligible X Premium / Grok subscription and select Grok 4.6; hosted X Search is enabled for models that advertise support | | `nanobot provider login xai-grok --set-main` | Authenticate an eligible X Premium / Grok subscription and select Grok 4.5; hosted X Search is enabled for models that advertise support |
| `nanobot provider login github-copilot --set-main` | Authenticate GitHub Copilot and select its current default model | | `nanobot provider login github-copilot --set-main` | Authenticate GitHub Copilot and select its current default model |
| `nanobot provider logout openai-codex` | Remove OpenAI Codex OAuth state | | `nanobot provider logout openai-codex` | Remove OpenAI Codex OAuth state |
| `nanobot provider logout xai-grok --config <path>` | Remove the selected nanobot instance's xAI OAuth state | | `nanobot provider logout xai-grok --config <path>` | Remove the selected nanobot instance's xAI OAuth state |
+28 -25
View File
@@ -188,7 +188,7 @@ These variables are process-level switches. Set them in the same terminal, servi
| Variable | Default | Description | | Variable | Default | Description |
|----------|---------|-------------| |----------|---------|-------------|
| `NANOBOT_MAX_CONCURRENT_REQUESTS` | Unlimited | Maximum concurrently running inbound agent requests. Set a positive integer to apply a cap; unset, `0`, or a negative value means unlimited. | | `NANOBOT_MAX_CONCURRENT_REQUESTS` | `3` | Maximum concurrently running inbound agent requests. Must be an integer; set `0` or a negative value for unlimited. |
| `NANOBOT_LLM_TIMEOUT_S` | `300` | Wall-clock timeout, in seconds. Ordinary requests use this value; streaming requests use the greater of 300 seconds or twice this value. Set `0` to disable. Sustained-goal turns bypass this wall-clock cap. | | `NANOBOT_LLM_TIMEOUT_S` | `300` | Wall-clock timeout, in seconds. Ordinary requests use this value; streaming requests use the greater of 300 seconds or twice this value. Set `0` to disable. Sustained-goal turns bypass this wall-clock cap. |
| `NANOBOT_STREAM_IDLE_TIMEOUT_S` | `90` | Streaming idle timeout, in seconds, used by streaming providers. Invalid or non-positive values are ignored; values above `3600` are clamped. | | `NANOBOT_STREAM_IDLE_TIMEOUT_S` | `90` | Streaming idle timeout, in seconds, used by streaming providers. Invalid or non-positive values are ignored; values above `3600` are clamped. |
| `NANOBOT_OPENAI_COMPAT_TIMEOUT_S` | `120` | HTTP request timeout, in seconds, for OpenAI-compatible providers. Invalid or non-positive values are ignored. | | `NANOBOT_OPENAI_COMPAT_TIMEOUT_S` | `120` | HTTP request timeout, in seconds, for OpenAI-compatible providers. Invalid or non-positive values are ignored. |
@@ -729,11 +729,6 @@ Then run:
nanobot agent -m "Hello!" nanobot agent -m "Hello!"
``` ```
The WebUI model selector loads the models available to the signed-in account
from Codex's online catalog. Context-window and reasoning-effort metadata come
from that response; if discovery is unavailable, nanobot keeps a small built-in
fallback instead of emptying the selector.
Codex Fast mode can be enabled from the WebUI provider settings, or with: Codex Fast mode can be enabled from the WebUI provider settings, or with:
```json ```json
@@ -769,14 +764,11 @@ nanobot provider login xai-grok --set-main
nanobot agent -m "Hello from Grok." nanobot agent -m "Hello from Grok."
``` ```
The default model is `xai-grok/grok-4.6` with a 500,000-token context window. The default model is `xai-grok/grok-4.5` with a 500,000-token context window.
The provider reads and caches xAI's online model catalog for both WebUI model The provider reads xAI's model catalog and includes the server-hosted `x_search`
selection and runtime capabilities. Newly available models appear automatically; tool only when the selected model advertises `supportsBackendSearch`. Models
when discovery fails, the last successful catalog or built-in fallback remains without that capability continue normally without hosted X Search. When enabled,
available. The server-hosted `x_search` tool is included only when the selected searches run inside xAI's Responses API and citations arrive as inline links.
model advertises support. Models without that capability continue normally
without hosted X Search. When enabled, searches run inside xAI's Responses API
and citations arrive as inline links.
Hosted X Search is on by default to preserve this behavior. It can be turned off in the Hosted X Search is on by default to preserve this behavior. It can be turned off in the
WebUI provider settings or with `providers.xaiGrok.extraBody.tools: []`. WebUI provider settings or with `providers.xaiGrok.extraBody.tools: []`.
@@ -813,10 +805,6 @@ a nanobot update.
GitHub Copilot uses OAuth instead of API keys. Requires a [GitHub account with a plan](https://github.com/features/copilot/plans) configured. No `providers.github_copilot` block is needed in `config.json`; `nanobot provider login` stores the OAuth session outside config. GitHub Copilot uses OAuth instead of API keys. Requires a [GitHub account with a plan](https://github.com/features/copilot/plans) configured. No `providers.github_copilot` block is needed in `config.json`; `nanobot provider login` stores the OAuth session outside config.
After login, the WebUI loads the account-specific Copilot model catalog online.
Only models compatible with nanobot's current chat-completions or Responses
transport are shown.
For GitHub Enterprise / Copilot for Business, set the endpoint overrides you need before login: For GitHub Enterprise / Copilot for Business, set the endpoint overrides you need before login:
```bash ```bash
export NANOBOT_GITHUB_COPILOT_CLIENT_ID="your-enterprise-client-id" export NANOBOT_GITHUB_COPILOT_CLIENT_ID="your-enterprise-client-id"
@@ -2225,7 +2213,7 @@ The notification gate runs on a built-in system prompt. Advanced users can overr
## Subagent Concurrency ## Subagent Concurrency
By default, nanobot allows four subagents to run at the same time. Additional subagents wait for capacity instead of being rejected. Lower the limit if a local model server cannot hold multiple KV caches, or raise it when the provider can handle more parallel work: By default, nanobot only allows one spawned subagent at a time. When the limit is reached, the `spawn` tool returns an error so the agent can decide to wait or rearrange its work. This protects local LLM servers from loading multiple KV caches at once. If your provider can handle more parallel work, raise the limit:
```json ```json
{ {
@@ -2237,11 +2225,22 @@ By default, nanobot allows four subagents to run at the same time. Additional su
} }
``` ```
The deprecated `agents.defaults.failOnToolError` field is silently ignored when present in older configs. Subagents also stop immediately when one of their tools returns an execution error. That default keeps failures visible to the parent agent. If your subagent workflows use tools that can fail transiently and should be retried or worked around by the model, disable hard-stop behavior:
```json
{
"agents": {
"defaults": {
"failOnToolError": false
}
}
}
```
| Option | Default | Description | | Option | Default | Description |
|--------|---------|-------------| |--------|---------|-------------|
| `agents.defaults.maxConcurrentSubagents` | `4` | Maximum number of subagents that may run at the same time. Additional tasks wait for capacity. | | `agents.defaults.maxConcurrentSubagents` | `1` | Maximum number of spawned subagents that may run at the same time. Attempts to spawn beyond this limit return an error. |
| `agents.defaults.failOnToolError` | `true` | Stop a spawned subagent when a tool execution fails. Set to `false` to return tool errors to the subagent model so it can recover within the same run. |
## Auto Compact ## Auto Compact
@@ -2268,12 +2267,16 @@ When a user is idle for longer than a configured threshold, nanobot **proactivel
How it works: How it works:
1. **Idle detection**: On each idle tick (~1 s), checks whether an idle-session scan is due. By default, the full scan runs at most once per minute. 1. **Idle detection**: On each idle tick (~1 s), checks whether an idle-session scan is due. By default, the full scan runs at most once per minute.
2. **Background compaction**: Older context is summarized while the most recent messages remain available. 2. **Background compaction**: Idle sessions summarize the older live prefix via LLM and keep the most recent legal suffix (currently 8 messages).
3. **Session preservation**: The complete session history remains stored for later inspection and reuse. 3. **Summary injection**: When the user returns, the summary is injected as runtime context (one-shot, not persisted) alongside the retained recent suffix.
4. **Restart-safe resume**: The compacted context remains available after a process restart. 4. **Restart-safe resume**: The summary is also mirrored into session metadata so it can still be recovered after a process restart.
> [!NOTE] > [!NOTE]
> Auto compact shortens the context sent to the model without deleting the session's structured message history. > Mental model: "summarize older context, keep the freshest live turns, **and overwrite the session file with the compact form.**" It is not a full `session.clear()`, but it is a write — not a soft cursor move.
>
> Concretely, auto compact rewrites `sessions/<key>.jsonl` in place: older messages (including their structured `tool_calls` / `tool_call_id` / `reasoning_content`) are replaced by just the retained recent suffix (currently 8 messages), while the archived prefix is preserved only as a plain-text summary appended to `memory/history.jsonl` (or a `[RAW] ...` flattened dump if LLM summarization fails). The original structured JSON of those turns is no longer recoverable from the session file.
>
> This differs from the **token-driven soft consolidation** that fires when a prompt exceeds the context budget: that path only advances an internal `last_consolidated` cursor and leaves the session file untouched, so the raw tool-call trail stays on disk and can still be replayed or audited. If you rely on that trail for debugging or auditing, set `idleCompactAfterMinutes` to `0` and let only the token-driven path run.
## Timezone ## Timezone
+3 -1
View File
@@ -29,7 +29,9 @@ Memory moves through nanobot in two stages.
### Stage 1: Consolidator ### Stage 1: Consolidator
When a conversation grows large, the `Consolidator` summarizes older turns and appends the result to `memory/history.jsonl`, while keeping recent conversation available. Each summary preserves useful long-term facts and a short handoff for active work. When a conversation grows large enough to pressure the context window, nanobot does not try to carry every old message forever.
Instead, the `Consolidator` summarizes the oldest safe slice of the conversation and appends that summary to `memory/history.jsonl`.
This file is: This file is:
+5 -3
View File
@@ -4,11 +4,11 @@ Let the agent sense and adjust its own runtime state — like asking a coworker
## Why You Need It ## Why You Need It
Normal tools let the agent operate on the outside world (read/write files, search code). But the agent knows nothing about itself — it doesn't know which model it's running on, which workspace it can access, or which runtime limits apply. Normal tools let the agent operate on the outside world (read/write files, search code). But the agent knows nothing about itself — it doesn't know which model it's running on, how many iterations are left, or how many tokens it has consumed.
My tool fills this gap. With it, the agent can: My tool fills this gap. With it, the agent can:
- **Know who it is**: What model am I using? Where is my workspace? What is my per-turn iteration limit? - **Know who it is**: What model am I using? Where is my workspace? How many iterations remain?
- **Adapt on the fly**: Complex task? Expand the context window. Simple chat? Switch to a faster model. - **Adapt on the fly**: Complex task? Expand the context window. Simple chat? Switch to a faster model.
- **Remember across turns**: Store notes in your scratchpad that persist into the next conversation turn. - **Remember across turns**: Store notes in your scratchpad that persist into the next conversation turn.
@@ -44,6 +44,7 @@ my(action="check")
# workspace: PosixPath('/tmp/workspace') # workspace: PosixPath('/tmp/workspace')
# provider_retry_mode: 'standard' # provider_retry_mode: 'standard'
# max_tool_result_chars: 16000 # max_tool_result_chars: 16000
# _current_iteration: 3
# _last_usage: {'prompt_tokens': 45000, 'completion_tokens': 8000} # _last_usage: {'prompt_tokens': 45000, 'completion_tokens': 8000}
# Note: prompt_tokens is cumulative across all turns, not current context window occupancy. # Note: prompt_tokens is cumulative across all turns, not current context window occupancy.
``` ```
@@ -67,7 +68,7 @@ my(action="check", key="web_config.enable")
|----------|-----| |----------|-----|
| "What model are you using?" | `check("model")` | | "What model are you using?" | `check("model")` |
| "Which model preset is active?" | `check("model_preset")` | | "Which model preset is active?" | `check("model_preset")` |
| "What is the per-turn iteration limit?" | `check("max_iterations")` | | "How many more tool calls can you make?" | `check("max_iterations")` minus `check("_current_iteration")` |
| "How many tokens has this conversation used?" | `check("_last_usage")` — cumulative across all turns | | "How many tokens has this conversation used?" | `check("_last_usage")` — cumulative across all turns |
| "Where is your working directory?" | `check("workspace")` | | "Where is your working directory?" | `check("workspace")` |
| "Show me your full config" | `check()` | | "Show me your full config" | `check()` |
@@ -204,6 +205,7 @@ Can be checked but not set:
| Subagent manager | `subagents` | Observable, but replacing breaks the system | | Subagent manager | `subagents` | Observable, but replacing breaks the system |
| Execution config | `exec_config` | Can check sandbox/enable status, cannot change it | | Execution config | `exec_config` | Can check sandbox/enable status, cannot change it |
| Web config | `web_config` | Can check enable status, cannot change it | | Web config | `web_config` | Can check enable status, cannot change it |
| Iteration counter | `_current_iteration` | Updated by runner only |
### Sensitive field protection ### Sensitive field protection
+3 -15
View File
@@ -572,23 +572,15 @@ For OpenAI Codex:
nanobot provider login openai-codex --set-main nanobot provider login openai-codex --set-main
``` ```
The WebUI reads the account's Codex model catalog online, including current
context-window and reasoning-effort metadata. A small compatible catalog remains
available when the service cannot be reached.
For an eligible X Premium / Grok subscription: For an eligible X Premium / Grok subscription:
```bash ```bash
nanobot provider login xai-grok --set-main nanobot provider login xai-grok --set-main
``` ```
This selects `xai-grok/grok-4.6`. The WebUI model selector reads xAI's online This selects `xai-grok/grok-4.5`. The provider reads xAI's model catalog and
model catalog, so newly available subscription models appear without a nanobot exposes the hosted `x_search` tool only when the selected model advertises
release. Online metadata is cached and enriched with nanobot's curated labels; `supportsBackendSearch`; otherwise the model runs without hosted X Search.
if xAI is temporarily unavailable, nanobot uses the last successful catalog or
a small built-in fallback instead of emptying the selector. The same catalog
controls whether the provider exposes the hosted `x_search` tool; models that do
not advertise support continue without hosted X Search.
When enabled, Grok can search current X posts and return inline source links When enabled, Grok can search current X posts and return inline source links
without invoking a local nanobot tool. Credentials are stored under the without invoking a local nanobot tool. Credentials are stored under the
active instance's `auth/xai.json` (normally `~/.nanobot/auth/xai.json`), not in active instance's `auth/xai.json` (normally `~/.nanobot/auth/xai.json`), not in
@@ -607,10 +599,6 @@ For GitHub Copilot:
nanobot provider login github-copilot --set-main nanobot provider login github-copilot --set-main
``` ```
The WebUI reads the models enabled for the signed-in Copilot account. nanobot
lists entries that support its current Copilot chat-completions or Responses
transport and hides models that it cannot route safely.
Each command authenticates the selected provider and makes its current default model active. OpenAI Codex and eligible GitHub Copilot models participate in [Responses state retention](./configuration.md#responses-state-and-compaction), while native compaction remains provider-capability-specific. OAuth providers are not valid automatic fallbacks. See [`troubleshooting.md`](./troubleshooting.md#provider-and-model-problems) for proxy, headless-login, model-name, and config-key errors. Each command authenticates the selected provider and makes its current default model active. OpenAI Codex and eligible GitHub Copilot models participate in [Responses state retention](./configuration.md#responses-state-and-compaction), while native compaction remains provider-capability-specific. OAuth providers are not valid automatic fallbacks. See [`troubleshooting.md`](./troubleshooting.md#provider-and-model-problems) for proxy, headless-login, model-name, and config-key errors.
## Provider Resolution ## Provider Resolution
+6 -6
View File
@@ -103,19 +103,19 @@ Use `nanobot gateway logs`, `restart`, and `stop` to manage that background gate
If you do not want the browser or need to isolate a WebUI problem, send one message directly: If you do not want the browser or need to isolate a WebUI problem, send one message directly:
```bash ```bash
nanobot -m "Hello!" nanobot agent -m "Hello!"
``` ```
Then start an interactive terminal chat with: Then start an interactive terminal chat with:
```bash ```bash
nanobot nanobot agent
``` ```
In interactive mode, `Enter` sends and `Shift+Enter` inserts a newline (`Ctrl+J` is the In interactive mode, `Enter` sends and `Shift+Enter` inserts a newline (`Ctrl+J` is the
universal fallback). While nanobot is working, `Enter` sends immediately, `Tab` waits until the universal fallback). While a turn is running,
current response is finished, and `Option+Up` on macOS (`Alt+Up` on Windows/Linux) edits the `Enter` steers it, `Tab` queues a follow-up, and `Option+Up` on macOS (`Alt+Up` on
latest waiting message. Exit Windows/Linux) edits the latest queued message. Exit
with `exit`, `/exit`, `:q`, or `Ctrl+D`. with `exit`, `/exit`, `:q`, or `Ctrl+D`.
## Choose One Next Step ## Choose One Next Step
@@ -173,7 +173,7 @@ nanobot webui
``` ```
The source path follows current `main` and can be newer than the published package. The editable The source path follows current `main` and can be newer than the published package. The editable
install keeps Python pointed at the checkout; `nanobot` runs `tui/` with Bun, and install keeps Python pointed at the checkout; `nanobot agent` runs `tui/` with Bun, and
`nanobot webui` automatically rebuilds `webui/` when its bundled assets are stale. All normal `nanobot webui` automatically rebuilds `webui/` when its bundled assets are stale. All normal
commands remain the same as a stable install. For development details, follow commands remain the same as a stable install. For development details, follow
[`../CONTRIBUTING.md`](../CONTRIBUTING.md). [`../CONTRIBUTING.md`](../CONTRIBUTING.md).
+1 -3
View File
@@ -23,9 +23,7 @@ one is missing, starts or joins the same on-demand gateway used by the native
TUI, and opens the browser. With a fresh config, TUI, and opens the browser. With a fresh config,
it can open before a model is configured so you can finish setup in **Settings it can open before a model is configured so you can finish setup in **Settings
→ Models**. The first-run path binds the WebUI to `127.0.0.1` by default, so → Models**. The first-run path binds the WebUI to `127.0.0.1` by default, so
it is not available from other devices on your LAN. While the launcher remains it is not available from other devices on your LAN.
attached, it mirrors new log output from that exact gateway instance in the
terminal without replaying older logs.
After model setup, explicitly promote the shared gateway when you do not want to keep a client open: After model setup, explicitly promote the shared gateway when you do not want to keep a client open:
+97 -10
View File
@@ -2,7 +2,9 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Collection import asyncio
import inspect
from collections.abc import Awaitable, Collection
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Coroutine from typing import TYPE_CHECKING, Any, Callable, Coroutine
@@ -47,8 +49,26 @@ class AutoCompact:
return idle_seconds >= self._ttl * 60 return idle_seconds >= self._ttl * 60
def _has_unarchived_messages(self, key: str) -> bool: def _has_unarchived_messages(self, key: str) -> bool:
session = self.sessions.get_or_create(key) return self._session_has_unarchived_messages(self.sessions.get_or_create(key))
return session.last_archived < len(session.messages)
@staticmethod
def _session_has_unarchived_messages(session: Session) -> bool:
return session.last_consolidated < len(session.messages)
def _has_native_async_session_method(self, name: str) -> bool:
"""Check the manager's real class, not mock-generated instance attributes."""
method = inspect.getattr_static(type(self.sessions), name, None)
return inspect.iscoroutinefunction(method)
async def _list_sessions_nonblocking(self) -> list[dict[str, Any]]:
if self._has_native_async_session_method("list_sessions_async"):
return await self.sessions.list_sessions_async()
return await asyncio.to_thread(self.sessions.list_sessions)
async def _get_or_create_nonblocking(self, key: str) -> Session:
if self._has_native_async_session_method("get_or_create_async"):
return await self.sessions.get_or_create_async(key)
return await asyncio.to_thread(self.sessions.get_or_create, key)
@classmethod @classmethod
def _is_internal_session(cls, key: str) -> bool: def _is_internal_session(cls, key: str) -> bool:
@@ -79,6 +99,31 @@ class AutoCompact:
self._archiving.add(key) self._archiving.add(key)
schedule_background(self._archive(key, runtime=runtime)) schedule_background(self._archive(key, runtime=runtime))
async def check_expired_async(
self,
schedule_background: Callable[[Coroutine[Any, Any, None]], None],
resolve_runtime: Callable[[Session], Awaitable[LLMRuntime]],
active_session_keys: Collection[str] = (),
) -> None:
"""Schedule idle archival without blocking the event loop."""
now = datetime.now()
active_keys = set(active_session_keys)
for info in await self._list_sessions_nonblocking():
key = info.get("key", "")
if not key or self._is_internal_session(key) or key in self._archiving:
continue
if key in active_keys or not self._is_expired(info.get("updated_at"), now):
continue
session = await self._get_or_create_nonblocking(key)
if not self._session_has_unarchived_messages(session):
continue
try:
runtime = await resolve_runtime(session)
except (KeyError, ValueError):
continue
self._archiving.add(key)
schedule_background(self._archive_async(key, runtime=runtime))
async def _archive(self, key: str, *, runtime: LLMRuntime) -> None: async def _archive(self, key: str, *, runtime: LLMRuntime) -> None:
if self._is_internal_session(key): if self._is_internal_session(key):
self._archiving.discard(key) self._archiving.discard(key)
@@ -90,18 +135,38 @@ class AutoCompact:
max_suffix=self._RECENT_SUFFIX_MESSAGES, max_suffix=self._RECENT_SUFFIX_MESSAGES,
) )
if summary and summary != "(nothing)": if summary and summary != "(nothing)":
session = self.sessions.get_or_create(key) self._record_stored_summary(key, self.sessions.get_or_create(key))
stored = session_summary_from_metadata(
session.metadata,
fallback_last_active=session.updated_at,
)
if stored is not None:
self._summaries[key] = stored
except Exception: except Exception:
logger.exception("Auto-compact: failed for {}", key) logger.exception("Auto-compact: failed for {}", key)
finally: finally:
self._archiving.discard(key) self._archiving.discard(key)
async def _archive_async(self, key: str, *, runtime: LLMRuntime) -> None:
if self._is_internal_session(key):
self._archiving.discard(key)
return
try:
summary = await self.consolidator.compact_idle_session(
key,
runtime=runtime,
max_suffix=self._RECENT_SUFFIX_MESSAGES,
)
if summary and summary != "(nothing)":
session = await self._get_or_create_nonblocking(key)
self._record_stored_summary(key, session)
except Exception:
logger.exception("Auto-compact: failed for {}", key)
finally:
self._archiving.discard(key)
def _record_stored_summary(self, key: str, session: Session) -> None:
stored = session_summary_from_metadata(
session.metadata,
fallback_last_active=session.updated_at,
)
if stored is not None:
self._summaries[key] = stored
def prepare_session(self, session: Session, key: str) -> tuple[Session, SessionSummary | None]: def prepare_session(self, session: Session, key: str) -> tuple[Session, SessionSummary | None]:
if self._is_internal_session(key): if self._is_internal_session(key):
self._archiving.discard(key) self._archiving.discard(key)
@@ -110,6 +175,28 @@ class AutoCompact:
if key in self._archiving or self._is_expired(session.updated_at): if key in self._archiving or self._is_expired(session.updated_at):
logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving) logger.info("Auto-compact: reloading session {} (archiving={})", key, key in self._archiving)
session = self.sessions.get_or_create(key) session = self.sessions.get_or_create(key)
return self._prepared_summary(session, key)
async def prepare_session_async(
self,
session: Session,
key: str,
) -> tuple[Session, SessionSummary | None]:
"""Prepare a session without blocking on a reload."""
if self._is_internal_session(key):
self._archiving.discard(key)
self._summaries.pop(key, None)
return session, 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 = await self._get_or_create_nonblocking(key)
return self._prepared_summary(session, key)
def _prepared_summary(
self,
session: Session,
key: str,
) -> tuple[Session, SessionSummary | None]:
# Hot path: summary from in-memory dict (process hasn't restarted). # Hot path: summary from in-memory dict (process hasn't restarted).
entry = self._summaries.pop(key, None) entry = self._summaries.pop(key, None)
if entry: if entry:
+32 -2
View File
@@ -13,6 +13,19 @@ class AutomationTurnError(RuntimeError):
"""Raised when an automation turn reaches the agent and finishes with an error.""" """Raised when an automation turn reaches the agent and finishes with an error."""
class AutomationTurnAcceptedCancellation(asyncio.CancelledError):
"""Cancellation raised after an automation turn was accepted for processing.
Callers must not replay the turn: the accepted agent work now has independent
ownership and may continue after the submitting task is cancelled.
"""
def _consume_future_exception(future: asyncio.Future[object]) -> None:
if not future.cancelled():
future.exception()
async def publish_next_deferred_turn( async def publish_next_deferred_turn(
*, *,
deferred_queues: dict[str, list[InboundMessage]], deferred_queues: dict[str, list[InboundMessage]],
@@ -70,19 +83,36 @@ class AutomationTurnCoordinator:
future: asyncio.Future[OutboundMessage | None] = loop.create_future() future: asyncio.Future[OutboundMessage | None] = loop.create_future()
self._waiters[turn_id] = future self._waiters[turn_id] = future
self._pending_messages_by_turn_id[turn_id] = msg self._pending_messages_by_turn_id[turn_id] = msg
accepted = False
try: try:
if self._is_running(): if self._is_running():
await self._publish_inbound(msg) await self._publish_inbound(msg)
accepted = True
else: else:
await self._dispatch(msg) # Direct dispatch is given independent task ownership for the
# same reason as publishing to the inbound queue: once admitted,
# cancelling this submitter must not cancel and then replay the
# already-running agent turn.
dispatch_future: asyncio.Future[object] = asyncio.ensure_future(
self._dispatch(msg)
)
dispatch_future.add_done_callback(_consume_future_exception)
accepted = True
await asyncio.shield(dispatch_future)
try: try:
return await future return await future
except asyncio.CancelledError: except asyncio.CancelledError as exc:
if accepted:
raise AutomationTurnAcceptedCancellation(*exc.args) from None
raise raise
except AutomationTurnError: except AutomationTurnError:
raise raise
except Exception as exc: except Exception as exc:
raise AutomationTurnError(str(exc) or exc.__class__.__name__) from exc raise AutomationTurnError(str(exc) or exc.__class__.__name__) from exc
except asyncio.CancelledError as exc:
if accepted and not isinstance(exc, AutomationTurnAcceptedCancellation):
raise AutomationTurnAcceptedCancellation(*exc.args) from None
raise
finally: finally:
self._waiters.pop(turn_id, None) self._waiters.pop(turn_id, None)
self._pending_messages_by_turn_id.pop(turn_id, None) self._pending_messages_by_turn_id.pop(turn_id, None)
+78 -66
View File
@@ -30,7 +30,11 @@ from nanobot.security.workspace_access import WorkspaceScopeResolver
from nanobot.session.keys import last_channel_from_metadata from nanobot.session.keys import last_channel_from_metadata
from nanobot.session.manager import Session from nanobot.session.manager import Session
from nanobot.session.summary import SessionSummary from nanobot.session.summary import SessionSummary
from nanobot.utils.helpers import detect_image_mime, load_bundled_template from nanobot.utils.helpers import (
detect_image_mime,
load_bundled_template,
truncate_text_to_tokens,
)
from nanobot.utils.prompt_templates import render_template from nanobot.utils.prompt_templates import render_template
@@ -71,29 +75,14 @@ class PersistedPromptContextResolver:
return channel, scope.project_path return channel, scope.project_path
@dataclass(frozen=True, slots=True)
class TranscriptInput:
"""Raw turn inputs from which ``ContextBuilder`` assembles a transcript."""
history: list[dict[str, Any]]
current_message: str | None
media: Sequence[str] | None = None
current_role: str = "user"
session_summary: SessionSummary | None = None
runtime_context_blocks: Sequence[RuntimeContextBlock] | None = None
@property
def message_count(self) -> int:
"""Number of boundary-preserving messages in the assembled transcript."""
return 1 + len(self.history) + (self.current_message is not None)
class ContextBuilder: class ContextBuilder:
"""Builds the context (system prompt + messages) for the agent.""" """Builds the context (system prompt + messages) for the agent."""
BOOTSTRAP_FILES = ["AGENTS.md", "SOUL.md", "USER.md"] BOOTSTRAP_FILES = ["AGENTS.md", "SOUL.md", "USER.md"]
_SKIPPABLE_DEFAULTS = {"AGENTS.md", "USER.md"} _SKIPPABLE_DEFAULTS = {"AGENTS.md", "USER.md"}
_RUNTIME_CONTEXT_TAG = RUNTIME_CONTEXT_TAG _RUNTIME_CONTEXT_TAG = RUNTIME_CONTEXT_TAG
_MAX_RECENT_HISTORY = 50
_MAX_HISTORY_TOKENS = 8_000 # hard cap on recent history section size (tokens)
_RUNTIME_CONTEXT_END = RUNTIME_CONTEXT_END _RUNTIME_CONTEXT_END = RUNTIME_CONTEXT_END
def __init__(self, workspace: Path, timezone: str | None = None, disabled_skills: list[str] | None = None): def __init__(self, workspace: Path, timezone: str | None = None, disabled_skills: list[str] | None = None):
@@ -109,6 +98,9 @@ class ContextBuilder:
session_summary: SessionSummary | None = None, session_summary: SessionSummary | None = None,
workspace: Path | None = None, workspace: Path | None = None,
include_memory: bool = True, include_memory: bool = True,
include_memory_recent_history: bool = True,
session_key: str | None = None,
unified_session: bool = False,
) -> str: ) -> str:
"""Build the system prompt from identity, bootstrap files, memory, and skills.""" """Build the system prompt from identity, bootstrap files, memory, and skills."""
root = workspace or self.workspace root = workspace or self.workspace
@@ -146,6 +138,29 @@ class ContextBuilder:
if skills_summary: if skills_summary:
parts.append(render_template("agent/skills_section.md", skills_summary=skills_summary)) parts.append(render_template("agent/skills_section.md", skills_summary=skills_summary))
if include_memory_recent_history:
entries = self.memory.read_recent_history_for_prompt(
since_cursor=self.memory.get_last_dream_cursor(),
session_key=session_key,
unified_session=unified_session,
)
if entries:
capped = entries[-self._MAX_RECENT_HISTORY:]
capped = self._without_duplicate_session_summary(
capped,
session_key=session_key,
session_summary=session_summary,
)
if capped:
history_text = "\n".join(
f"- [{e['timestamp']}] {e['content']}" for e in capped
)
history_text = truncate_text_to_tokens(
history_text,
self._MAX_HISTORY_TOKENS,
)
parts.append("# Recent History\n\n" + history_text)
if session_summary: if session_summary:
parts.append( parts.append(
"[Archived Context Summary]\n\n" "[Archived Context Summary]\n\n"
@@ -155,6 +170,25 @@ class ContextBuilder:
return "\n\n---\n\n".join(parts) return "\n\n---\n\n".join(parts)
@staticmethod
def _without_duplicate_session_summary(
entries: list[dict[str, Any]],
*,
session_key: str | None,
session_summary: SessionSummary | None,
) -> list[dict[str, Any]]:
"""Drop the history entry already represented by the session summary."""
if not session_summary:
return entries
for index in range(len(entries) - 1, -1, -1):
entry = entries[index]
if (
entry.get("session_key") == session_key
and entry.get("content") == session_summary["text"]
):
return [*entries[:index], *entries[index + 1:]]
return entries
def _get_identity(self, channel: str | None = None, workspace: Path | None = None) -> str: def _get_identity(self, channel: str | None = None, workspace: Path | None = None) -> str:
"""Get the core identity section.""" """Get the core identity section."""
root = workspace or self.workspace root = workspace or self.workspace
@@ -244,68 +278,46 @@ class ContextBuilder:
runtime_context_blocks: Sequence[RuntimeContextBlock] | None = None, runtime_context_blocks: Sequence[RuntimeContextBlock] | None = None,
workspace: Path | None = None, workspace: Path | None = None,
include_memory: bool = True, include_memory: bool = True,
include_memory_recent_history: bool = True,
session_key: str | None = None,
unified_session: bool = False,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
"""Compatibility wrapper for callers that need merged adjacent roles.""" """Build the complete message list for an LLM call."""
messages = self.build_transcript(
TranscriptInput(
history=history,
current_message=current_message,
media=media,
current_role=current_role,
session_summary=session_summary,
runtime_context_blocks=runtime_context_blocks,
),
channel=channel,
workspace=workspace,
include_memory=include_memory,
)
current = messages[-1]
if len(messages) < 2 or messages[-2].get("role") != current.get("role"):
return messages
merged = dict(messages[-2])
merged["content"] = self._merge_message_content(
merged.get("content"),
current.get("content"),
)
current_meta = current.get("_meta")
if current.get("role") == "user" and isinstance(current_meta, dict):
internal_meta = dict(merged.get("_meta") or {})
internal_meta.update(cast(dict[str, Any], current_meta))
merged["_meta"] = internal_meta
return [*messages[:-2], merged]
def build_transcript(
self,
transcript: TranscriptInput,
*,
channel: str | None = None,
workspace: Path | None = None,
include_memory: bool = True,
) -> list[dict[str, Any]]:
"""Build a model transcript while preserving the fresh-turn boundary."""
root = workspace or self.workspace root = workspace or self.workspace
messages: list[dict[str, Any]] = [ messages: list[dict[str, Any]] = [
{ {
"role": "system", "role": "system",
"content": self.build_system_prompt( "content": self.build_system_prompt(
channel=channel, channel=channel,
session_summary=transcript.session_summary, session_summary=session_summary,
workspace=root, workspace=root,
include_memory=include_memory, include_memory=include_memory,
include_memory_recent_history=include_memory_recent_history,
session_key=session_key,
unified_session=unified_session,
), ),
}, },
*transcript.history, *history,
] ]
if transcript.current_message is None:
return messages
current = self.build_current_message( current = self.build_current_message(
transcript.current_message, current_message,
media=list(transcript.media) if transcript.media else None, media=media,
current_role=transcript.current_role, current_role=current_role,
runtime_context_blocks=transcript.runtime_context_blocks, runtime_context_blocks=runtime_context_blocks,
) )
if messages[-1].get("role") == current_role:
last = dict(messages[-1])
last["content"] = self._merge_message_content(
last.get("content"),
current.get("content"),
)
current_meta = current.get("_meta")
if current_role == "user" and isinstance(current_meta, dict):
internal_meta = dict(last.get("_meta") or {})
internal_meta.update(cast(dict[str, Any], current_meta))
last["_meta"] = internal_meta
messages[-1] = last
return messages
messages.append(current) messages.append(current)
return messages return messages
+134 -101
View File
@@ -13,7 +13,6 @@ from typing import TYPE_CHECKING, Any, cast
from loguru import logger from loguru import logger
from nanobot.providers.base import LLMUsage
from nanobot.utils.helpers import ( from nanobot.utils.helpers import (
estimate_message_tokens, estimate_message_tokens,
estimate_prompt_tokens_chain, estimate_prompt_tokens_chain,
@@ -28,6 +27,12 @@ if TYPE_CHECKING:
from nanobot.providers.base import LLMProvider from nanobot.providers.base import LLMProvider
SNIP_SAFETY_BUFFER = 1024 SNIP_SAFETY_BUFFER = 1024
MICROCOMPACT_MIN_CHARS = 500
INFLIGHT_COMPACT_TARGET_RATIO = 0.85
COMPACTABLE_TOOLS = frozenset({
"read_file", "exec", "grep", "find_files",
"web_search", "web_fetch", "list_dir", "list_exec_sessions",
})
# read_file is the recovery path for persisted results; exempting it prevents persist->read->persist loops. # read_file is the recovery path for persisted results; exempting it prevents persist->read->persist loops.
TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"}) TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"})
BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]" BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]"
@@ -36,27 +41,6 @@ PLACEHOLDER_TEXTS = frozenset({
}) })
class ContextWindowExceededError(RuntimeError):
"""Raised before a locally fitted request that still exceeds its budget."""
def __init__(
self,
*,
session_key: str | None,
estimated_tokens: int,
input_budget: int,
source: str,
) -> None:
self.session_key = session_key
self.estimated_tokens = estimated_tokens
self.input_budget = input_budget
self.source = source
super().__init__(
"Model input still exceeds the local context budget after request fitting "
f"for {session_key or 'default'}: {estimated_tokens}/{input_budget} via {source}"
)
def _tool_call_name_is_valid(tool_call: Any) -> bool: def _tool_call_name_is_valid(tool_call: Any) -> bool:
"""Whether a persisted OpenAI-style tool_call carries a usable name. """Whether a persisted OpenAI-style tool_call carries a usable name.
@@ -83,6 +67,7 @@ class ContextGovernanceConfig:
context_window_tokens: int | None = None context_window_tokens: int | None = None
context_block_limit: int | None = None context_block_limit: int | None = None
max_tokens: int | None = None max_tokens: int | None = None
inflight_start_index: int = 0
class ContextGovernor: class ContextGovernor:
@@ -92,85 +77,17 @@ class ContextGovernor:
self, self,
config: ContextGovernanceConfig, config: ContextGovernanceConfig,
messages: list[dict[str, Any]], messages: list[dict[str, Any]],
compacted_tool_call_ids: set[str],
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
updated = self.strip_placeholder_assistant_messages(messages) updated = self.strip_placeholder_assistant_messages(messages)
updated = self.strip_malformed_tool_calls(updated) updated = self.strip_malformed_tool_calls(updated)
updated = self.drop_orphan_tool_results(updated) updated = self.drop_orphan_tool_results(updated)
updated = self.backfill_missing_tool_results(updated) updated = self.backfill_missing_tool_results(updated)
return self.apply_tool_result_budget(config, updated) updated = self.apply_tool_result_budget(config, updated)
updated = self.compact_inflight_overflow(config, updated, compacted_tool_call_ids)
def fit_to_budget( updated = self.snip_history(config, updated)
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
*,
tool_definitions: list[dict[str, Any]] | None,
) -> list[dict[str, Any]]:
"""Fit a model-facing copy while keeping the source transcript intact."""
updated = self.snip_history(
config,
messages,
tool_definitions=tool_definitions,
force=True,
)
updated = self.drop_orphan_tool_results(updated) updated = self.drop_orphan_tool_results(updated)
updated = self.backfill_missing_tool_results(updated) return self.backfill_missing_tool_results(updated)
if not config.context_window_tokens:
return updated
budget = self.input_budget(config)
estimated, source = estimate_prompt_tokens_chain(
config.provider,
config.model,
updated,
tool_definitions,
)
if budget > 0 and estimated <= budget:
return updated
raise ContextWindowExceededError(
session_key=config.session_key,
estimated_tokens=estimated,
input_budget=budget,
source=source,
)
def fit_request(
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
usage: LLMUsage | None,
*,
usage_matches_messages: bool,
tool_definitions: list[dict[str, Any]] | None,
request_context_tokens: int | None = None,
) -> tuple[list[dict[str, Any]], bool]:
"""Fit the request when its measured or estimated input is pressured."""
if not config.context_window_tokens:
return messages, False
budget = self.input_budget(config)
if (
request_context_tokens is None
and usage_matches_messages
and usage is not None
and usage.context_tokens is not None
):
pressured = budget <= 0 or usage.context_tokens >= budget
else:
estimated, _ = estimate_prompt_tokens_chain(
config.provider,
config.model,
messages,
tool_definitions,
)
if request_context_tokens is not None:
estimated = max(estimated, request_context_tokens)
pressured = budget <= 0 or estimated >= budget
if not pressured:
return messages, False
return self.fit_to_budget(
config,
messages,
tool_definitions=tool_definitions,
), True
@staticmethod @staticmethod
def input_budget(config: ContextGovernanceConfig) -> int: def input_budget(config: ContextGovernanceConfig) -> int:
@@ -409,13 +326,71 @@ class ContextGovernor:
updated[idx]["content"] = normalized updated[idx]["content"] = normalized
return updated return updated
def compact_inflight_overflow(
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
compacted_tool_call_ids: set[str],
) -> list[dict[str, Any]]:
"""Compact in-flight tool results only when the request would overflow."""
budget = self.input_budget(config)
if budget <= 0:
return messages
tools = config.tools.get_definitions()
updated = self._apply_recorded_compactions(messages, compacted_tool_call_ids)
estimate, source = estimate_prompt_tokens_chain(
config.provider,
config.model,
updated,
tools,
)
if estimate <= budget:
return updated
target = int(budget * INFLIGHT_COMPACT_TARGET_RATIO)
candidates = self._inflight_compaction_candidates(
config,
updated,
compacted_tool_call_ids,
)
if not candidates:
return updated
for candidate_idx, (idx, tool_call_id) in enumerate(candidates):
is_newest_candidate = candidate_idx == len(candidates) - 1
if is_newest_candidate and estimate <= budget:
break
if tool_call_id in compacted_tool_call_ids:
continue
if updated is messages:
updated = [dict(m) for m in messages]
compacted_tool_call_ids.add(tool_call_id)
self._compact_tool_result_at(updated, idx)
estimate, source = estimate_prompt_tokens_chain(
config.provider,
config.model,
updated,
tools,
)
if estimate <= target:
break
logger.debug(
"In-flight context compaction for {}: prompt={} budget={} target={} via {}, ids={}",
config.session_key or "default",
estimate,
budget,
target,
source,
len(compacted_tool_call_ids),
)
return updated
def snip_history( def snip_history(
self, self,
config: ContextGovernanceConfig, config: ContextGovernanceConfig,
messages: list[dict[str, Any]], messages: list[dict[str, Any]],
*,
tool_definitions: list[dict[str, Any]] | None,
force: bool = False,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
if not messages or not config.context_window_tokens: if not messages or not config.context_window_tokens:
return messages return messages
@@ -424,13 +399,14 @@ class ContextGovernor:
if budget <= 0: if budget <= 0:
return messages return messages
tools = config.tools.get_definitions()
estimate, _ = estimate_prompt_tokens_chain( estimate, _ = estimate_prompt_tokens_chain(
config.provider, config.provider,
config.model, config.model,
messages, messages,
tool_definitions, tools,
) )
if not force and estimate <= budget: if estimate <= budget:
return messages return messages
system_messages = [dict(msg) for msg in messages if msg.get("role") == "system"] system_messages = [dict(msg) for msg in messages if msg.get("role") == "system"]
@@ -443,7 +419,7 @@ class ContextGovernor:
config.provider, config.provider,
config.model, config.model,
system_messages, system_messages,
tool_definitions, tools,
) )
remaining_budget = max(0, budget - max(system_tokens, fixed_tokens)) remaining_budget = max(0, budget - max(system_tokens, fixed_tokens))
kept: list[dict[str, Any]] = [] kept: list[dict[str, Any]] = []
@@ -458,6 +434,16 @@ class ContextGovernor:
return system_messages + self._legal_history_tail(kept, non_system) return system_messages + self._legal_history_tail(kept, non_system)
@staticmethod
def _tool_result_compaction_message(message: dict[str, Any]) -> str:
name = message.get("name", "tool")
return (
f"Error: The previous {name} result was compacted to fit context because it was too "
"large. Do not repeat the same call unchanged. Retry with a narrower path, query, "
"range, or result limit, use another tool, or tell the user the task cannot fit in "
"the available context."
)
def _legal_history_tail( def _legal_history_tail(
self, self,
kept: list[dict[str, Any]], kept: list[dict[str, Any]],
@@ -476,3 +462,50 @@ class ContextGovernor:
if messages[idx].get("role") == "user": if messages[idx].get("role") == "user":
return messages[idx:] return messages[idx:]
return [] return []
def _apply_recorded_compactions(
self,
messages: list[dict[str, Any]],
compacted_tool_call_ids: set[str],
) -> list[dict[str, Any]]:
if not compacted_tool_call_ids:
return messages
updated = messages
for idx, msg in enumerate(messages):
if msg.get("role") != "tool":
continue
tool_call_id = msg.get("tool_call_id")
if not tool_call_id or str(tool_call_id) not in compacted_tool_call_ids:
continue
compaction_message = self._tool_result_compaction_message(msg)
if msg.get("content") == compaction_message:
continue
if updated is messages:
updated = [dict(m) for m in messages]
updated[idx]["content"] = compaction_message
return updated
def _inflight_compaction_candidates(
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
compacted_tool_call_ids: set[str],
) -> list[tuple[int, str]]:
compactable: list[tuple[int, str]] = []
for idx, msg in enumerate(messages):
if idx < config.inflight_start_index:
continue
if msg.get("role") != "tool" or msg.get("name") not in COMPACTABLE_TOOLS:
continue
tool_call_id = msg.get("tool_call_id")
if not tool_call_id or str(tool_call_id) in compacted_tool_call_ids:
continue
content = msg.get("content")
if not isinstance(content, str) or len(content) < MICROCOMPACT_MIN_CHARS:
continue
compactable.append((idx, str(tool_call_id)))
return compactable
def _compact_tool_result_at(self, messages: list[dict[str, Any]], idx: int) -> None:
messages[idx]["content"] = self._tool_result_compaction_message(messages[idx])
+321 -173
View File
@@ -14,7 +14,6 @@ from collections.abc import Coroutine, Iterable, Mapping
from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress
from dataclasses import dataclass, field from dataclasses import dataclass, field
from enum import Enum, auto from enum import Enum, auto
from functools import partial
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, Awaitable, Callable, TypeVar, cast from typing import TYPE_CHECKING, Any, Awaitable, Callable, TypeVar, cast
@@ -24,24 +23,20 @@ from nanobot.agent import context as agent_context
from nanobot.agent import model_presets as preset_helpers from nanobot.agent import model_presets as preset_helpers
from nanobot.agent.autocompact import AutoCompact from nanobot.agent.autocompact import AutoCompact
from nanobot.agent.automation_turns import publish_next_deferred_turn from nanobot.agent.automation_turns import publish_next_deferred_turn
from nanobot.agent.context import ContextBuilder, PersistedPromptContextResolver, TranscriptInput from nanobot.agent.context import ContextBuilder, PersistedPromptContextResolver
from nanobot.agent.cron_turns import CronTurnCoordinator from nanobot.agent.cron_turns import CronTurnCoordinator
from nanobot.agent.hook import AgentHook, AgentTurnHookFactory from nanobot.agent.hook import AgentHook, AgentTurnHookFactory
from nanobot.agent.memory import Consolidator from nanobot.agent.memory import Consolidator
from nanobot.agent.model_runtime import ModelRuntimeResolver from nanobot.agent.model_runtime import ModelRuntimeResolver
from nanobot.agent.runner import ( from nanobot.agent.runner import _MAX_INJECTIONS_PER_TURN, AgentRunner, AgentRunSpec
_MAX_INJECTIONS_PER_TURN,
AgentRunner,
AgentRunResult,
AgentRunSpec,
)
from nanobot.agent.subagent import SubagentManager from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context from nanobot.agent.tools.context import RequestContext, bind_request_context, reset_request_context
from nanobot.agent.tools.exec_session import ExecSessionManager from nanobot.agent.tools.exec_session import ExecSessionManager
from nanobot.agent.tools.file_state import FileStateStore, bind_file_states, reset_file_states from nanobot.agent.tools.file_state import FileStateStore, bind_file_states, reset_file_states
from nanobot.agent.tools.message import capture_message_deliveries from nanobot.agent.tools.message import MessageTool
from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.registry import ToolRegistry
from nanobot.agent.tools.runtime_control import AgentRuntimeControl from nanobot.agent.tools.runtime_control import AgentRuntimeControl
from nanobot.agent.tools.self import MyTool
from nanobot.agent.turn_delivery import ( from nanobot.agent.turn_delivery import (
TurnDelivery, TurnDelivery,
TurnDeliveryFactory, TurnDeliveryFactory,
@@ -72,10 +67,12 @@ from nanobot.security.workspace_access import (
reset_workspace_scope, reset_workspace_scope,
) )
from nanobot.session import turn_continuation from nanobot.session import turn_continuation
from nanobot.session.async_compat import call_session_manager
from nanobot.session.automation_turns import automation_history_overrides from nanobot.session.automation_turns import automation_history_overrides
from nanobot.session.goal_state import ( from nanobot.session.goal_state import (
goal_state_runtime_lines, goal_state_runtime_lines,
runner_wall_llm_timeout_s, runner_wall_llm_timeout_s,
sustained_goal_active,
) )
from nanobot.session.history_visibility import HIDDEN_HISTORY_META from nanobot.session.history_visibility import HIDDEN_HISTORY_META
from nanobot.session.keys import UNIFIED_SESSION_KEY, remember_last_channel from nanobot.session.keys import UNIFIED_SESSION_KEY, remember_last_channel
@@ -116,7 +113,6 @@ if TYPE_CHECKING:
_T = TypeVar("_T") _T = TypeVar("_T")
_SUBAGENT_PROVIDER_TASK_META = "subagent_provider_task_id" _SUBAGENT_PROVIDER_TASK_META = "subagent_provider_task_id"
_SUBAGENT_TERMINAL_WAIT_SECONDS = 300.0
class TurnKind(Enum): class TurnKind(Enum):
@@ -136,7 +132,7 @@ class TurnContext:
session: Session | None = None session: Session | None = None
history: list[dict[str, Any]] = field(default_factory=list) history: list[dict[str, Any]] = field(default_factory=list)
transcript_input: TranscriptInput | None = None initial_messages: list[dict[str, Any]] = field(default_factory=list)
provider_state: ProviderConversationState | None = field(default=None, repr=False) provider_state: ProviderConversationState | None = field(default=None, repr=False)
request_context: RequestContext | None = None request_context: RequestContext | None = None
runtime_context_blocks: list[RuntimeContextBlock] = field(default_factory=list) runtime_context_blocks: list[RuntimeContextBlock] = field(default_factory=list)
@@ -145,6 +141,7 @@ class TurnContext:
final_content: str | None = None final_content: str | None = None
all_messages: list[dict[str, Any]] = field(default_factory=list) all_messages: list[dict[str, Any]] = field(default_factory=list)
stop_reason: str = "" stop_reason: str = ""
had_injections: bool = False
streamed_content: bool = False streamed_content: bool = False
input_persisted_early: bool = False input_persisted_early: bool = False
@@ -199,10 +196,19 @@ class AgentLoop:
5. Sends responses back 5. Sends responses back
""" """
@property
def current_iteration(self) -> int:
return self._current_iteration
@property @property
def tool_names(self) -> list[str]: def tool_names(self) -> list[str]:
return self.tools.tool_names return self.tools.tool_names
@property
def last_usage(self) -> LLMUsage | None:
"""Latest aggregate usage exposed through the runtime-control snapshot."""
return self._last_usage
@property @property
def provider(self) -> LLMProvider: def provider(self) -> LLMProvider:
"""Provider selected for future turn admissions.""" """Provider selected for future turn admissions."""
@@ -265,6 +271,7 @@ class AgentLoop:
context_window_tokens: int | None = None, context_window_tokens: int | None = None,
context_block_limit: int | None = None, context_block_limit: int | None = None,
max_tool_result_chars: int | None = None, max_tool_result_chars: int | None = None,
fail_on_tool_error: bool | None = None,
provider_retry_mode: str = "standard", provider_retry_mode: str = "standard",
tool_hint_max_length: int | None = None, tool_hint_max_length: int | None = None,
cron_service: CronService | None = None, cron_service: CronService | None = None,
@@ -274,6 +281,7 @@ class AgentLoop:
channels_config: ChannelsConfig | None = None, channels_config: ChannelsConfig | None = None,
timezone: str | None = None, timezone: str | None = None,
session_ttl_minutes: int = 0, session_ttl_minutes: int = 0,
consolidation_ratio: float = 0.5,
hooks: list[AgentHook] | None = None, hooks: list[AgentHook] | None = None,
hook_factories: list[AgentTurnHookFactory] | None = None, hook_factories: list[AgentTurnHookFactory] | None = None,
unified_session: bool = False, unified_session: bool = False,
@@ -372,6 +380,7 @@ class AgentLoop:
default_restrict_to_workspace=restrict_to_workspace, default_restrict_to_workspace=restrict_to_workspace,
) )
self._start_time = time.time() self._start_time = time.time()
self._last_usage: LLMUsage | None = None
self._extra_hooks: list[AgentHook] = hooks or [] self._extra_hooks: list[AgentHook] = hooks or []
self._hook_factories: list[AgentTurnHookFactory] = hook_factories or [] self._hook_factories: list[AgentTurnHookFactory] = hook_factories or []
@@ -396,6 +405,7 @@ class AgentLoop:
disabled_skills=disabled_skills, disabled_skills=disabled_skills,
max_iterations=self.max_iterations, max_iterations=self.max_iterations,
max_concurrent_subagents=max_concurrent_subagents, max_concurrent_subagents=max_concurrent_subagents,
fail_on_tool_error=fail_on_tool_error,
llm_wall_timeout_for_session=lambda sk: runner_wall_llm_timeout_s(self.sessions, sk), llm_wall_timeout_for_session=lambda sk: runner_wall_llm_timeout_s(self.sessions, sk),
) )
self._unified_session = unified_session self._unified_session = unified_session
@@ -430,8 +440,8 @@ class AgentLoop:
("cron", self._cron_turns), ("cron", self._cron_turns),
("local trigger", self._local_trigger_turns), ("local trigger", self._local_trigger_turns),
) )
# NANOBOT_MAX_CONCURRENT_REQUESTS: unset or <=0 means unlimited. # NANOBOT_MAX_CONCURRENT_REQUESTS: <=0 means unlimited; default 3.
_max = int(os.environ.get("NANOBOT_MAX_CONCURRENT_REQUESTS", "0")) _max = int(os.environ.get("NANOBOT_MAX_CONCURRENT_REQUESTS", "3"))
self._concurrency_gate: asyncio.Semaphore | None = ( self._concurrency_gate: asyncio.Semaphore | None = (
asyncio.Semaphore(_max) if _max > 0 else None asyncio.Semaphore(_max) if _max > 0 else None
) )
@@ -444,6 +454,8 @@ class AgentLoop:
workspace_scopes=self.workspace_scopes, workspace_scopes=self.workspace_scopes,
unified_session=unified_session, unified_session=unified_session,
), ),
consolidation_ratio=consolidation_ratio,
unified_session=unified_session,
) )
self.auto_compact = AutoCompact( self.auto_compact = AutoCompact(
sessions=self.sessions, sessions=self.sessions,
@@ -455,6 +467,7 @@ class AgentLoop:
if model_preset: if model_preset:
self.set_model_preset(model_preset, publish_update=False) self.set_model_preset(model_preset, publish_update=False)
self._register_default_tools(provider_snapshot_loader=provider_snapshot_loader) self._register_default_tools(provider_snapshot_loader=provider_snapshot_loader)
self._current_iteration: int = 0
self.commands = CommandRouter() self.commands = CommandRouter()
register_builtin_commands(self.commands) register_builtin_commands(self.commands)
@@ -506,6 +519,7 @@ class AgentLoop:
context_window_tokens=context_window_tokens, context_window_tokens=context_window_tokens,
context_block_limit=defaults.context_block_limit, context_block_limit=defaults.context_block_limit,
max_tool_result_chars=defaults.max_tool_result_chars, max_tool_result_chars=defaults.max_tool_result_chars,
fail_on_tool_error=defaults.fail_on_tool_error,
provider_retry_mode=defaults.provider_retry_mode, provider_retry_mode=defaults.provider_retry_mode,
tool_hint_max_length=defaults.tool_hint_max_length, tool_hint_max_length=defaults.tool_hint_max_length,
restrict_to_workspace=config.tools.restrict_to_workspace, restrict_to_workspace=config.tools.restrict_to_workspace,
@@ -515,6 +529,7 @@ class AgentLoop:
disabled_skills=defaults.disabled_skills, disabled_skills=defaults.disabled_skills,
session_ttl_minutes=defaults.session_ttl_minutes, session_ttl_minutes=defaults.session_ttl_minutes,
idle_compact_check_interval_seconds=defaults.idle_compact_check_interval_seconds, idle_compact_check_interval_seconds=defaults.idle_compact_check_interval_seconds,
consolidation_ratio=defaults.consolidation_ratio,
tools_config=config.tools, tools_config=config.tools,
model_presets=preset_helpers.configured_model_presets(config), model_presets=preset_helpers.configured_model_presets(config),
model_preset=defaults.model_preset, model_preset=defaults.model_preset,
@@ -526,6 +541,33 @@ class AgentLoop:
**extra, **extra,
) )
async def _get_or_create_session(self, key: str) -> Session:
"""Use native async session loading, with a compatibility fallback."""
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
"""Use native async session saving, with a compatibility fallback."""
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
async def _save_runtime_checkpoint(self, session: Session) -> None:
"""Use native async checkpoint saving, with a compatibility fallback."""
await call_session_manager(
self.sessions,
"save_runtime_checkpoint_async",
self.sessions.save_runtime_checkpoint,
session,
)
def _sync_subagent_runtime_limits(self) -> None: def _sync_subagent_runtime_limits(self) -> None:
"""Keep subagent runtime limits aligned with mutable loop settings.""" """Keep subagent runtime limits aligned with mutable loop settings."""
self.subagents.max_iterations = self.max_iterations self.subagents.max_iterations = self.max_iterations
@@ -565,6 +607,30 @@ class AgentLoop:
self.sessions.save(session) self.sessions.save(session)
return self.llm_runtime() return self.llm_runtime()
async def runtime_for_session_async(
self,
session: Session,
*,
recover_removed: bool = True,
) -> LLMRuntime:
"""Resolve a session runtime without blocking on recovery persistence."""
name = model_preset_from_metadata(session.metadata)
if name is None:
return self.llm_runtime()
try:
return self.runtime_resolver.resolve_preset(name)
except KeyError:
if not recover_removed or name in self.runtime_resolver.model_presets:
raise
logger.warning(
"Session '{}' references removed model preset '{}'; falling back to default",
session.key,
name,
)
session.metadata.pop(SESSION_MODEL_PRESET_METADATA_KEY, None)
await self._save_session(session)
return self.llm_runtime()
def set_session_model_preset( def set_session_model_preset(
self, self,
session_key: str, session_key: str,
@@ -577,6 +643,18 @@ class AgentLoop:
self.sessions.save(session) self.sessions.save(session)
return runtime return runtime
async def set_session_model_preset_async(
self,
session_key: str,
name: str,
) -> LLMRuntime:
"""Validate and persist one session's preset selection without blocking."""
runtime = self.runtime_resolver.resolve_preset(name)
session = await self._get_or_create_session(session_key)
session.metadata[SESSION_MODEL_PRESET_METADATA_KEY] = runtime.model_preset
await self._save_session(session)
return runtime
def _publish_runtime_selection( def _publish_runtime_selection(
self, self,
runtime: LLMRuntime, runtime: LLMRuntime,
@@ -639,11 +717,20 @@ class AgentLoop:
timezone=self.context.timezone or "UTC", timezone=self.context.timezone or "UTC",
workspace_sandbox=self.workspace_scopes.sandbox_status, workspace_sandbox=self.workspace_scopes.sandbox_status,
runtime_events=self.runtime_events, runtime_events=self.runtime_events,
runtime_control=AgentRuntimeControl(self),
) )
loader = ToolLoader() loader = ToolLoader()
registered = loader.load(ctx, self.tools) registered = loader.load(ctx, self.tools)
# MyTool receives only the explicit runtime-control capability.
if self.tools_config.my.enable:
self.tools.register(
MyTool(
runtime_control=AgentRuntimeControl(self),
modify_allowed=self.tools_config.my.allow_set,
)
)
registered.append("my")
logger.info("Registered {} tools: {}", len(registered), registered) logger.info("Registered {} tools: {}", len(registered), registered)
def register_runtime_context_provider( def register_runtime_context_provider(
@@ -680,17 +767,14 @@ class AgentLoop:
session_key=session_key, session_key=session_key,
) )
def _persist_user_message_early( def _stage_user_message_early(
self, self,
msg: InboundMessage, msg: InboundMessage,
session: Session, session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None, runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any, **kwargs: Any,
) -> bool: ) -> bool:
"""Persist the triggering user message before the turn starts. """Add the triggering user message and recovery markers in memory."""
Returns True if the message was persisted.
"""
if not turn_continuation.should_persist_user_message(msg.metadata): if not turn_continuation.should_persist_user_message(msg.metadata):
return False return False
media_paths = [ media_paths = [
@@ -719,19 +803,61 @@ class AgentLoop:
followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY) followup_id = msg.metadata.get(PENDING_FOLLOWUP_ID_KEY)
if isinstance(followup_id, str) and followup_id: if isinstance(followup_id, str) and followup_id:
acknowledge_pending_followups(session, [followup_id]) acknowledge_pending_followups(session, [followup_id])
self.sessions.save(session)
return True return True
return False return False
def _build_transcript_input(self, ctx: TurnContext) -> TranscriptInput: def _persist_user_message_early(
"""Capture the persisted history and fresh input as separate transcript parts.""" self,
msg: InboundMessage,
session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any,
) -> bool:
"""Synchronously persist the user message for compatibility callers."""
persisted = self._stage_user_message_early(
msg,
session,
runtime_context_blocks,
**kwargs,
)
if persisted:
self.sessions.save(session)
return persisted
async def _persist_user_message_early_async(
self,
msg: InboundMessage,
session: Session,
runtime_context_blocks: list[RuntimeContextBlock] | None = None,
**kwargs: Any,
) -> bool:
"""Persist the user message without blocking the event loop."""
persisted = self._stage_user_message_early(
msg,
session,
runtime_context_blocks,
**kwargs,
)
if persisted:
await self._save_session(session)
return persisted
def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]:
"""Build the initial message list for the LLM turn."""
assert ctx.session is not None assert ctx.session is not None
return TranscriptInput( scope = self.workspace_scopes.for_message(ctx.msg, ctx.session.metadata)
return self.context.build_messages(
history=ctx.history, history=ctx.history,
current_message=ctx.msg.content, current_message=ctx.msg.content,
media=ctx.msg.media if ctx.kind is TurnKind.USER and ctx.msg.media else None, media=ctx.msg.media if ctx.kind is TurnKind.USER and ctx.msg.media else None,
channel=ctx.delivery.route.channel,
session_summary=ctx.pending_summary, session_summary=ctx.pending_summary,
workspace=scope.project_path,
runtime_context_blocks=ctx.runtime_context_blocks, runtime_context_blocks=ctx.runtime_context_blocks,
include_memory=ctx.session.policy.persist,
include_memory_recent_history=not ctx.ephemeral,
session_key=ctx.session.key,
unified_session=self._unified_session,
) )
def _request_context_for_turn(self, ctx: TurnContext) -> RequestContext: def _request_context_for_turn(self, ctx: TurnContext) -> RequestContext:
@@ -813,7 +939,7 @@ class AgentLoop:
if tool is None: if tool is None:
content = "Shell execution is disabled in this nanobot configuration." content = "Shell execution is disabled in this nanobot configuration."
else: else:
session = ctx.session or self.sessions.get_or_create(ctx.key) session = ctx.session or await AgentLoop._get_or_create_session(self, ctx.key)
scope = self.workspace_scopes.for_turn( scope = self.workspace_scopes.for_turn(
channel=ctx.msg.channel, channel=ctx.msg.channel,
message_metadata=metadata, message_metadata=metadata,
@@ -922,7 +1048,7 @@ class AgentLoop:
async def _run_agent_loop( async def _run_agent_loop(
self, self,
transcript_input: TranscriptInput, initial_messages: list[dict[str, Any]],
on_progress: Callable[..., Awaitable[None]] | None = None, on_progress: Callable[..., 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,
@@ -930,6 +1056,12 @@ class AgentLoop:
*, *,
runtime: LLMRuntime, runtime: LLMRuntime,
session: Session | None = None, session: Session | None = None,
channel: str = "cli",
chat_id: str = "direct",
message_id: str | None = None,
metadata: dict[str, Any] | None = None,
session_key: str | None = None,
original_user_text: str | None = None,
pending_queue: asyncio.Queue[InboundMessage] | None = None, pending_queue: asyncio.Queue[InboundMessage] | None = None,
ephemeral: bool = False, ephemeral: bool = False,
run_extra_hooks_for_ephemeral: bool = False, run_extra_hooks_for_ephemeral: bool = False,
@@ -939,7 +1071,7 @@ class AgentLoop:
tools: ToolRegistry | None = None, tools: ToolRegistry | None = None,
request_context: RequestContext | None = None, request_context: RequestContext | None = None,
provider_state: ProviderConversationState | None = None, provider_state: ProviderConversationState | None = None,
) -> AgentRunResult: ) -> tuple[str | None, list[str], list[dict[str, Any]], 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.
@@ -947,7 +1079,7 @@ class AgentLoop:
``resuming=True`` means the active turn continues. ``merge_next=True`` means ``resuming=True`` means the active turn continues. ``merge_next=True`` means
the next text segment belongs to the same user-visible assistant message. the next text segment belongs to the same user-visible assistant message.
Returns the complete result produced by ``AgentRunner``. Returns (final_content, tools_used, messages, stop_reason, had_injections).
""" """
self._sync_subagent_runtime_limits() self._sync_subagent_runtime_limits()
@@ -965,14 +1097,17 @@ class AgentLoop:
public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = ( public_payload[self._PROVIDER_STATE_CHECKPOINT_VERSION_KEY] = (
self._PROVIDER_STATE_CHECKPOINT_VERSION self._PROVIDER_STATE_CHECKPOINT_VERSION
) )
self._set_runtime_checkpoint(session, public_payload) await self._set_runtime_checkpoint_async(session, public_payload)
async def _drain_pending( async def _drain_pending(*, limit: int = _MAX_INJECTIONS_PER_TURN) -> list[dict[str, Any]]:
*, """Drain follow-up messages from the pending queue.
limit: int = _MAX_INJECTIONS_PER_TURN,
first_msg: InboundMessage | None = None, When no messages are immediately available but sub-agents
) -> list[dict[str, Any]]: spawned in this dispatch are still running, blocks until at
"""Drain only messages that are already available.""" least one result arrives (or timeout). This keeps the runner
loop alive so subsequent sub-agent completions are consumed
in-order rather than dispatched separately.
"""
if pending_queue is None: if pending_queue is None:
return [] return []
@@ -1044,77 +1179,52 @@ class AgentLoop:
return row return row
items: list[dict[str, Any]] = [] items: list[dict[str, Any]] = []
if first_msg is not None:
items.append(await _to_user_message(first_msg))
while len(items) < limit: while len(items) < limit:
try: try:
items.append(await _to_user_message(pending_queue.get_nowait())) items.append(await _to_user_message(pending_queue.get_nowait()))
except asyncio.QueueEmpty: except asyncio.QueueEmpty:
break break
# Block if nothing drained but sub-agents spawned in this dispatch
# are still running. Keeps the runner loop alive so subsequent
# completions are injected in-order rather than dispatched separately.
if (not items
and session is not None
and self.subagents.get_running_count_by_session(session.key) > 0):
try:
msg = await asyncio.wait_for(pending_queue.get(), timeout=300)
except asyncio.TimeoutError:
logger.warning(
"Timeout waiting for sub-agent completion in session {}",
session.key,
)
return items
items.append(await _to_user_message(msg))
while len(items) < limit:
try:
items.append(await _to_user_message(pending_queue.get_nowait()))
except asyncio.QueueEmpty:
break
return items return items
terminal_wait_deadline: float | None = None active_session_key = session.key if session else session_key
async def _wait_for_pending(
*,
limit: int = _MAX_INJECTIONS_PER_TURN,
) -> list[dict[str, Any]]:
"""Wait for a pending result only when the runner is ready to exit."""
nonlocal terminal_wait_deadline
items = await _drain_pending(limit=limit)
if (
items
or pending_queue is None
or session is None
or self.subagents.get_running_count_by_session(session.key) == 0
):
return items
now = asyncio.get_running_loop().time()
if terminal_wait_deadline is None:
terminal_wait_deadline = now + _SUBAGENT_TERMINAL_WAIT_SECONDS
remaining = terminal_wait_deadline - now
if remaining <= 0:
return []
try:
msg = await asyncio.wait_for(pending_queue.get(), timeout=remaining)
except asyncio.TimeoutError:
logger.warning(
"Timeout waiting for sub-agent completion before session {} exits",
session.key,
)
return []
return await _drain_pending(limit=limit, first_msg=msg)
request_ctx = request_context or RequestContext(
channel="cli",
chat_id="direct",
session_key=session.key if session is not None else None,
runtime=runtime,
)
active_session_key = session.key if session else request_ctx.session_key
request_metadata = request_ctx.metadata
effective_scope = self.workspace_scopes.for_turn( effective_scope = self.workspace_scopes.for_turn(
channel=request_ctx.channel, channel=channel,
message_metadata=request_metadata, message_metadata=metadata,
session_metadata=session.metadata if session is not None else None, session_metadata=session.metadata if session is not None else None,
) )
transcript_builder = partial(
self.context.build_transcript,
channel=request_ctx.channel,
workspace=effective_scope.project_path,
include_memory=session.policy.persist if session is not None else True,
)
if request_context is None:
request_ctx = dataclasses.replace(
request_ctx,
workspace=effective_scope.project_path,
)
effective_tools = tools or self.tools effective_tools = tools or self.tools
request_ctx = request_context or RequestContext(
channel=channel,
chat_id=chat_id,
message_id=message_id,
session_key=active_session_key,
original_user_text=original_user_text,
runtime=runtime,
metadata=dict(metadata or {}),
workspace=effective_scope.project_path,
)
file_state_token = bind_file_states(self._file_state_store.for_session(active_session_key)) file_state_token = bind_file_states(self._file_state_store.for_session(active_session_key))
request_token = bind_request_context(request_ctx) request_token = bind_request_context(request_ctx)
workspace_token = bind_workspace_scope(effective_scope) workspace_token = bind_workspace_scope(effective_scope)
@@ -1139,14 +1249,15 @@ class AgentLoop:
on_progress=on_progress, on_progress=on_progress,
on_stream=on_stream, on_stream=on_stream,
on_stream_end=on_stream_end, on_stream_end=on_stream_end,
channel=request_ctx.channel, channel=channel,
chat_id=request_ctx.chat_id, chat_id=chat_id,
message_id=request_ctx.message_id, message_id=message_id,
metadata=request_metadata, metadata=metadata,
attributes=dict(request_ctx.attributes), attributes=dict(request_ctx.attributes),
session_key=active_session_key, session_key=active_session_key,
workspace=effective_scope.project_path, workspace=effective_scope.project_path,
tool_hint_max_length=self.tool_hint_max_length, tool_hint_max_length=self.tool_hint_max_length,
on_iteration=lambda iteration: setattr(self, "_current_iteration", iteration),
registered_hook_factories=self._hook_factories, registered_hook_factories=self._hook_factories,
turn_hook_factories=list(hook_factories or []), turn_hook_factories=list(hook_factories or []),
registered_hooks=self._extra_hooks, registered_hooks=self._extra_hooks,
@@ -1155,42 +1266,43 @@ class AgentLoop:
run_extra_hooks_for_ephemeral=run_extra_hooks_for_ephemeral, run_extra_hooks_for_ephemeral=run_extra_hooks_for_ephemeral,
)) ))
result = await self.runner.run(AgentRunSpec( result = await self.runner.run(AgentRunSpec(
initial_messages=None, initial_messages=initial_messages,
tools=effective_tools, tools=effective_tools,
runtime=runtime, runtime=runtime,
max_iterations=self.max_iterations, max_iterations=self.max_iterations,
max_tool_result_chars=self.max_tool_result_chars, max_tool_result_chars=self.max_tool_result_chars,
transcript_input=transcript_input,
transcript_builder=transcript_builder,
hook=hook, hook=hook,
error_message="Sorry, I encountered an error calling the AI model.",
concurrent_tools=True, concurrent_tools=True,
workspace=effective_scope.project_path, workspace=effective_scope.project_path,
session_key=session.key if session else None, session_key=session.key if session else None,
context_block_limit=self.context_block_limit, context_block_limit=self.context_block_limit,
provider_retry_mode=self.provider_retry_mode, provider_retry_mode=self.provider_retry_mode,
progress_callback=on_progress,
stream_progress_deltas=on_stream is not None,
retry_wait_callback=on_retry_wait, retry_wait_callback=on_retry_wait,
checkpoint_callback=_checkpoint, checkpoint_callback=_checkpoint,
injection_callback=_drain_pending, injection_callback=_drain_pending,
terminal_injection_callback=_wait_for_pending,
# Sustained goals may legitimately exceed NANOBOT_LLM_TIMEOUT_S; idle stall # Sustained goals may legitimately exceed NANOBOT_LLM_TIMEOUT_S; idle stall
# is still capped by NANOBOT_STREAM_IDLE_TIMEOUT_S in streaming providers. # is still capped by NANOBOT_STREAM_IDLE_TIMEOUT_S in streaming providers.
llm_timeout_s=runner_wall_llm_timeout_s( llm_timeout_s=runner_wall_llm_timeout_s(
self.sessions, self.sessions,
session.key if session is not None else request_ctx.session_key, session.key if session is not None else session_key,
metadata=session_metadata, metadata=session_metadata,
message_metadata=request_metadata, message_metadata=metadata,
), ),
continuation_callback=_goal_continue, goal_active_predicate=lambda: sustained_goal_active(session.metadata) if session is not None else False,
goal_continue_message=_goal_continue,
finalize_on_max_iterations=turn_continuation.should_finalize_on_max_iterations( finalize_on_max_iterations=turn_continuation.should_finalize_on_max_iterations(
pending_queue_available=pending_queue is not None and session is not None, pending_queue_available=pending_queue is not None and session is not None,
session_metadata=session_metadata, session_metadata=session_metadata,
message_metadata=request_metadata, message_metadata=metadata,
), ),
provider_state=provider_state, provider_state=provider_state,
llm_usage_source=source_from_request( llm_usage_source=source_from_request(
active_session_key, active_session_key,
channel=request_ctx.channel, channel=channel,
metadata=request_metadata, metadata=metadata,
), ),
)) ))
finally: finally:
@@ -1198,6 +1310,7 @@ class AgentLoop:
reset_workspace_scope(workspace_token) reset_workspace_scope(workspace_token)
reset_request_context(request_token) reset_request_context(request_token)
reset_file_states(file_state_token) reset_file_states(file_state_token)
self._last_usage = result.usage
if session is not None and not ephemeral: if session is not None and not ephemeral:
session.provider_state = result.provider_state session.provider_state = result.provider_state
if result.stop_reason == "max_iterations": if result.stop_reason == "max_iterations":
@@ -1206,7 +1319,7 @@ class AgentLoop:
stop_reason=result.stop_reason, stop_reason=result.stop_reason,
pending_queue_available=pending_queue is not None and session is not None, pending_queue_available=pending_queue is not None and session is not None,
session_metadata=session_metadata, session_metadata=session_metadata,
message_metadata=request_metadata, message_metadata=metadata,
) )
# Push final content through stream so streaming channels (e.g. Feishu) # Push final content through stream so streaming channels (e.g. Feishu)
# update the card instead of leaving it empty. # update the card instead of leaving it empty.
@@ -1220,20 +1333,35 @@ class AgentLoop:
await on_stream_end(resuming=False) await on_stream_end(resuming=False)
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 return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
def _check_expired_sessions_if_due(self) -> None: def _idle_compact_scan_due(self) -> bool:
"""Scan idle sessions no more often than the configured interval."""
now = time.monotonic() now = time.monotonic()
if now < self._next_idle_compact_check_at: if now < self._next_idle_compact_check_at:
return return False
self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s
return True
def _check_expired_sessions_if_due(self) -> None:
"""Synchronously scan idle sessions for compatibility with direct callers."""
if not self._idle_compact_scan_due():
return
self.auto_compact.check_expired( self.auto_compact.check_expired(
self.schedule_background, self.schedule_background,
self.runtime_for_session, self.runtime_for_session,
active_session_keys=self._pending_queues.keys(), active_session_keys=self._pending_queues.keys(),
) )
async def _check_expired_sessions_if_due_async(self) -> None:
"""Scan idle sessions without blocking the event loop."""
if not self._idle_compact_scan_due():
return
await self.auto_compact.check_expired_async(
self.schedule_background,
self.runtime_for_session_async,
active_session_keys=self._pending_queues.keys(),
)
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."""
self._running = True self._running = True
@@ -1244,7 +1372,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._check_expired_sessions_if_due() await self._check_expired_sessions_if_due_async()
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.
@@ -1322,7 +1450,7 @@ class AgentLoop:
) )
continue continue
pending_msg = routed_msg pending_msg = routed_msg
session = self.sessions.get_or_create(effective_key) session = await self._get_or_create_session(effective_key)
followup_id = record_pending_followup(session, pending_msg) followup_id = record_pending_followup(session, pending_msg)
if followup_id is not None: if followup_id is not None:
pending_msg = dataclasses.replace( pending_msg = dataclasses.replace(
@@ -1332,7 +1460,7 @@ class AgentLoop:
PENDING_FOLLOWUP_ID_KEY: followup_id, PENDING_FOLLOWUP_ID_KEY: followup_id,
}, },
) )
self.sessions.save(session) await self._save_session(session)
try: try:
self._pending_queues[effective_key].put_nowait(pending_msg) self._pending_queues[effective_key].put_nowait(pending_msg)
except asyncio.QueueFull: except asyncio.QueueFull:
@@ -1443,10 +1571,10 @@ class AgentLoop:
raise raise
try: try:
key = self._effective_session_key(msg) key = self._effective_session_key(msg)
session = self.sessions.get_or_create(key) session = await self._get_or_create_session(key)
if restore_runtime_checkpoint(session): if self._restore_runtime_checkpoint(session):
self._clear_pending_user_turn(session) self._clear_pending_user_turn(session)
self.sessions.save(session) await self._save_session(session)
logger.info( logger.info(
"Restored partial context for cancelled session {}", "Restored partial context for cancelled session {}",
key, key,
@@ -1720,12 +1848,18 @@ class AgentLoop:
msg: InboundMessage, msg: InboundMessage,
final_content: str, final_content: str,
stop_reason: str, stop_reason: str,
had_injections: bool,
streamed_content: bool, streamed_content: bool,
*, *,
log_content: bool = True, log_content: bool = True,
turn_latency_ms: int | None = None, turn_latency_ms: int | None = None,
) -> OutboundMessage | None: ) -> OutboundMessage | None:
"""Assemble the final outbound message from turn results.""" """Assemble the final outbound message from turn results."""
# MessageTool suppression
if (mt := self.tools.get("message")) and isinstance(mt, MessageTool) and mt._sent_in_turn:
if not had_injections or stop_reason == "empty_final_response":
return None
if log_content: if log_content:
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)
@@ -1765,7 +1899,7 @@ class AgentLoop:
if ctx.session is None: if ctx.session is None:
raise RuntimeError("required session is not active") raise RuntimeError("required session is not active")
else: else:
ctx.session = self.sessions.get_or_create(ctx.session_key) ctx.session = await self._get_or_create_session(ctx.session_key)
session = ctx.session session = ctx.session
ctx.ephemeral = ctx.ephemeral or not session.policy.persist ctx.ephemeral = ctx.ephemeral or not session.policy.persist
tools = ctx.tools or self.tools tools = ctx.tools or self.tools
@@ -1795,17 +1929,17 @@ class AgentLoop:
if ctx.kind is TurnKind.USER: if ctx.kind is TurnKind.USER:
self.workspace_scopes.persist_message_scope(session, msg) self.workspace_scopes.persist_message_scope(session, msg)
if restore_runtime_checkpoint(session): if self._restore_runtime_checkpoint(session):
self.sessions.save(session) await self._save_session(session)
if ( if (
RECOVERY_INBOUND_METADATA_KEY not in msg.metadata RECOVERY_INBOUND_METADATA_KEY not in msg.metadata
and restore_pending_interruption(session) and restore_pending_interruption(session)
): ):
self.sessions.save(session) await self._save_session(session)
async def _compact_session(self, ctx: TurnContext) -> None: async def _compact_session(self, ctx: TurnContext) -> None:
session = ctx.require_session() session = ctx.require_session()
ctx.session, pending = self.auto_compact.prepare_session( ctx.session, pending = await self.auto_compact.prepare_session_async(
session, session,
ctx.session_key, ctx.session_key,
) )
@@ -1842,14 +1976,14 @@ class AgentLoop:
# them out of LLM context. /new is excluded because it # them out of LLM context. /new is excluded because it
# intentionally clears the session. # intentionally clears the session.
if cmd_ctx.raw.lower() != "/new": if cmd_ctx.raw.lower() != "/new":
ctx.input_persisted_early = self._persist_user_message_early( ctx.input_persisted_early = await self._persist_user_message_early_async(
ctx.msg, session, _command=True ctx.msg, session, _command=True
) )
session.add_message( session.add_message(
"assistant", result.content, _command=True "assistant", result.content, _command=True
) )
self._clear_pending_user_turn(session) self._clear_pending_user_turn(session)
self.sessions.save(session) await self._save_session(session)
if not ctx.ephemeral: if not ctx.ephemeral:
await self.runtime_event_publisher.session_turn_persisted( await self.runtime_event_publisher.session_turn_persisted(
ctx.msg, ctx.msg,
@@ -1864,7 +1998,7 @@ class AgentLoop:
session = ctx.require_session() session = ctx.require_session()
runtime = ctx.runtime runtime = ctx.runtime
if runtime is None: if runtime is None:
runtime = self.runtime_for_session(session) runtime = await self.runtime_for_session_async(session)
ctx.runtime = runtime ctx.runtime = runtime
if ctx.session_key.startswith("dream:"): if ctx.session_key.startswith("dream:"):
logger.info( logger.info(
@@ -1879,15 +2013,12 @@ class AgentLoop:
session, session,
runtime=runtime, runtime=runtime,
) )
# Token consolidation may have committed a replacement checkpoint
# after the compact stage captured its summary for this request.
ctx.session, ctx.pending_summary = self.auto_compact.prepare_session(
session,
ctx.session_key,
)
session = ctx.require_session()
is_subagent = ctx.kind is TurnKind.SYSTEM and ctx.msg.sender_id == "subagent" is_subagent = ctx.kind is TurnKind.SYSTEM and ctx.msg.sender_id == "subagent"
if ctx.kind is TurnKind.USER and (message_tool := self.tools.get("message")):
if isinstance(message_tool, MessageTool):
message_tool.start_turn()
_hist_kwargs: dict[str, Any] = { _hist_kwargs: dict[str, Any] = {
"max_tokens": self._replay_token_budget(runtime), "max_tokens": self._replay_token_budget(runtime),
"extend_to_user": is_subagent, "extend_to_user": is_subagent,
@@ -1911,7 +2042,7 @@ class AgentLoop:
# provider compatibility or prompt assembly work. A compatible # provider compatibility or prompt assembly work. A compatible
# staged state replaces this in a second atomic save below. # staged state replaces this in a second atomic save below.
session.provider_state = None session.provider_state = None
self.sessions.save(session) await self._save_session(session)
ctx.input_persisted_early = True ctx.input_persisted_early = True
await ctx.delivery.runtime_admitted(runtime) await ctx.delivery.runtime_admitted(runtime)
@@ -1965,7 +2096,7 @@ class AgentLoop:
elif stored_state is not None: elif stored_state is not None:
session.provider_state = None session.provider_state = None
if ctx.kind is TurnKind.USER: if ctx.kind is TurnKind.USER:
ctx.input_persisted_early = self._persist_user_message_early( ctx.input_persisted_early = await self._persist_user_message_early_async(
ctx.msg, ctx.msg,
session, session,
runtime_context_blocks=ctx.runtime_context_blocks, runtime_context_blocks=ctx.runtime_context_blocks,
@@ -1975,8 +2106,8 @@ class AgentLoop:
elif subagent_followup_persisted and staged_provider_state: elif subagent_followup_persisted and staged_provider_state:
# Upgrade the replay-safe baseline to the resumable state before # Upgrade the replay-safe baseline to the resumable state before
# prompt assembly and the first model checkpoint. # prompt assembly and the first model checkpoint.
self.sessions.save(session) await self._save_session(session)
ctx.transcript_input = self._build_transcript_input(ctx) ctx.initial_messages = self._build_initial_messages(ctx)
if ctx.on_progress is None: if ctx.on_progress is None:
ctx.on_progress = ctx.delivery.progress_callback() ctx.on_progress = ctx.delivery.progress_callback()
@@ -1988,36 +2119,36 @@ class AgentLoop:
if ctx.visible_run_started_at is None: if ctx.visible_run_started_at is None:
ctx.visible_run_started_at = time.time() ctx.visible_run_started_at = time.time()
await ctx.delivery.running(started_at=ctx.visible_run_started_at) await ctx.delivery.running(started_at=ctx.visible_run_started_at)
assert ctx.transcript_input is not None result = await self._run_agent_loop(
with capture_message_deliveries() as message_sends: ctx.initial_messages,
result = await self._run_agent_loop( runtime=runtime,
ctx.transcript_input, on_progress=ctx.on_progress,
runtime=runtime, on_stream=ctx.on_stream,
on_progress=ctx.on_progress, on_stream_end=ctx.on_stream_end,
on_stream=ctx.on_stream, on_retry_wait=ctx.on_retry_wait,
on_stream_end=ctx.on_stream_end, session=ctx.session,
on_retry_wait=ctx.on_retry_wait, channel=ctx.delivery.route.channel,
session=ctx.session, chat_id=ctx.delivery.route.chat_id,
pending_queue=ctx.pending_queue, message_id=ctx.msg.metadata.get("message_id"),
ephemeral=ctx.ephemeral, metadata=ctx.msg.metadata,
run_extra_hooks_for_ephemeral=ctx.run_extra_hooks_for_ephemeral, session_key=ctx.session_key,
hooks=ctx.hooks, original_user_text=ctx.original_user_text,
hook_factories=ctx.hook_factories, pending_queue=ctx.pending_queue,
turn_scopes=ctx.turn_scopes, ephemeral=ctx.ephemeral,
tools=ctx.tools, run_extra_hooks_for_ephemeral=ctx.run_extra_hooks_for_ephemeral,
request_context=ctx.request_context, hooks=ctx.hooks,
provider_state=ctx.provider_state, hook_factories=ctx.hook_factories,
) turn_scopes=ctx.turn_scopes,
ctx.final_content = result.final_content tools=ctx.tools,
ctx.all_messages = result.messages request_context=ctx.request_context,
ctx.stop_reason = result.stop_reason provider_state=ctx.provider_state,
if ( )
ctx.kind is TurnKind.USER final_content, _, all_msgs, stop_reason, had_injections = result
and (ctx.delivery.route.channel, ctx.delivery.route.chat_id) in message_sends ctx.final_content = final_content
and (not result.had_injections or result.stop_reason == "empty_final_response") ctx.all_messages = all_msgs
): ctx.stop_reason = stop_reason
ctx.suppress_response = True ctx.had_injections = had_injections
ctx.usage = result.usage ctx.usage = self._last_usage
ctx.delivery.record_usage(ctx.usage) ctx.delivery.record_usage(ctx.usage)
if ctx.kind is TurnKind.USER: if ctx.kind is TurnKind.USER:
await turn_continuation.maybe_continue_turn(ctx) await turn_continuation.maybe_continue_turn(ctx)
@@ -2051,6 +2182,9 @@ class AgentLoop:
turn_latency_ms=ctx.turn_latency_ms, turn_latency_ms=ctx.turn_latency_ms,
) )
ctx.delivery.record_latency(ctx.turn_latency_ms) ctx.delivery.record_latency(ctx.turn_latency_ms)
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
await self._save_session(session)
if not ctx.ephemeral: if not ctx.ephemeral:
self.schedule_background( self.schedule_background(
self.consolidator.maybe_consolidate_by_tokens( self.consolidator.maybe_consolidate_by_tokens(
@@ -2058,10 +2192,6 @@ class AgentLoop:
runtime=runtime, runtime=runtime,
) )
) )
self._clear_pending_user_turn(session)
self._clear_runtime_checkpoint(session)
self.sessions.save(session)
if not ctx.ephemeral:
await self.runtime_event_publisher.session_turn_persisted( await self.runtime_event_publisher.session_turn_persisted(
ctx.msg, ctx.msg,
ctx.session_key, ctx.session_key,
@@ -2085,6 +2215,7 @@ class AgentLoop:
ctx.delivery.delivery_message, ctx.delivery.delivery_message,
cast(str, ctx.final_content), cast(str, ctx.final_content),
ctx.stop_reason, ctx.stop_reason,
ctx.had_injections,
ctx.streamed_content, ctx.streamed_content,
log_content=ctx.require_session().policy.log_content, log_content=ctx.require_session().policy.log_content,
turn_latency_ms=ctx.turn_latency_ms, turn_latency_ms=ctx.turn_latency_ms,
@@ -2273,11 +2404,24 @@ class AgentLoop:
) )
return True return True
def _set_runtime_checkpoint(self, session: Session, payload: dict[str, Any]) -> None: def _set_runtime_checkpoint(
"""Persist the latest in-flight turn state into session metadata.""" self,
session: Session,
payload: dict[str, Any],
) -> None:
"""Synchronously persist a checkpoint for compatibility callers."""
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
self.sessions.save_runtime_checkpoint(session) self.sessions.save_runtime_checkpoint(session)
async def _set_runtime_checkpoint_async(
self,
session: Session,
payload: dict[str, Any],
) -> None:
"""Persist the latest in-flight turn state without blocking the event loop."""
session.metadata[self._RUNTIME_CHECKPOINT_KEY] = payload
await self._save_runtime_checkpoint(session)
def _mark_pending_user_turn(self, session: Session) -> None: def _mark_pending_user_turn(self, session: Session) -> None:
session.metadata[self._PENDING_USER_TURN_KEY] = True session.metadata[self._PENDING_USER_TURN_KEY] = True
@@ -2288,6 +2432,10 @@ class AgentLoop:
if self._RUNTIME_CHECKPOINT_KEY in session.metadata: if self._RUNTIME_CHECKPOINT_KEY in session.metadata:
session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None) session.metadata.pop(self._RUNTIME_CHECKPOINT_KEY, None)
def _restore_runtime_checkpoint(self, session: Session) -> bool:
"""Materialize an unfinished turn into session history before a new request."""
return restore_runtime_checkpoint(session)
async def process_direct( async def process_direct(
self, self,
content: str, content: str,
+317 -322
View File
@@ -1,4 +1,4 @@
"""Memory storage, transcript archiving, and legacy consolidation coordination.""" """Memory system: pure file I/O store and lightweight Consolidator."""
# Tool schemas are installed by the ``@tool_parameters`` class decorator at # Tool schemas are installed by the ``@tool_parameters`` class decorator at
# runtime; static analyzers cannot observe that it clears ``parameters`` from # runtime; static analyzers cannot observe that it clears ``parameters`` from
@@ -22,6 +22,7 @@ from loguru import logger
from nanobot.llm_usage.context import llm_usage_source from nanobot.llm_usage.context import llm_usage_source
from nanobot.runtime_context import public_history_messages from nanobot.runtime_context import public_history_messages
from nanobot.session.async_compat import call_session_manager
from nanobot.session.manager import ( from nanobot.session.manager import (
MIN_COMPACTED_REPLAY_MESSAGES, MIN_COMPACTED_REPLAY_MESSAGES,
Session, Session,
@@ -32,10 +33,10 @@ from nanobot.utils.gitstore import GitStore
from nanobot.utils.helpers import ( from nanobot.utils.helpers import (
content_with_media_breadcrumbs, content_with_media_breadcrumbs,
ensure_dir, ensure_dir,
estimate_message_tokens,
estimate_prompt_tokens_chain, estimate_prompt_tokens_chain,
strip_think, strip_think,
truncate_text, truncate_text,
truncate_text_to_tokens,
) )
from nanobot.utils.prompt_templates import render_template from nanobot.utils.prompt_templates import render_template
from nanobot.utils.workspace_prompts import ( from nanobot.utils.workspace_prompts import (
@@ -66,6 +67,8 @@ class MemoryStore:
# durable files are tiny in practice (~5 KB total), but a runaway file must # durable files are tiny in practice (~5 KB total), but a runaway file must
# not unbounded the prompt. # not unbounded the prompt.
_DREAM_FILE_EMBED_CAP = 8000 _DREAM_FILE_EMBED_CAP = 8000
_INTERNAL_HISTORY_SESSION_PREFIXES = ("cron:", "dream:")
_INTERNAL_HISTORY_SESSION_KEYS = {"heartbeat"}
_LEGACY_ENTRY_START_RE = re.compile(r"^\[(\d{4}-\d{2}-\d{2}[^\]]*)\]\s*") _LEGACY_ENTRY_START_RE = re.compile(r"^\[(\d{4}-\d{2}-\d{2}[^\]]*)\]\s*")
_LEGACY_TIMESTAMP_RE = re.compile(r"^\[(\d{4}-\d{2}-\d{2} \d{2}:\d{2})\]\s*") _LEGACY_TIMESTAMP_RE = re.compile(r"^\[(\d{4}-\d{2}-\d{2} \d{2}:\d{2})\]\s*")
_LEGACY_RAW_MESSAGE_RE = re.compile( _LEGACY_RAW_MESSAGE_RE = re.compile(
@@ -259,29 +262,6 @@ class MemoryStore:
# -- history.jsonl — append-only, JSONL format --------------------------- # -- history.jsonl — append-only, JSONL format ---------------------------
def _normalize_history_entry(
self,
entry: str,
*,
max_chars: int | None = None,
) -> str:
"""Return the exact bounded, model-safe text accepted by the journal."""
limit = max_chars if max_chars is not None else _HISTORY_ENTRY_HARD_CAP
raw = entry.rstrip()
content = strip_think(raw)
if len(content) > limit:
if not self._oversize_logged:
self._oversize_logged = True
logger.warning(
"history entry exceeds {} chars ({}); truncating. "
"Usually means a caller forgot its own cap; "
"further occurrences suppressed.",
limit,
len(content),
)
content = truncate_text(content, limit)
return content
def append_history( def append_history(
self, self,
entry: str, entry: str,
@@ -296,16 +276,27 @@ class MemoryStore:
persisted. If the cleaned content is empty but the raw entry wasn't, persisted. If the cleaned content is empty but the raw entry wasn't,
the record is persisted with an empty string rather than falling back the record is persisted with an empty string rather than falling back
to the raw leak otherwise `strip_think`'s guarantees would be to the raw leak otherwise `strip_think`'s guarantees would be
undone when Dream consumes the journal entry. undone by history replay / consolidation downstream.
A defensive cap (*max_chars*, default ``_HISTORY_ENTRY_HARD_CAP``) is A defensive cap (*max_chars*, default ``_HISTORY_ENTRY_HARD_CAP``) is
applied as a final safety net: individual callers should cap their own applied as a final safety net: individual callers should cap their own
content more tightly; this default only exists to catch unintentional content more tightly; this default only exists to catch unintentional
large writes (e.g. an LLM echoing its input back as a "summary"). large writes (e.g. an LLM echoing its input back as a "summary").
""" """
limit = max_chars if max_chars is not None else _HISTORY_ENTRY_HARD_CAP
ts = datetime.now().strftime("%Y-%m-%d %H:%M") ts = datetime.now().strftime("%Y-%m-%d %H:%M")
raw = entry.rstrip() raw = entry.rstrip()
content = self._normalize_history_entry(entry, max_chars=max_chars) if len(raw) > limit:
if not self._oversize_logged:
self._oversize_logged = True
logger.warning(
"history entry exceeds {} chars ({}); truncating. "
"Usually means a caller forgot its own cap; "
"further occurrences suppressed.",
limit, len(raw),
)
raw = truncate_text(raw, limit)
content = strip_think(raw)
# Cursor allocation and the append must be atomic: concurrent writers # Cursor allocation and the append must be atomic: concurrent writers
# could otherwise read the same current cursor and emit duplicates. # could otherwise read the same current cursor and emit duplicates.
with self._append_lock: with self._append_lock:
@@ -313,7 +304,7 @@ class MemoryStore:
if raw and not content: if raw and not content:
logger.debug( logger.debug(
"history entry {} stripped to empty (likely template leak); " "history entry {} stripped to empty (likely template leak); "
"persisting empty content to avoid re-polluting Dream input", "persisting empty content to avoid re-polluting context",
cursor, cursor,
) )
record = {"cursor": cursor, "timestamp": ts, "content": content} record = {"cursor": cursor, "timestamp": ts, "content": content}
@@ -403,6 +394,36 @@ class MemoryStore:
"""Return history entries with a valid cursor > *since_cursor*.""" """Return history entries with a valid cursor > *since_cursor*."""
return [e for e, c in self._iter_valid_entries() if c > since_cursor] return [e for e, c in self._iter_valid_entries() if c > since_cursor]
@classmethod
def _is_internal_history_session(cls, session_key: str | None) -> bool:
if not session_key:
return False
return (
session_key in cls._INTERNAL_HISTORY_SESSION_KEYS
or session_key.startswith(cls._INTERNAL_HISTORY_SESSION_PREFIXES)
)
def read_recent_history_for_prompt(
self,
since_cursor: int,
*,
session_key: str | None,
unified_session: bool = False,
) -> list[dict[str, Any]]:
"""Return unprocessed history entries safe to inject into a turn prompt."""
entries = self.read_unprocessed_history(since_cursor=since_cursor)
if session_key is None:
return entries
if not unified_session:
return [e for e in entries if e.get("session_key") == session_key]
return [
entry
for entry in entries
if (entry_session := entry.get("session_key")) == session_key
or not self._is_internal_history_session(entry_session)
]
def compact_history(self) -> None: def compact_history(self) -> None:
"""Drop oldest processed entries without discarding pending Dream input.""" """Drop oldest processed entries without discarding pending Dream input."""
if self.max_history_entries <= 0: if self.max_history_entries <= 0:
@@ -699,28 +720,21 @@ class MemoryStore:
*, *,
max_chars: int | None = None, max_chars: int | None = None,
session_key: str | None = None, session_key: str | None = None,
) -> str: ) -> None:
"""Persist and return a bounded raw checkpoint when summarization degrades.""" """Fallback: dump raw messages to history.jsonl without LLM summarization."""
checkpoint = self._build_raw_checkpoint(messages, max_chars=max_chars) limit = max_chars if max_chars is not None else _RAW_ARCHIVE_MAX_CHARS
self.append_history(checkpoint, session_key=session_key) formatted = truncate_text(
self._format_messages(public_history_messages(messages)),
limit,
)
self.append_history(
f"[RAW] {len(messages)} messages\n"
f"{formatted}",
session_key=session_key,
)
logger.warning( logger.warning(
"Memory consolidation degraded: raw-archived {} messages", len(messages) "Memory consolidation degraded: raw-archived {} messages", len(messages)
) )
return checkpoint
def _build_raw_checkpoint(
self,
messages: list[dict[str, Any]],
*,
max_chars: int | None = None,
) -> str:
"""Build the same bounded checkpoint as :meth:`raw_archive` without writing it."""
limit = max_chars if max_chars is not None else _RAW_ARCHIVE_MAX_CHARS
checkpoint = (
f"[RAW] {len(messages)} messages\n"
f"{self._format_messages(public_history_messages(messages))}"
)
return self._normalize_history_entry(checkpoint, max_chars=limit)
# ------------------------------------------------------------------ # ------------------------------------------------------------------
# Dream helpers # Dream helpers
@@ -772,215 +786,21 @@ class MemoryStore:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Memory ingestion and legacy context-pressure coordination # Consolidator — lightweight token-budget triggered consolidation
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Raw fallbacks use a tighter cap. Completed model summaries may scale with the # Individual history.jsonl writers cap their own payloads tightly; the
# configured generation budget, while append_history() still enforces the # _HISTORY_ENTRY_HARD_CAP at append_history() is a belt-and-suspenders default
# emergency hard cap against pathological provider output. # that catches any new caller that forgot to set its own cap.
_RAW_ARCHIVE_MAX_CHARS = 16_000 # fallback dump (LLM failed) _RAW_ARCHIVE_MAX_CHARS = 16_000 # fallback dump (LLM failed)
_HISTORY_ENTRY_HARD_CAP = 64_000 # emergency cap in append_history _ARCHIVE_SUMMARY_MAX_CHARS = 8_000 # LLM-produced consolidation summary
_HISTORY_ENTRY_HARD_CAP = 64_000 # emergency cap in append_history
class MemoryArchiver:
"""Write durable transcript batches to the Memory ingestion journal.
The archiver deliberately has no SessionManager dependency: it may read a
captured transcript batch and append to history.jsonl, but it cannot mutate
provider continuation state or advance a session watermark.
"""
def __init__(
self,
store: MemoryStore,
build_messages: Callable[..., list[dict[str, Any]]],
get_tool_definitions: Callable[[], list[dict[str, Any]]],
resolve_prompt_context: Callable[[Session], tuple[str | None, Path | None]] | None = None,
) -> None:
self.store = store
self._build_messages = build_messages
self._get_tool_definitions = get_tool_definitions
self._resolve_prompt_context = resolve_prompt_context
def _raw_checkpoint(
self,
messages: list[dict[str, Any]],
*,
session_key: str,
previous_summary: str | None,
max_tokens: int,
) -> str:
"""Persist the failed chunk and return a bounded replacement checkpoint."""
raw = self.store.raw_archive(messages, session_key=session_key)
token_limit = max(1, max_tokens)
if not previous_summary:
return truncate_text_to_tokens(raw, token_limit)
combined = (
"[Previous archived context]\n"
f"{previous_summary}\n\n"
"[Newly archived raw context]\n"
f"{raw}"
)
bounded = truncate_text_to_tokens(combined, token_limit)
if bounded == combined:
return combined
# Keep evidence from both sides when their full concatenation cannot fit.
section_limit = max(1, (token_limit - 32) // 2)
return truncate_text_to_tokens(
"[Previous archived context]\n"
f"{truncate_text_to_tokens(previous_summary, section_limit)}\n\n"
"[Newly archived raw context]\n"
f"{truncate_text_to_tokens(raw, section_limit)}",
token_limit,
)
async def archive(
self,
messages: list[dict[str, Any]],
*,
runtime: LLMRuntime,
session_key: str,
request_messages: list[dict[str, Any]],
request_tools: list[dict[str, Any]],
previous_summary: str | None = None,
) -> str | None:
"""Execute a prepared archive request and persist its result."""
if not messages:
return None
def raw_fallback() -> str:
return self._raw_checkpoint(
messages,
session_key=session_key,
previous_summary=previous_summary,
max_tokens=runtime.generation.max_tokens,
)
try:
with llm_usage_source("dream"):
response = await runtime.provider.chat_with_retry(
model=runtime.model,
messages=request_messages,
tools=request_tools,
temperature=runtime.generation.temperature,
max_tokens=runtime.generation.max_tokens,
reasoning_effort=runtime.generation.reasoning_effort,
)
except Exception:
logger.warning("Memory archive provider call failed, raw-dumping to history")
return raw_fallback()
if response.finish_reason in {"error", "length"}:
logger.warning(
"Memory archive provider did not complete ({}), raw-dumping to history",
response.finish_reason,
)
return raw_fallback()
if response.has_tool_calls is True:
logger.warning("Memory archive provider returned tool calls, raw-dumping to history")
return raw_fallback()
summary = response.content
if not summary or not summary.strip():
logger.warning("Memory archive provider returned no summary, raw-dumping to history")
return raw_fallback()
summary = self.store._normalize_history_entry(summary)
if not summary:
logger.warning("Memory archive provider summary was not safe to replay, raw-dumping")
return raw_fallback()
if summary == "(nothing)":
return "(nothing)"
self.store.append_history(summary, session_key=session_key)
return summary
async def archive_session(
self,
session: Session,
*,
archive_end: int,
runtime: LLMRuntime,
input_token_budget: int,
) -> str | None:
"""Archive a captured session prefix without mutating the session."""
messages = list(session.messages[session.last_archived:archive_end])
if not messages:
return None
session_summary = session_summary_from_metadata(
session.metadata,
fallback_last_active=session.updated_at,
)
previous_summary = session_summary["text"] if session_summary else None
def raw_fallback() -> str:
return self._raw_checkpoint(
messages,
session_key=session.key,
previous_summary=previous_summary,
max_tokens=runtime.generation.max_tokens,
)
if input_token_budget <= 0:
logger.debug(
"Memory archive has no safe input budget for {}; raw-dumping",
session.key,
)
return raw_fallback()
prefix = Session(
key=session.key,
messages=list(session.messages[:archive_end]),
last_consolidated=session.last_archived,
)
history = prefix.get_history(max_tokens=input_token_budget)
archive_history = Session(
key=session.key,
messages=messages,
).get_history()
if not archive_history or history[-len(archive_history):] != archive_history:
logger.debug(
"Memory archive cannot replay the full chunk for {}; raw-dumping",
session.key,
)
return raw_fallback()
prompt = render_template("agent/consolidator_archive.md", strip=True)
channel = session.key.split(":", 1)[0] if ":" in session.key else None
workspace: Path | None = None
if self._resolve_prompt_context is not None:
channel, workspace = self._resolve_prompt_context(session)
request_messages = self._build_messages(
history=history,
current_message=prompt,
channel=channel,
session_summary=session_summary,
workspace=workspace,
)
tools = self._get_tool_definitions()
estimated, source = estimate_prompt_tokens_chain(
runtime.provider,
runtime.model,
request_messages,
tools,
)
if estimated > input_token_budget:
logger.debug(
"Memory archive prefix exceeds budget for {}; raw-dumping: {}/{} via {}",
session.key,
estimated,
input_token_budget,
source,
)
return raw_fallback()
return await self.archive(
messages,
runtime=runtime,
session_key=session.key,
request_messages=request_messages,
request_tools=tools,
previous_summary=previous_summary,
)
class Consolidator: class Consolidator:
"""Legacy context-pressure coordinator backed by a MemoryArchiver.""" """Summarize compacted messages into history.jsonl."""
_MAX_CONSOLIDATION_ROUNDS = 5
_SAFETY_BUFFER = 1024 # extra headroom for tokenizer estimation drift _SAFETY_BUFFER = 1024 # extra headroom for tokenizer estimation drift
@@ -991,21 +811,36 @@ class Consolidator:
build_messages: Callable[..., list[dict[str, Any]]], build_messages: Callable[..., list[dict[str, Any]]],
get_tool_definitions: Callable[[], list[dict[str, Any]]], get_tool_definitions: Callable[[], list[dict[str, Any]]],
resolve_prompt_context: Callable[[Session], tuple[str | None, Path | None]] | None = None, resolve_prompt_context: Callable[[Session], tuple[str | None, Path | None]] | None = None,
consolidation_ratio: float = 0.5,
unified_session: bool = False,
): ):
self.store = store self.store = store
self.sessions = sessions self.sessions = sessions
self.consolidation_ratio = consolidation_ratio
self.unified_session = unified_session
self._build_messages = build_messages self._build_messages = build_messages
self._get_tool_definitions = get_tool_definitions self._get_tool_definitions = get_tool_definitions
self.archiver = MemoryArchiver( self._resolve_prompt_context = resolve_prompt_context
store=store,
build_messages=build_messages,
get_tool_definitions=get_tool_definitions,
resolve_prompt_context=resolve_prompt_context,
)
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = ( self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
weakref.WeakValueDictionary() weakref.WeakValueDictionary()
) )
async def _get_or_create_session(self, key: str) -> Session:
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
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())
@@ -1013,19 +848,24 @@ class Consolidator:
def pick_consolidation_boundary( def pick_consolidation_boundary(
self, self,
session: Session, session: Session,
) -> int | None: tokens_to_remove: int,
"""Return the fixed user-led boundary before the recent replay tail.""" ) -> tuple[int, int] | None:
if not session.messages: """Pick a user-turn boundary that removes enough old prompt tokens."""
start = session.last_consolidated
if start >= len(session.messages) or tokens_to_remove <= 0:
return None return None
boundary = max(0, len(session.messages) - MIN_COMPACTED_REPLAY_MESSAGES)
while boundary > 0 and session.messages[boundary].get("role") != "user": removed_tokens = 0
boundary -= 1 last_boundary: tuple[int, int] | None = None
if ( for idx in range(start, len(session.messages)):
boundary <= session.last_archived message = session.messages[idx]
or session.messages[boundary].get("role") != "user" if idx > start and message.get("role") == "user":
): last_boundary = (idx, removed_tokens)
return None if removed_tokens >= tokens_to_remove:
return boundary return last_boundary
removed_tokens += estimate_message_tokens(message)
return last_boundary
@staticmethod @staticmethod
def _full_replay_history( def _full_replay_history(
@@ -1036,18 +876,13 @@ class Consolidator:
return [] return []
return session.get_history() return session.get_history()
@staticmethod async def _persist_last_summary(self, session: Session, summary: str | None) -> None:
def _set_last_summary( if summary and summary != "(nothing)":
session: Session,
summary: str,
*,
last_active: datetime | None = None,
) -> None:
if summary != "(nothing)":
session.metadata["_last_summary"] = { session.metadata["_last_summary"] = {
"text": summary, "text": summary,
"last_active": (last_active or session.updated_at).isoformat(), "last_active": session.updated_at.isoformat(),
} }
await self._save_session(session)
def estimate_session_prompt_tokens( def estimate_session_prompt_tokens(
self, self,
@@ -1067,6 +902,8 @@ class Consolidator:
current_message="[token-probe]", current_message="[token-probe]",
channel=channel, channel=channel,
session_summary=summary, session_summary=summary,
session_key=session.key,
unified_session=self.unified_session,
) )
return estimate_prompt_tokens_chain( return estimate_prompt_tokens_chain(
runtime.provider, runtime.provider,
@@ -1083,6 +920,58 @@ class Consolidator:
- self._SAFETY_BUFFER - self._SAFETY_BUFFER
) )
async def archive(
self,
messages: list[dict[str, Any]],
*,
runtime: LLMRuntime,
session_key: str,
request_messages: list[dict[str, Any]],
request_tools: list[dict[str, Any]],
) -> str | None:
"""Execute a prepared consolidation request and persist its result."""
if not messages:
return None
try:
with llm_usage_source("dream"):
response = await runtime.provider.chat_with_retry(
model=runtime.model,
messages=request_messages,
tools=request_tools,
tool_choice="none",
temperature=runtime.generation.temperature,
max_tokens=runtime.generation.max_tokens,
reasoning_effort=runtime.generation.reasoning_effort,
)
except Exception:
logger.warning("Consolidation provider call failed, raw-dumping to history")
self.store.raw_archive(messages, session_key=session_key)
return None
if response.finish_reason in {"error", "length"}:
logger.warning(
"Consolidation provider did not complete ({}), raw-dumping to history",
response.finish_reason,
)
self.store.raw_archive(messages, session_key=session_key)
return None
if response.has_tool_calls is True:
logger.warning("Consolidation provider returned tool calls, raw-dumping to history")
self.store.raw_archive(messages, session_key=session_key)
return None
summary = response.content
if not summary or not summary.strip():
logger.warning("Consolidation provider returned no summary, raw-dumping to history")
self.store.raw_archive(messages, session_key=session_key)
return None
if summary.strip() == "(nothing)":
return "(nothing)"
self.store.append_history(
summary,
max_chars=_ARCHIVE_SUMMARY_MAX_CHARS,
session_key=session_key,
)
return summary
async def archive_session( async def archive_session(
self, self,
session: Session, session: Session,
@@ -1090,12 +979,82 @@ class Consolidator:
archive_end: int, archive_end: int,
runtime: LLMRuntime, runtime: LLMRuntime,
) -> str | None: ) -> str | None:
"""Compatibility wrapper for the extracted MemoryArchiver.""" """Archive a session prefix by appending a consolidation instruction."""
return await self.archiver.archive_session( messages = list(session.messages[session.last_consolidated:archive_end])
session, if not messages:
archive_end=archive_end, return None
budget = self._input_token_budget(runtime)
if budget <= 0:
logger.debug(
"Consolidation has no safe input budget for {}; raw-dumping",
session.key,
)
self.store.raw_archive(messages, session_key=session.key)
return None
prefix = Session(
key=session.key,
messages=list(session.messages[:archive_end]),
last_consolidated=session.last_consolidated,
)
history = prefix.get_history(max_tokens=budget)
archive_history = Session(
key=session.key,
messages=messages,
).get_history()
if (
not archive_history
or history[-len(archive_history):] != archive_history
):
logger.debug(
"Consolidation cannot replay the full chunk for {}; raw-dumping",
session.key,
)
self.store.raw_archive(messages, session_key=session.key)
return None
prompt = render_template(
"agent/consolidator_archive.md",
strip=True,
archive_count=len(archive_history),
)
channel = session.key.split(":", 1)[0] if ":" in session.key else None
workspace: Path | None = None
if self._resolve_prompt_context is not None:
channel, workspace = self._resolve_prompt_context(session)
request_messages = self._build_messages(
history=history,
current_message=prompt,
channel=channel,
session_summary=session_summary_from_metadata(
session.metadata,
fallback_last_active=session.updated_at,
),
workspace=workspace,
session_key=session.key,
unified_session=self.unified_session,
)
tools = self._get_tool_definitions()
estimated, source = estimate_prompt_tokens_chain(
runtime.provider,
runtime.model,
request_messages,
tools,
)
if estimated > budget:
logger.debug(
"Consolidation prefix exceeds budget for {}; raw-dumping: {}/{} via {}",
session.key,
estimated,
budget,
source,
)
self.store.raw_archive(messages, session_key=session.key)
return None
return await self.archive(
messages,
runtime=runtime, runtime=runtime,
input_token_budget=self._input_token_budget(runtime), session_key=session.key,
request_messages=request_messages,
request_tools=tools,
) )
async def maybe_consolidate_by_tokens( async def maybe_consolidate_by_tokens(
@@ -1104,71 +1063,104 @@ class Consolidator:
*, *,
runtime: LLMRuntime, runtime: LLMRuntime,
) -> None: ) -> None:
"""Archive one fixed old prefix when the prompt exceeds the safe budget. """Loop: archive old messages until prompt fits within safe budget.
The budget reserves space for completion tokens and a safety buffer The budget reserves space for completion tokens and a safety buffer
so the LLM request never exceeds the context window. so the LLM request never exceeds the context window.
""" """
if runtime.context_window_tokens <= 0:
return
lock = self.get_lock(session.key) lock = self.get_lock(session.key)
async with lock: async with lock:
# Refresh session reference: AutoCompact may have replaced it. # Refresh session reference: AutoCompact may have replaced it.
fresh = self.sessions.get_or_create(session.key) fresh = await self._get_or_create_session(session.key)
if fresh is not session: if fresh is not session:
session = fresh session = fresh
if runtime.context_window_tokens <= 0:
return
if not session.messages: if not session.messages:
return return
budget = self._input_token_budget(runtime) budget = self._input_token_budget(runtime)
target = int(budget * self.consolidation_ratio)
last_summary: str | None = None
estimated, source = self.estimate_session_prompt_tokens( estimated, source = self.estimate_session_prompt_tokens(
session, session,
runtime=runtime, runtime=runtime,
) )
if estimated <= 0: if estimated <= 0:
await self._persist_last_summary(session, last_summary)
return return
if estimated < budget: if estimated < budget:
unarchived_count = len(session.messages) - session.last_archived unconsolidated_count = len(session.messages) - session.last_consolidated
logger.debug( logger.debug(
"Token consolidation idle {}: {}/{} via {}, msgs={}", "Token consolidation idle {}: {}/{} via {}, msgs={}",
session.key, session.key,
estimated, estimated,
runtime.context_window_tokens, runtime.context_window_tokens,
source, source,
unarchived_count, unconsolidated_count,
) )
await self._persist_last_summary(session, last_summary)
return return
end_idx = self.pick_consolidation_boundary(session) for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
if end_idx is None: if estimated <= target:
logger.debug( break
"Token consolidation: no safe fixed boundary for {}",
boundary = self.pick_consolidation_boundary(session, max(1, estimated - target))
if boundary is None:
logger.debug(
"Token consolidation: no safe boundary for {} (round {})",
session.key,
round_num,
)
break
end_idx = boundary[0]
chunk = session.messages[session.last_consolidated:end_idx]
if not chunk:
break
logger.info(
"Token consolidation round {} for {}: {}/{} via {}, chunk={} msgs",
round_num,
session.key, session.key,
estimated,
runtime.context_window_tokens,
source,
len(chunk),
) )
return summary = await self.archive_session(
session,
archive_end=end_idx,
runtime=runtime,
)
# Advance the cursor either way: on success the chunk was
# summarized; on failure archive_session() raw-archived it as
# a breadcrumb. Re-archiving the same chunk on the next call
# would just emit duplicate [RAW] entries.
if summary:
last_summary = summary
session.last_consolidated = end_idx
session.provider_state = None
await self._save_session(session)
if not summary:
# LLM is degraded — stop hammering it this call;
# the next invocation can retry a fresh chunk.
break
chunk = session.messages[session.last_archived:end_idx] estimated, source = self.estimate_session_prompt_tokens(
if not chunk: session,
return runtime=runtime,
)
if estimated <= 0:
break
logger.info( # Persist the last summary to session metadata so it can be injected
"Token consolidation for {}: {}/{} via {}, chunk={} msgs", # into the runtime context on the next prepare_session() call, aligning
session.key, # the summary injection strategy with AutoCompact._archive().
estimated, await self._persist_last_summary(session, last_summary)
runtime.context_window_tokens,
source,
len(chunk),
)
summary = await self.archive_session(
session,
archive_end=end_idx,
runtime=runtime,
)
if summary is None:
return
self._set_last_summary(session, summary)
session.last_archived = end_idx
self.sessions.save(session)
async def compact_idle_session( async def compact_idle_session(
self, self,
@@ -1193,9 +1185,9 @@ class Consolidator:
lock = self.get_lock(session_key) lock = self.get_lock(session_key)
async with lock: async with lock:
self.sessions.invalidate(session_key) self.sessions.invalidate(session_key)
session = self.sessions.get_or_create(session_key) session = await self._get_or_create_session(session_key)
archive_start = session.last_archived archive_start = session.last_consolidated
messages_to_archive = list(session.messages[archive_start:]) messages_to_archive = list(session.messages[archive_start:])
if not messages_to_archive: if not messages_to_archive:
return "" return ""
@@ -1207,15 +1199,18 @@ class Consolidator:
archive_end=archive_end, archive_end=archive_end,
runtime=runtime, runtime=runtime,
) )
if summary is None:
return None
self._set_last_summary(session, summary, last_active=last_active) if summary and summary != "(nothing)":
session.metadata["_last_summary"] = {
"text": summary,
"last_active": last_active.isoformat(),
}
# A turn can append while the provider call is in flight. Advance only # A turn can append while the provider call is in flight. Advance only
# through the captured batch so new messages remain eligible next time. # through the captured batch so new messages remain eligible next time.
session.last_archived = archive_end session.last_consolidated = archive_end
self.sessions.save(session) session.provider_state = None
await self._save_session(session)
visible = session.get_history( visible = session.get_history(
max_messages=MIN_COMPACTED_REPLAY_MESSAGES, max_messages=MIN_COMPACTED_REPLAY_MESSAGES,
+4
View File
@@ -31,6 +31,7 @@ class AgentProgressHook(AgentHook):
*, *,
session_key: str | None = None, session_key: str | None = None,
tool_hint_max_length: int = 40, tool_hint_max_length: int = 40,
on_iteration: Callable[[int], None] | None = None,
) -> None: ) -> None:
super().__init__(reraise=True) super().__init__(reraise=True)
self._on_progress = on_progress self._on_progress = on_progress
@@ -38,6 +39,7 @@ class AgentProgressHook(AgentHook):
self._on_stream_end = on_stream_end self._on_stream_end = on_stream_end
self._session_key = session_key self._session_key = session_key
self._tool_hint_max_length = tool_hint_max_length self._tool_hint_max_length = tool_hint_max_length
self._on_iteration = on_iteration
self._stream_buf = "" self._stream_buf = ""
self._think_extractor = IncrementalThinkExtractor() self._think_extractor = IncrementalThinkExtractor()
self._reasoning_open = False self._reasoning_open = False
@@ -94,6 +96,8 @@ class AgentProgressHook(AgentHook):
self._think_extractor.reset() self._think_extractor.reset()
async def before_iteration(self, context: AgentHookContext) -> None: async def before_iteration(self, context: AgentHookContext) -> None:
if self._on_iteration:
self._on_iteration(context.iteration)
logger.debug( logger.debug(
"Starting agent loop iteration {} for session {}", "Starting agent loop iteration {} for session {}",
context.iteration, context.iteration,
+478 -255
View File
File diff suppressed because it is too large Load Diff
+36 -35
View File
@@ -13,7 +13,7 @@ from typing import Any, Callable, NotRequired, TypedDict
from loguru import logger from loguru import logger
from nanobot.agent.hook import AgentHook, AgentHookContext from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.runner import AgentRunner, AgentRunSpec from nanobot.agent.runner import AgentRunner, AgentRunResult, AgentRunSpec
from nanobot.agent.tools.base import ToolResult from nanobot.agent.tools.base import ToolResult
from nanobot.agent.tools.context import ( from nanobot.agent.tools.context import (
RequestContext, RequestContext,
@@ -55,8 +55,7 @@ class SubagentStatus:
label: str label: str
task_description: str task_description: str
started_at: float # time.monotonic() started_at: float # time.monotonic()
# queued | initializing | awaiting_tools | tools_completed | final_response | done | error phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error
phase: str = "initializing"
iteration: int = 0 iteration: int = 0
tool_events: list[dict[str, str]] = field(default_factory=list) tool_events: list[dict[str, str]] = field(default_factory=list)
usage: LLMUsage | None = None usage: LLMUsage | None = None
@@ -105,6 +104,7 @@ class SubagentManager:
disabled_skills: list[str] | None = None, disabled_skills: list[str] | None = None,
max_iterations: int | None = None, max_iterations: int | None = None,
max_concurrent_subagents: int | None = None, max_concurrent_subagents: int | None = None,
fail_on_tool_error: bool | None = None,
llm_wall_timeout_for_session: Callable[[str | None], float | None] | None = None, llm_wall_timeout_for_session: Callable[[str | None], float | None] | None = None,
): ):
if workspace is None: if workspace is None:
@@ -148,7 +148,11 @@ class SubagentManager:
if max_concurrent_subagents is not None if max_concurrent_subagents is not None
else defaults.max_concurrent_subagents else defaults.max_concurrent_subagents
) )
self._run_slots = asyncio.Semaphore(self.max_concurrent_subagents) self.fail_on_tool_error = (
fail_on_tool_error
if fail_on_tool_error is not None
else defaults.fail_on_tool_error
)
self.runner = AgentRunner() self.runner = AgentRunner()
self._exec_session_manager = ExecSessionManager() self._exec_session_manager = ExecSessionManager()
self._llm_wall_timeout_for_session = llm_wall_timeout_for_session self._llm_wall_timeout_for_session = llm_wall_timeout_for_session
@@ -342,7 +346,7 @@ class SubagentManager:
self._session_tasks.setdefault(session_key, set()).add(task_id) self._session_tasks.setdefault(session_key, set()).add(task_id)
try: try:
result = await inline_task result = await inline_task
if status.phase == "error" or status.stop_reason == "error": if status.phase == "error" or status.stop_reason in {"error", "tool_error"}:
return ToolResult.error(result) return ToolResult.error(result)
return result return result
finally: finally:
@@ -365,35 +369,6 @@ class SubagentManager:
workspace_scope: WorkspaceScope | None = None, workspace_scope: WorkspaceScope | None = None,
*, *,
announce: bool = True, announce: bool = True,
) -> str:
"""Wait for capacity, then execute one subagent task."""
status.phase = "queued"
async with self._run_slots:
status.phase = "initializing"
return await self._run_admitted_subagent(
task_id,
task,
label,
origin,
status,
runtime,
origin_message_id,
workspace_scope,
announce=announce,
)
async def _run_admitted_subagent(
self,
task_id: str,
task: str,
label: str,
origin: _SubagentOrigin,
status: SubagentStatus,
runtime: LLMRuntime,
origin_message_id: str | None = None,
workspace_scope: WorkspaceScope | None = None,
*,
announce: bool = True,
) -> str: ) -> str:
"""Execute the subagent task and announce the result.""" """Execute the subagent task and announce the result."""
logger.info("Subagent [{}] starting task: {}", task_id, label) logger.info("Subagent [{}] starting task: {}", task_id, label)
@@ -441,6 +416,7 @@ class SubagentManager:
max_iterations_message="Task completed but no final response was generated.", max_iterations_message="Task completed but no final response was generated.",
finalize_on_max_iterations=False, finalize_on_max_iterations=False,
error_message=None, error_message=None,
fail_on_tool_error=self.fail_on_tool_error,
checkpoint_callback=_on_checkpoint, checkpoint_callback=_on_checkpoint,
session_key=sess_key, session_key=sess_key,
workspace=root, workspace=root,
@@ -457,7 +433,11 @@ class SubagentManager:
status.phase = "done" status.phase = "done"
status.stop_reason = result.stop_reason status.stop_reason = result.stop_reason
if result.stop_reason == "error": if result.stop_reason == "tool_error":
status.tool_events = list(result.tool_events)
final_result = self._format_partial_progress(result)
final_status = "error"
elif result.stop_reason == "error":
final_result = result.error or "Error: subagent execution failed." final_result = result.error or "Error: subagent execution failed."
final_status = "error" final_status = "error"
else: else:
@@ -538,6 +518,27 @@ class SubagentManager:
await self.bus.publish_inbound(msg) await self.bus.publish_inbound(msg)
logger.debug("Subagent [{}] announced result to {}:{}", task_id, origin['channel'], origin['chat_id']) logger.debug("Subagent [{}] announced result to {}:{}", task_id, origin['channel'], origin['chat_id'])
@staticmethod
def _format_partial_progress(result: AgentRunResult) -> str:
completed = [e for e in result.tool_events if e["status"] == "ok"]
failure = next((e for e in reversed(result.tool_events) if e["status"] == "error"), None)
lines: list[str] = []
if completed:
lines.append("Completed steps:")
for event in completed[-3:]:
lines.append(f"- {event['name']}: {event['detail']}")
if failure:
if lines:
lines.append("")
lines.append("Failure:")
lines.append(f"- {failure['name']}: {failure['detail']}")
if result.error and not failure:
if lines:
lines.append("")
lines.append("Failure:")
lines.append(f"- {result.error}")
return "\n".join(lines) or (result.error or "Error: subagent execution failed.")
def _build_subagent_prompt(self, workspace: Path | None = None) -> str: def _build_subagent_prompt(self, workspace: Path | None = None) -> str:
"""Build a focused system prompt for the subagent.""" """Build a focused system prompt for the subagent."""
from nanobot.agent.skills import SkillsLoader from nanobot.agent.skills import SkillsLoader
+22 -8
View File
@@ -4,7 +4,10 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import inspect
from pathlib import Path from pathlib import Path
from typing import TypedDict
from pydantic import Field from pydantic import Field
@@ -24,6 +27,14 @@ from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_li
from nanobot.security.workspace_access import current_tool_workspace from nanobot.security.workspace_access import current_tool_workspace
class _CliAppRunKwargs(TypedDict):
args: list[str]
json_output: bool
working_dir: str | None
timeout: int | None
restrict_to_workspace: bool
class CliAppsToolConfig(Base): class CliAppsToolConfig(Base):
"""CLI Apps tool configuration.""" """CLI Apps tool configuration."""
@@ -147,14 +158,17 @@ class CliAppsTool(Tool):
) )
workspace = access.project_path or self.workspace workspace = access.project_path or self.workspace
manager = CliAppManager(workspace=workspace, runtime=self.runtime) manager = CliAppManager(workspace=workspace, runtime=self.runtime)
run_kwargs: _CliAppRunKwargs = {
"args": args or [],
"json_output": bool(json),
"working_dir": working_dir,
"timeout": timeout,
"restrict_to_workspace": access.restrict_to_workspace,
}
try: try:
return manager.run( run_async = inspect.getattr_static(type(manager), "run_async", None)
name, if inspect.iscoroutinefunction(run_async):
args=args or [], return await manager.run_async(name, **run_kwargs)
json_output=bool(json), return await asyncio.to_thread(manager.run, name, **run_kwargs)
working_dir=working_dir,
timeout=timeout,
restrict_to_workspace=access.restrict_to_workspace,
)
except CliAppError as exc: except CliAppError as exc:
return ToolResult.error(f"Error: {exc.message}") return ToolResult.error(f"Error: {exc.message}")
-2
View File
@@ -11,7 +11,6 @@ if TYPE_CHECKING:
from nanobot.agent.subagent import SubagentManager from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.exec_session import ExecSessionManager from nanobot.agent.tools.exec_session import ExecSessionManager
from nanobot.agent.tools.file_state import FileStates from nanobot.agent.tools.file_state import FileStates
from nanobot.agent.tools.runtime_control import RuntimeControl
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.bus.runtime_events import RuntimeEventBus from nanobot.bus.runtime_events import RuntimeEventBus
from nanobot.config.schema import ProviderConfig, ToolsConfig from nanobot.config.schema import ProviderConfig, ToolsConfig
@@ -91,4 +90,3 @@ class ToolContext:
timezone: str = "UTC" timezone: str = "UTC"
workspace_sandbox: WorkspaceSandboxStatus | None = None workspace_sandbox: WorkspaceSandboxStatus | None = None
runtime_events: RuntimeEventBus | None = None runtime_events: RuntimeEventBus | None = None
runtime_control: RuntimeControl | None = None
+29 -2
View File
@@ -143,14 +143,41 @@ class CronTool(Tool):
tz: str | None = None, tz: str | None = None,
at: str | None = None, at: str | None = None,
job_id: str | None = None, job_id: str | None = None,
) -> str:
if action == "add" and self._in_cron_context.get():
return ToolResult.error(
"Error: cannot schedule new jobs from within a cron job execution"
)
return await self._cron.run_sync(
self._execute_sync,
action,
name,
message,
every_seconds,
cron_expr,
tz,
at,
job_id,
)
def _execute_sync(
self,
action: str,
name: str | None,
message: str,
every_seconds: int | None,
cron_expr: str | None,
tz: str | None,
at: str | None,
job_id: str | None,
) -> str: ) -> str:
if action == "add": if action == "add":
if self._in_cron_context.get(): if self._in_cron_context.get():
return ToolResult.error("Error: cannot schedule new jobs from within a cron job execution") return ToolResult.error("Error: cannot schedule new jobs from within a cron job execution")
return self._add_job(name, message, every_seconds, cron_expr, tz, at) return self._add_job(name, message, every_seconds, cron_expr, tz, at)
elif action == "list": if action == "list":
return self._list_jobs() return self._list_jobs()
elif action == "remove": if action == "remove":
return self._remove_job(job_id) return self._remove_job(job_id)
return f"Unknown action: {action}" return f"Unknown action: {action}"
+100 -94
View File
@@ -22,8 +22,7 @@ from nanobot.agent.tools.schema import (
DEFAULT_YIELD_MS = 1000 DEFAULT_YIELD_MS = 1000
MAX_YIELD_MS = 30_000 MAX_YIELD_MS = 30_000
DEFAULT_WAIT_FOR_MS = 10_000 DEFAULT_WAIT_FOR_MS = 10_000
DEFAULT_UNTIL_EXIT_MS = 600_000 MAX_WAIT_FOR_MS = 120_000
MAX_WAIT_FOR_MS = 600_000
DEFAULT_MAX_OUTPUT_CHARS = 10_000 DEFAULT_MAX_OUTPUT_CHARS = 10_000
MAX_OUTPUT_CHARS = 50_000 MAX_OUTPUT_CHARS = 50_000
OUTPUT_DRAIN_GRACE_S = 0.1 OUTPUT_DRAIN_GRACE_S = 0.1
@@ -496,39 +495,51 @@ def format_session_poll(session_id: str, poll: _SessionPoll) -> str:
@tool_parameters( @tool_parameters(
tool_parameters_schema( tool_parameters_schema(
session_id=StringSchema("Session ID returned by exec."), session_id=StringSchema("Session id returned by exec when yield_time_ms is used."),
input=StringSchema( chars=StringSchema(
"Text to send to stdin; omit to poll output.", "Bytes/text to write to stdin. Omit or pass an empty string to only poll recent output.",
nullable=True, nullable=True,
), ),
close_stdin=BooleanSchema( close_stdin=BooleanSchema(
description="Close stdin after sending input.", description="Close stdin after writing chars. Useful for commands waiting for EOF.",
default=False, default=False,
), ),
terminate=BooleanSchema( terminate=BooleanSchema(
description="Terminate the session; use alone.", description="Terminate the running exec session.",
default=False, default=False,
), ),
yield_time_ms=IntegerSchema(
description="Milliseconds to wait before returning recent output (default 1000, max 30000).",
minimum=0,
maximum=MAX_YIELD_MS,
),
wait_for=StringSchema( wait_for=StringSchema(
"Return when this text appears in output.", "Optional text to wait for in output before returning. "
min_length=1, "Useful for interactive commands and dev servers.",
nullable=True, nullable=True,
), ),
until_exit=BooleanSchema( wait_timeout_ms=IntegerSchema(
description="Wait for the process to exit.", description="Maximum milliseconds to wait for wait_for text (default 10000, max 120000).",
default=False,
),
timeout_ms=IntegerSchema(
description="Maximum wait: 1s normally, 10s for wait_for, 10m for until_exit.",
minimum=0, minimum=0,
maximum=MAX_WAIT_FOR_MS, maximum=MAX_WAIT_FOR_MS,
nullable=True, nullable=True,
), ),
max_output_chars=IntegerSchema(
description="Maximum output characters to return from this poll (default 10000, max 50000).",
minimum=1000,
maximum=MAX_OUTPUT_CHARS,
),
max_output_tokens=IntegerSchema(
description="Compatibility alias for max_output_chars. The current runtime uses a character budget.",
minimum=1000,
maximum=MAX_OUTPUT_CHARS,
nullable=True,
),
required=["session_id"], required=["session_id"],
) )
) )
class ExecSessionTool(Tool): class WriteStdinTool(Tool):
"""Interact with or wait for a running exec session.""" """Write to or poll a running exec session."""
_scopes = {"core", "subagent"} _scopes = {"core", "subagent"}
config_key = "exec" config_key = "exec"
@@ -560,103 +571,98 @@ class ExecSessionTool(Tool):
@property @property
def name(self) -> str: def name(self) -> str:
return "exec_session" return "write_stdin"
@property @property
def description(self) -> str: def description(self) -> str:
return "Manage a session returned by exec." return (
"Interact with a running exec session created by exec with "
"yield_time_ms. Use chars='' to poll without writing, chars to send "
"stdin, close_stdin=true to send EOF, or terminate=true to stop the "
"process. Use wait_for with wait_timeout_ms for dev servers, test "
"watchers, and prompts where you need to wait for expected output. "
"Do not use this to start new commands; start them with exec."
)
async def execute( # pyright: ignore[reportIncompatibleMethodOverride] async def execute( # pyright: ignore[reportIncompatibleMethodOverride]
self, self,
session_id: str, session_id: str,
input: str | None = None, chars: str | None = None,
close_stdin: bool = False, close_stdin: bool = False,
terminate: bool = False, terminate: bool = False,
yield_time_ms: int | None = None,
wait_for: str | None = None, wait_for: str | None = None,
until_exit: bool = False, wait_timeout_ms: int | None = None,
timeout_ms: int | None = None, max_output_chars: int | None = None,
max_output_tokens: int | None = None,
**kwargs: Any, **kwargs: Any,
) -> str: ) -> str:
try: try:
if wait_for == "": if max_output_chars is None:
return ToolResult.error("Error: wait_for must not be empty.") max_output_chars = max_output_tokens
if wait_for is not None and until_exit: output_limit = clamp_session_int(
return ToolResult.error( max_output_chars,
"Error: wait_for and until_exit are mutually exclusive." DEFAULT_MAX_OUTPUT_CHARS,
) 1000,
if terminate: MAX_OUTPUT_CHARS,
if any( )
( if wait_for:
input is not None, return await self._wait_for_output(
close_stdin,
wait_for is not None,
until_exit,
timeout_ms is not None,
)
):
return ToolResult.error("Error: terminate must be used alone.")
poll = await self._manager.write(
session_id=session_id, session_id=session_id,
chars=None, chars=chars,
close_stdin=False, close_stdin=close_stdin,
terminate=True, terminate=terminate,
yield_time_ms=0, wait_for=wait_for,
max_output_chars=DEFAULT_MAX_OUTPUT_CHARS, wait_timeout_ms=clamp_session_int(
owner_session_key=current_request_session_key(), wait_timeout_ms,
DEFAULT_WAIT_FOR_MS,
0,
MAX_WAIT_FOR_MS,
),
max_output_chars=output_limit,
) )
result = format_session_poll(session_id, poll) poll = await self._manager.write(
return ToolResult.error(result) if poll.timed_out else result
default_timeout_ms = (
DEFAULT_UNTIL_EXIT_MS
if until_exit
else DEFAULT_WAIT_FOR_MS
if wait_for is not None
else DEFAULT_YIELD_MS
)
return await self._wait(
session_id=session_id, session_id=session_id,
input=input, chars=chars,
close_stdin=close_stdin, close_stdin=close_stdin,
wait_for=wait_for, terminate=terminate,
until_exit=until_exit, yield_time_ms=clamp_session_int(yield_time_ms, DEFAULT_YIELD_MS, 0, MAX_YIELD_MS),
timeout_ms=clamp_session_int( max_output_chars=output_limit,
timeout_ms, owner_session_key=current_request_session_key(),
default_timeout_ms,
0,
MAX_WAIT_FOR_MS,
),
) )
result = format_session_poll(session_id, poll)
return ToolResult.error(result) if poll.timed_out else result
except KeyError: except KeyError:
return ToolResult.error(f"Error: exec session not found: {session_id!r}") return ToolResult.error(f"Error: exec session not found: {session_id!r}")
except Exception as exc: except Exception as exc:
return ToolResult.error(f"Error managing exec session: {exc}") return ToolResult.error(f"Error writing to exec session: {exc}")
async def _wait( async def _wait_for_output(
self, self,
*, *,
session_id: str, session_id: str,
input: str | None, chars: str | None,
close_stdin: bool, close_stdin: bool,
wait_for: str | None, terminate: bool,
until_exit: bool, wait_for: str,
timeout_ms: int, wait_timeout_ms: int,
max_output_chars: int,
) -> str: ) -> str:
deadline = time.monotonic() + (timeout_ms / 1000) deadline = time.monotonic() + (wait_timeout_ms / 1000)
aggregate = _BoundedOutputBuffer(DEFAULT_MAX_OUTPUT_CHARS) aggregate = _BoundedOutputBuffer(max_output_chars)
upstream_truncated = 0 upstream_truncated = 0
search_overlap = "" search_overlap = ""
first = True first = True
matched = False poll: _SessionPoll | None = None
while True: while True:
remaining_ms = max(0, int((deadline - time.monotonic()) * 1000)) remaining_ms = max(0, int((deadline - time.monotonic()) * 1000))
step_ms = min(MAX_YIELD_MS if until_exit else 500, remaining_ms) step_ms = min(500, remaining_ms)
poll = await self._manager.write( poll = await self._manager.write(
session_id=session_id, session_id=session_id,
chars=input if first else None, chars=chars if first else None,
close_stdin=close_stdin if first else False, close_stdin=close_stdin if first else False,
terminate=False, terminate=terminate if first else False,
yield_time_ms=step_ms, yield_time_ms=step_ms,
max_output_chars=MAX_OUTPUT_CHARS, max_output_chars=MAX_OUTPUT_CHARS,
owner_session_key=current_request_session_key(), owner_session_key=current_request_session_key(),
@@ -665,25 +671,20 @@ class ExecSessionTool(Tool):
upstream_truncated += poll.truncated_chars upstream_truncated += poll.truncated_chars
if poll.output: if poll.output:
aggregate.append(poll.output) aggregate.append(poll.output)
if wait_for is not None: searchable = search_overlap + poll.output
searchable = search_overlap + poll.output if wait_for in searchable:
matched = wait_for in searchable poll.output, aggregate_truncated = aggregate.drain()
overlap_chars = len(wait_for) - 1 poll.truncated_chars = upstream_truncated + aggregate_truncated
search_overlap = searchable[-overlap_chars:] if overlap_chars else "" result = format_session_poll(session_id, poll)
return ToolResult.error(result) if poll.timed_out else result
expired = time.monotonic() >= deadline overlap_chars = max(0, len(wait_for) - 1)
has_activity = wait_for is None and not until_exit and bool(poll.output) search_overlap = searchable[-overlap_chars:] if overlap_chars else ""
if poll.done or matched or has_activity or expired: if poll.done or remaining_ms <= 0:
poll.output, aggregate_truncated = aggregate.drain() poll.output, aggregate_truncated = aggregate.drain()
poll.truncated_chars = upstream_truncated + aggregate_truncated poll.truncated_chars = upstream_truncated + aggregate_truncated
result = format_session_poll(session_id, poll) result = format_session_poll(session_id, poll)
if wait_for is not None and not matched: if wait_for not in poll.output:
result += f"\nWait target not observed: {wait_for!r}" result += f"\nWait target not observed: {wait_for!r}"
elif until_exit and not poll.done:
result += (
f"\nWait timed out after {timeout_ms / 1000:g}s; "
"session remains active."
)
return ToolResult.error(result) if poll.timed_out else result return ToolResult.error(result) if poll.timed_out else result
@@ -721,7 +722,12 @@ class ListExecSessionsTool(Tool):
@property @property
def description(self) -> str: def description(self) -> str:
return "List active exec sessions." return (
"List active long-running exec sessions, including session_id, cwd, "
"elapsed time, idle time, remaining timeout, and command preview. "
"Use this to recover a session_id after context shifts before "
"polling, writing stdin, or terminating with write_stdin."
)
@property @property
def read_only(self) -> bool: def read_only(self) -> bool:
-293
View File
@@ -1,293 +0,0 @@
"""Execute tool calls and turn their outcomes into model observations."""
from __future__ import annotations
import asyncio
from collections.abc import Callable
from typing import Any, cast
from loguru import logger
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
from nanobot.providers.base import ToolCallRequest
from nanobot.utils.runtime import (
repeated_external_lookup_error,
repeated_workspace_violation_error,
)
_RETRY_HINT = "\n\n[Analyze the error above and try a different approach.]"
# SSRF is a hard security block at the tool boundary, but the agent turn
# should recover conversationally instead of aborting the runtime.
_SSRF_MARKERS: tuple[str, ...] = (
"internal/private url detected",
"private/internal address",
"private address",
)
_SSRF_BOUNDARY_NOTE = (
"This is a non-bypassable security boundary. Stop trying to access "
"private/internal URLs. Do not retry with curl, wget, encoded IPs, "
"alternate DNS, redirects, proxies, or another tool. Ask the user for "
"local files, logs, screenshots, or an explicit safe public URL instead. "
"If the user explicitly trusts this private URL, ask them to whitelist "
"the exact IP/CIDR via tools.ssrfWhitelist."
)
# Non-SSRF boundary markers returned to the model as recoverable tool errors.
_WORKSPACE_VIOLATION_MARKERS: tuple[str, ...] = (
"outside the configured workspace",
"outside allowed directory",
"working_dir is outside",
"working_dir could not be resolved",
"path outside working dir",
"path traversal detected",
)
def _with_retry_hint(payload: str) -> str:
"""Append the recovery hint exactly once."""
if payload.endswith(_RETRY_HINT):
return payload
return payload + _RETRY_HINT
async def execute_tool_calls(
tools: ToolRegistry,
tool_calls: list[ToolCallRequest],
*,
concurrent: bool,
external_lookup_counts: dict[str, int],
workspace_violation_counts: dict[str, int],
hook: AgentHook,
context: AgentHookContext,
) -> tuple[list[Any], list[dict[str, str]]]:
"""Execute one model response's tool calls in stable result order."""
tool_results: list[tuple[Any, dict[str, str]]] = []
for batch in _partition_tool_batches(tools, tool_calls, concurrent=concurrent):
if concurrent and len(batch) > 1:
batch_results = await asyncio.gather(*(
_execute_tool_call(
tools,
tool_call,
external_lookup_counts,
workspace_violation_counts,
hook,
context,
)
for tool_call in batch
))
tool_results.extend(batch_results)
else:
for tool_call in batch:
result = await _execute_tool_call(
tools,
tool_call,
external_lookup_counts,
workspace_violation_counts,
hook,
context,
)
tool_results.append(result)
results = [result for result, _event in tool_results]
events = [event for _result, event in tool_results]
return results, events
async def _execute_tool_call(
tools: ToolRegistry,
tool_call: ToolCallRequest,
external_lookup_counts: dict[str, int],
workspace_violation_counts: dict[str, int],
hook: AgentHook,
context: AgentHookContext,
) -> tuple[Any, dict[str, str]]:
lookup_error = repeated_external_lookup_error(
tool_call.name,
tool_call.arguments,
external_lookup_counts,
)
if lookup_error:
event = {
"name": tool_call.name,
"status": "error",
"detail": "repeated external lookup blocked",
}
return _with_retry_hint(lookup_error), event
prepare_call = cast(
Callable[[str, Any], object] | None,
getattr(tools, "prepare_call", None),
)
tool, params, prep_error = None, tool_call.arguments, None
if callable(prepare_call):
prepared = prepare_call(tool_call.name, tool_call.arguments)
if isinstance(prepared, tuple):
prepared_tuple = cast(tuple[object, ...], prepared)
if len(prepared_tuple) == 3:
tool, params, prep_error = cast(tuple[Any, Any, str | None], prepared_tuple)
if prep_error:
payload = _with_retry_hint(prep_error)
event = {
"name": tool_call.name,
"status": "error",
"detail": prep_error.split(": ", 1)[-1][:120],
}
handled = _classify_violation(
raw_text=prep_error,
soft_payload=payload,
event=event,
tool_call=tool_call,
workspace_violation_counts=workspace_violation_counts,
)
if handled is not None:
return handled
return payload, event
await hook.before_execute_tool(context, tool_call, tool, params)
try:
if tool is not None:
result = await tool.execute(**params)
else:
result = await tools.execute(tool_call.name, params)
except asyncio.CancelledError:
raise
except Exception as exc:
await hook.on_execute_tool_error(context, tool_call, tool, params, exc)
event = {
"name": tool_call.name,
"status": "error",
"detail": str(exc),
}
payload = _with_retry_hint(f"Error: {type(exc).__name__}: {exc}")
handled = _classify_violation(
raw_text=str(exc),
soft_payload=payload,
event=event,
tool_call=tool_call,
workspace_violation_counts=workspace_violation_counts,
)
if handled is not None:
return handled
return payload, event
if is_tool_error_result(result):
await hook.on_execute_tool_error(context, tool_call, tool, params, result)
payload = _with_retry_hint(result)
event = {
"name": tool_call.name,
"status": "error",
"detail": result.replace("\n", " ").strip()[:120],
}
handled = _classify_violation(
raw_text=result,
soft_payload=payload,
event=event,
tool_call=tool_call,
workspace_violation_counts=workspace_violation_counts,
)
if handled is not None:
return handled
return payload, event
await hook.after_execute_tool(context, tool_call, tool, params, result)
detail = "" if result is None else str(result)
detail = detail.replace("\n", " ").strip()
if not detail:
detail = "(empty)"
elif len(detail) > 120:
detail = detail[:120] + "..."
return result, {"name": tool_call.name, "status": "ok", "detail": detail}
def is_ssrf_violation(text: str) -> bool:
"""Return whether a tool error describes a blocked private-network request."""
if not text:
return False
lowered = text.lower()
return any(marker in lowered for marker in _SSRF_MARKERS)
def _is_workspace_violation(text: str) -> bool:
"""Return whether text describes any workspace or network boundary rejection."""
if not text:
return False
lowered = text.lower()
if is_ssrf_violation(lowered):
return True
return any(marker in lowered for marker in _WORKSPACE_VIOLATION_MARKERS)
def _classify_violation(
*,
raw_text: str,
soft_payload: str,
event: dict[str, str],
tool_call: ToolCallRequest,
workspace_violation_counts: dict[str, int],
) -> tuple[Any, dict[str, str]] | None:
if is_ssrf_violation(raw_text):
logger.warning(
"Tool {} blocked by SSRF guard; returning non-retryable tool error: {}",
tool_call.name,
raw_text.replace("\n", " ").strip()[:200],
)
event["detail"] = _event_detail("ssrf_violation: ", raw_text)
return _ssrf_soft_payload(raw_text), event
if _is_workspace_violation(raw_text):
escalation = repeated_workspace_violation_error(
tool_call.name,
tool_call.arguments,
workspace_violation_counts,
)
event["detail"] = _event_detail("workspace_violation: ", raw_text)
if escalation is not None:
logger.warning(
"Tool {} hit workspace boundary repeatedly; escalating hint",
tool_call.name,
)
event["detail"] = _event_detail(
"workspace_violation_escalated: ",
raw_text,
)
return escalation, event
return soft_payload, event
return None
def _ssrf_soft_payload(raw_text: str) -> str:
text = raw_text.strip() or "Error: request blocked by SSRF guard"
return f"{text}\n\n{_SSRF_BOUNDARY_NOTE}"
def _event_detail(prefix: str, text: str, limit: int = 160) -> str:
return (prefix + text.replace("\n", " ").strip())[:limit]
def _partition_tool_batches(
tools: ToolRegistry,
tool_calls: list[ToolCallRequest],
*,
concurrent: bool,
) -> list[list[ToolCallRequest]]:
if not concurrent:
return [[tool_call] for tool_call in tool_calls]
batches: list[list[ToolCallRequest]] = []
current: list[ToolCallRequest] = []
for tool_call in tool_calls:
get_tool = cast(Callable[[str], Any] | None, getattr(tools, "get", None))
tool = get_tool(tool_call.name) if callable(get_tool) else None
can_batch = bool(tool and tool.concurrency_safe)
if can_batch:
current.append(tool_call)
continue
if current:
batches.append(current)
current = []
batches.append([tool_call])
if current:
batches.append(current)
return batches
+135 -111
View File
@@ -2,9 +2,11 @@
# pyright: reportPrivateUsage=false, reportUnusedFunction=false # pyright: reportPrivateUsage=false, reportUnusedFunction=false
import asyncio
import difflib import difflib
import mimetypes import mimetypes
import os import os
import threading
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -21,6 +23,7 @@ from nanobot.agent.tools.schema import (
) )
from nanobot.config_base import Base from nanobot.config_base import Base
from nanobot.security.workspace_access import current_tool_workspace from nanobot.security.workspace_access import current_tool_workspace
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime from nanobot.utils.helpers import build_image_content_blocks, detect_image_mime
@@ -251,16 +254,16 @@ def _builtin_skill_read_path(path: str) -> Path | None:
tool_parameters_schema( tool_parameters_schema(
path=StringSchema("The file path to read"), path=StringSchema("The file path to read"),
offset=IntegerSchema( offset=IntegerSchema(
description="1-based text or extracted-document line (default 1)", description="Line number to start reading from (1-indexed, default 1)",
minimum=1, minimum=1,
), ),
limit=IntegerSchema( limit=IntegerSchema(
description="Maximum lines to return (default 2000)", description="Maximum number of lines to read (default 2000)",
minimum=1, minimum=1,
), ),
pages=StringSchema("PDF page number or range, e.g. '7' or '1-5' (max 20 pages)"), pages=StringSchema("Page range for PDF files, e.g. '1-5' (default: all, max 20 pages)"),
force=BooleanSchema( force=BooleanSchema(
description="Return an unchanged range again", description="Bypass same-file read deduplication and return content again.",
default=False, default=False,
), ),
required=["path"], required=["path"],
@@ -282,8 +285,18 @@ class ReadFileTool(_FsTool):
@property @property
def description(self) -> str: def description(self) -> str:
return ( return (
"Read text, images, PDFs, and Office documents by path. " "Read a file (text, image, or document). "
"Text is line-numbered; use offset/limit or pages for targeted ranges." "Text output format: LINE_NUM|CONTENT. "
"Images return visual content for analysis. "
"Supports PDF, DOCX, XLSX, PPTX documents. "
"Uploaded non-image attachments are referenced by path; read them "
"with this tool only when their contents are needed. "
"Use find_files/list_dir first when the path is uncertain. "
"Read the relevant range before editing so replacements or patches "
"are based on current content. "
"Use offset and limit for large text files. "
"Use force=true to re-read content even if unchanged. "
"Reads exceeding ~128K chars are truncated."
) )
@property @property
@@ -332,7 +345,7 @@ class ReadFileTool(_FsTool):
# Office document support # Office document support
if fp.suffix.lower() in {".docx", ".xlsx", ".pptx"}: if fp.suffix.lower() in {".docx", ".xlsx", ".pptx"}:
return self._read_office_doc(fp, offset, limit) return self._read_office_doc(fp)
raw = fp.read_bytes() raw = fp.read_bytes()
if not raw: if not raw:
@@ -454,8 +467,8 @@ class ReadFileTool(_FsTool):
max_pages=self._MAX_PDF_PAGES, max_pages=self._MAX_PDF_PAGES,
max_chars=self._MAX_CHARS, max_chars=self._MAX_CHARS,
) )
except PdfPageRangeError as e: except PdfPageRangeError:
return ToolResult.error(f"Error: Invalid page range '{pages}': {e!s}.") return ToolResult.error(f"Error: Invalid page range '{pages}'. Use format like '1-5'.")
except PdfSafetyError as e: except PdfSafetyError as e:
return ToolResult.error(f"Error reading PDF: {e}") return ToolResult.error(f"Error reading PDF: {e}")
except Exception as e: except Exception as e:
@@ -474,85 +487,24 @@ class ReadFileTool(_FsTool):
) )
return result return result
def _read_office_doc( def _read_office_doc(self, fp: Path) -> str:
self, from nanobot.utils.document import extract_text
fp: Path,
offset: int,
limit: int | None,
) -> str:
from nanobot.utils.document import open_document_line_source
offset = max(1, offset) result = extract_text(fp)
requested_limit = limit or self._DEFAULT_LIMIT
source_iterator = None
try:
source = open_document_line_source(fp)
if source is None:
return ToolResult.error(f"Error: Unsupported file format: {fp.suffix}")
source_iterator = source.lines
numbered: list[str] = []
output_chars = 0
total_seen = 0
end = offset - 1
has_more = False
line_was_clipped = False
for line in source_iterator: if result is None:
total_seen = line.extracted_line return ToolResult.error(f"Error: Unsupported file format: {fp.suffix}")
if line.extracted_line < offset:
continue
if len(numbered) >= requested_limit:
has_more = True
break
rendered = f"{line.extracted_line}| {line.text}" if result.startswith("[error:"):
extra = 1 if numbered else 0 return ToolResult.error(f"Error reading {fp.suffix.upper()} file: {result}")
if output_chars + extra + len(rendered) > self._MAX_CHARS:
if numbered:
has_more = True
break
prefix = f"{line.extracted_line}| "
available = max(0, self._MAX_CHARS - len(prefix) - 3)
rendered = f"{prefix}{line.text[:available]}..."
line_was_clipped = True
has_more = True
numbered.append(rendered)
output_chars += extra + len(rendered)
end = line.extracted_line
if line_was_clipped:
break
if not numbered: if not result:
if total_seen == 0: return f"({fp.suffix.upper().lstrip('.')} has no extractable text: {fp})"
return (
f"({fp.suffix.upper().lstrip('.')} has no extractable text: {fp})"
)
return ToolResult.error(
f"Error: offset {offset} is beyond end of extracted document "
f"({total_seen} lines)"
)
output = "\n".join(numbered) if len(result) > self._MAX_CHARS:
if has_more: result = result[:self._MAX_CHARS] + "\n\n(Document text truncated at ~128K chars)"
if line_was_clipped:
output += ( return result
"\n\n(Document text truncated at ~128K chars; line clipped. "
f"Use offset={end + 1} to continue.)"
)
else:
output += (
f"\n\n(Showing extracted lines {offset}-{end}. "
f"Use offset={end + 1} to continue.)"
)
else:
output += f"\n\n(End of document — {total_seen} extracted lines total)"
return output
except Exception as e:
return ToolResult.error(f"Error reading {fp.suffix.upper()} file: {e!s}")
finally:
close = getattr(source_iterator, "close", None)
if close is not None:
close()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -715,22 +667,31 @@ def _match_covers_line(match: _MatchSpan, line: int) -> bool:
return match.line <= line <= _match_end_line(match) return match.line <= line <= _match_end_line(match)
def _find_exact_matches(content: str, old_text: str) -> list[_MatchSpan]: def _find_exact_matches(
content: str,
old_text: str,
*,
max_matches: int | None = None,
) -> list[_MatchSpan]:
matches: list[_MatchSpan] = [] matches: list[_MatchSpan] = []
start = 0 search_start = 0
while True: line_start = 0
idx = content.find(old_text, start) line = 1
while max_matches is None or len(matches) < max_matches:
idx = content.find(old_text, search_start)
if idx == -1: if idx == -1:
break break
line += content.count("\n", line_start, idx)
matches.append( matches.append(
_MatchSpan( _MatchSpan(
start=idx, start=idx,
end=idx + len(old_text), end=idx + len(old_text),
text=content[idx : idx + len(old_text)], text=content[idx : idx + len(old_text)],
line=content.count("\n", 0, idx) + 1, line=line,
) )
) )
start = idx + max(1, len(old_text)) line_start = idx
search_start = idx + max(1, len(old_text))
return matches return matches
@@ -786,27 +747,36 @@ def _find_quote_matches(content: str, old_text: str) -> list[_MatchSpan]:
norm_content = _normalize_quotes(content) norm_content = _normalize_quotes(content)
norm_old = _normalize_quotes(old_text) norm_old = _normalize_quotes(old_text)
matches: list[_MatchSpan] = [] matches: list[_MatchSpan] = []
start = 0 search_start = 0
line_start = 0
line = 1
while True: while True:
idx = norm_content.find(norm_old, start) idx = norm_content.find(norm_old, search_start)
if idx == -1: if idx == -1:
break break
line += content.count("\n", line_start, idx)
matches.append( matches.append(
_MatchSpan( _MatchSpan(
start=idx, start=idx,
end=idx + len(old_text), end=idx + len(old_text),
text=content[idx : idx + len(old_text)], text=content[idx : idx + len(old_text)],
line=content.count("\n", 0, idx) + 1, line=line,
) )
) )
start = idx + max(1, len(norm_old)) line_start = idx
search_start = idx + max(1, len(norm_old))
return matches return matches
def _find_matches(content: str, old_text: str) -> list[_MatchSpan]: def _find_matches(
"""Locate all matches using progressively looser strategies.""" content: str,
old_text: str,
*,
max_exact_matches: int | None = None,
) -> list[_MatchSpan]:
"""Locate matches using progressively looser strategies."""
for matcher in ( for matcher in (
lambda: _find_exact_matches(content, old_text), lambda: _find_exact_matches(content, old_text, max_matches=max_exact_matches),
lambda: _find_trim_matches(content, old_text), lambda: _find_trim_matches(content, old_text),
lambda: _find_trim_matches(content, old_text, normalize_quotes=True), lambda: _find_trim_matches(content, old_text, normalize_quotes=True),
lambda: _find_quote_matches(content, old_text), lambda: _find_quote_matches(content, old_text),
@@ -861,10 +831,8 @@ def _best_window(old_text: str, content: str) -> tuple[float, int, list[str], li
@tool_parameters( @tool_parameters(
tool_parameters_schema( tool_parameters_schema(
path=StringSchema("The file path to edit"), path=StringSchema("The file path to edit"),
old_text=StringSchema("The text to find and replace; copy it from read_file."), old_text=StringSchema("The text to find and replace"),
new_text=StringSchema( new_text=StringSchema("The text to replace with"),
"The replacement text; must differ from old_text for an existing file."
),
replace_all=BooleanSchema(description="Replace all occurrences (default false)"), replace_all=BooleanSchema(description="Replace all occurrences (default false)"),
occurrence=IntegerSchema( occurrence=IntegerSchema(
description="Optional 1-based occurrence to replace when old_text appears multiple times.", description="Optional 1-based occurrence to replace when old_text appears multiple times.",
@@ -901,9 +869,15 @@ class EditFileTool(_FsTool):
@property @property
def description(self) -> str: def description(self) -> str:
return ( return (
"Perform a small, exact replacement in one file. " "Perform a small, exact replacement in one file by replacing "
"Prefer apply_patch for multi-file, structural, or generated edits. " "old_text with new_text. When replacing text in an existing file, "
"occurrence, line_hint, and replace_all=true are mutually exclusive." "old_text and new_text must be different. Use this for narrow text substitutions "
"with old_text copied from read_file. For multi-file, structural, "
"or generated code edits, prefer apply_patch. If old_text matches "
"multiple times, provide more context or set occurrence, line_hint, "
"replace_all, and expected_replacements. When editing from numbered "
"read_file output, set line_hint to the exact target line. "
"Shows closest-match diagnostics on failure."
) )
@staticmethod @staticmethod
@@ -916,6 +890,43 @@ class EditFileTool(_FsTool):
new_text: str | None = None, new_text: str | None = None,
replace_all: bool = False, occurrence: int | None = None, replace_all: bool = False, occurrence: int | None = None,
line_hint: int | None = None, expected_replacements: int | None = None, **kwargs: Any, line_hint: int | None = None, expected_replacements: int | None = None, **kwargs: Any,
) -> str:
cancelled = threading.Event()
commit_lock = threading.Lock()
try:
return await asyncio.to_thread(
self._execute_sync,
path=path,
old_text=old_text,
new_text=new_text,
replace_all=replace_all,
occurrence=occurrence,
line_hint=line_hint,
expected_replacements=expected_replacements,
cancelled=cancelled,
commit_lock=commit_lock,
)
except asyncio.CancelledError:
cancelled.set()
# If a commit already started, do not report cancellation until the
# file bytes and FileStates record are settled. Otherwise, taking
# the lock first guarantees the worker observes ``cancelled`` before
# it can mutate the target.
await shield_and_drain(
asyncio.to_thread(self._wait_for_commit, commit_lock)
)
raise
@staticmethod
def _wait_for_commit(commit_lock: threading.Lock) -> None:
with commit_lock:
pass
def _execute_sync(
self, *, path: str | None, old_text: str | None,
new_text: str | None, replace_all: bool, occurrence: int | None,
line_hint: int | None, expected_replacements: int | None,
cancelled: threading.Event, commit_lock: threading.Lock,
) -> str: ) -> str:
try: try:
if not path: if not path:
@@ -939,9 +950,12 @@ class EditFileTool(_FsTool):
# Create-file semantics: old_text='' + file doesn't exist → create # Create-file semantics: old_text='' + file doesn't exist → create
if not file_exists: if not file_exists:
if old_text == "": if old_text == "":
fp.parent.mkdir(parents=True, exist_ok=True) with commit_lock:
fp.write_text(new_text, encoding="utf-8") if cancelled.is_set():
self._file_states.record_write(fp) return ToolResult.error("Error: edit_file cancelled.")
fp.parent.mkdir(parents=True, exist_ok=True)
fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp)
return f"Successfully created {fp}" return f"Successfully created {fp}"
return self._file_not_found_msg(path, fp) return self._file_not_found_msg(path, fp)
@@ -959,8 +973,11 @@ class EditFileTool(_FsTool):
content = raw.decode("utf-8") content = raw.decode("utf-8")
if content.strip(): if content.strip():
return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.") return ToolResult.error(f"Error: Cannot create file — {path} already exists and is not empty.")
fp.write_text(new_text, encoding="utf-8") with commit_lock:
self._file_states.record_write(fp) if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.write_text(new_text, encoding="utf-8")
self._file_states.record_write(fp)
return f"Successfully edited {fp}" return f"Successfully edited {fp}"
# Read-before-edit check # Read-before-edit check
@@ -970,7 +987,11 @@ class EditFileTool(_FsTool):
uses_crlf = b"\r\n" in raw uses_crlf = b"\r\n" in raw
content = raw.decode("utf-8").replace("\r\n", "\n") content = raw.decode("utf-8").replace("\r\n", "\n")
norm_old = old_text.replace("\r\n", "\n") norm_old = old_text.replace("\r\n", "\n")
matches = _find_matches(content, norm_old) matches = _find_matches(
content,
norm_old,
max_exact_matches=occurrence,
)
if not matches: if not matches:
return self._not_found_msg(old_text, content, path) return self._not_found_msg(old_text, content, path)
@@ -1047,8 +1068,11 @@ class EditFileTool(_FsTool):
if uses_crlf: if uses_crlf:
new_content = new_content.replace("\n", "\r\n") new_content = new_content.replace("\n", "\r\n")
fp.write_bytes(new_content.encode("utf-8")) with commit_lock:
self._file_states.record_write(fp) if cancelled.is_set():
return ToolResult.error("Error: edit_file cancelled.")
fp.write_bytes(new_content.encode("utf-8"))
self._file_states.record_write(fp)
msg = f"Successfully edited {fp}" msg = f"Successfully edited {fp}"
if warning: if warning:
msg = f"{warning}\n{msg}" msg = f"{warning}\n{msg}"
+52 -22
View File
@@ -17,6 +17,7 @@ from nanobot.agent.tools.context import RequestContext, ToolContext, current_req
from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema from nanobot.agent.tools.schema import StringSchema, tool_parameters_schema
from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext from nanobot.bus.runtime_events import GoalStateChanged, RuntimeEventBus, RuntimeEventContext
from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_lines from nanobot.runtime_context import RuntimeContextBlock, wrap_runtime_context_lines
from nanobot.session.async_compat import call_session_manager
from nanobot.session.goal_state import ( from nanobot.session.goal_state import (
GOAL_STATE_KEY, GOAL_STATE_KEY,
MAX_GOAL_OBJECTIVE_CHARS, MAX_GOAL_OBJECTIVE_CHARS,
@@ -28,6 +29,7 @@ from nanobot.session.goal_state import (
sustained_goal_active, sustained_goal_active,
) )
from nanobot.session.turn_continuation import reset_goal_continuation_rounds from nanobot.session.turn_continuation import reset_goal_continuation_rounds
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.prompt_templates import render_template from nanobot.utils.prompt_templates import render_template
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -60,36 +62,68 @@ class _GoalToolsMixin:
self._sessions = sessions self._sessions = sessions
self._runtime_events = runtime_events self._runtime_events = runtime_events
def _session(self): async def _get_or_create_session(self, key: str):
return await call_session_manager(
self._sessions,
"get_or_create_async",
self._sessions.get_or_create,
key,
)
async def _save_session(self, session: Any) -> None:
await call_session_manager(
self._sessions,
"save_async",
self._sessions.save,
session,
)
async def _session(self):
request_ctx = current_request_context() request_ctx = current_request_context()
if request_ctx is None: if request_ctx is None:
return None return None
key = request_ctx.session_key key = request_ctx.session_key
if not key: if not key:
return None return None
return self._sessions.get_or_create(key) return await self._get_or_create_session(key)
def _goal_mutation_allowed(self) -> bool: def _goal_mutation_allowed(self) -> bool:
return current_request_context() is not None and goal_mutation_allowed() return current_request_context() is not None and goal_mutation_allowed()
def _save_goal_state( async def _save_goal_state(
self, self,
sess: Any, sess: Any,
blob: dict[str, Any], blob: dict[str, Any],
*, *,
reset_continuation: bool = False, reset_continuation: bool = False,
revoke_permission: bool = False,
) -> None: ) -> None:
previous_metadata = deepcopy(sess.metadata) previous_metadata = deepcopy(sess.metadata)
sess.metadata[GOAL_STATE_KEY] = blob saved = False
discard_legacy_goal_state_key(sess.metadata)
if reset_continuation: async def save_and_publish() -> None:
reset_goal_continuation_rounds(sess.metadata) nonlocal saved
sess.metadata[GOAL_STATE_KEY] = blob
discard_legacy_goal_state_key(sess.metadata)
if reset_continuation:
reset_goal_continuation_rounds(sess.metadata)
try:
await self._save_session(sess)
except BaseException:
sess.metadata.clear()
sess.metadata.update(previous_metadata)
raise
saved = True
await self._publish_goal_state_changed(sess.metadata)
try: try:
self._sessions.save(sess) await shield_and_drain(save_and_publish())
except BaseException: finally:
sess.metadata.clear() # This ContextVar belongs to the caller task, not the settlement task.
sess.metadata.update(previous_metadata) # Apply the post-save permission effect here even when cancellation was
raise # delayed until the durable save and runtime notification completed.
if revoke_permission and saved:
revoke_goal_mutation_permission()
async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None: async def _publish_goal_state_changed(self, metadata: dict[str, Any]) -> None:
runtime_events = self._runtime_events runtime_events = self._runtime_events
@@ -175,7 +209,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
) -> RuntimeContextBlock | None: ) -> RuntimeContextBlock | None:
if not request.session_key: if not request.session_key:
return None return None
session = self._sessions.get_or_create(request.session_key) session = await self._get_or_create_session(request.session_key)
goal_start_requested = explicit_goal_requested(request.metadata) goal_start_requested = explicit_goal_requested(request.metadata)
goal_active = sustained_goal_active(session.metadata) goal_active = sustained_goal_active(session.metadata)
if not goal_start_requested and not goal_active: if not goal_start_requested and not goal_active:
@@ -197,7 +231,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
ui_summary: str | None = None, ui_summary: str | None = None,
**kwargs: Any, **kwargs: Any,
) -> str: ) -> str:
sess = self._session() sess = await self._session()
if sess is None: if sess is None:
return ToolResult.error( return ToolResult.error(
"Error: create_goal requires an active chat session (missing routing context)." "Error: create_goal requires an active chat session (missing routing context)."
@@ -225,8 +259,7 @@ class CreateGoalTool(Tool, _GoalToolsMixin):
"ui_summary": summary, "ui_summary": summary,
"started_at": _iso_now(), "started_at": _iso_now(),
} }
self._save_goal_state(sess, blob, reset_continuation=True) await self._save_goal_state(sess, blob, reset_continuation=True)
await self._publish_goal_state_changed(sess.metadata)
extra = f"\nSummary line: {summary}" if summary else "" extra = f"\nSummary line: {summary}" if summary else ""
return ( return (
"Goal recorded. Keep working toward the objective using ordinary tools. " "Goal recorded. Keep working toward the objective using ordinary tools. "
@@ -305,7 +338,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
ui_summary: str | None = None, ui_summary: str | None = None,
**kwargs: Any, **kwargs: Any,
) -> str: ) -> str:
sess = self._session() sess = await self._session()
if sess is None: if sess is None:
return ToolResult.error("Error: update_goal requires an active chat session.") return ToolResult.error("Error: update_goal requires an active chat session.")
prior = parse_goal_state(goal_state_raw(sess.metadata)) prior = parse_goal_state(goal_state_raw(sess.metadata))
@@ -340,8 +373,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
"previous_objective": str(prior.get("objective") or ""), "previous_objective": str(prior.get("objective") or ""),
"recap": (recap or "").strip(), "recap": (recap or "").strip(),
} }
self._save_goal_state(sess, blob, reset_continuation=True) await self._save_goal_state(sess, blob, reset_continuation=True)
await self._publish_goal_state_changed(sess.metadata)
extra = f"\nSummary line: {summary}" if summary else "" extra = f"\nSummary line: {summary}" if summary else ""
return "Goal replaced. Continue toward the new objective using ordinary tools." + extra return "Goal replaced. Continue toward the new objective using ordinary tools." + extra
@@ -359,9 +391,7 @@ class UpdateGoalTool(Tool, _GoalToolsMixin):
} }
if normalized == "complete": if normalized == "complete":
blob["completed_at"] = ended blob["completed_at"] = ended
self._save_goal_state(sess, blob) await self._save_goal_state(sess, blob, revoke_permission=True)
revoke_goal_mutation_permission()
await self._publish_goal_state_changed(sess.metadata)
tail = (recap or "").strip() tail = (recap or "").strip()
label = { label = {
+5 -5
View File
@@ -20,10 +20,10 @@ from nanobot.agent.tools.base import Tool, ToolResult
from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.registry import ToolRegistry
from nanobot.security.network import ( from nanobot.security.network import (
PinnedDNSAsyncTransport, PinnedDNSAsyncTransport,
async_resolve_url_target,
async_validate_url_target,
env_proxy_applies_to_url, env_proxy_applies_to_url,
httpx_env_proxy_mounts, httpx_env_proxy_mounts,
resolve_url_target,
validate_url_target,
) )
from nanobot.utils.cancellation import task_is_cancelling from nanobot.utils.cancellation import task_is_cancelling
@@ -249,7 +249,7 @@ async def _probe_http_url(url: str, timeout: float = 3.0) -> bool:
port = parsed.port port = parsed.port
if not port: if not port:
port = 443 if parsed.scheme == "https" else 80 port = 443 if parsed.scheme == "https" else 80
ok, _, resolved_ips = resolve_url_target(url) ok, _, resolved_ips = await async_resolve_url_target(url)
if not ok: if not ok:
return False return False
if env_proxy_applies_to_url(url): if env_proxy_applies_to_url(url):
@@ -298,7 +298,7 @@ def _pinned_transport_kwargs() -> dict[str, Any]:
async def _validate_mcp_request_url(request: httpx.Request) -> None: async def _validate_mcp_request_url(request: httpx.Request) -> None:
"""Validate each outgoing MCP HTTP request, including redirect targets.""" """Validate each outgoing MCP HTTP request, including redirect targets."""
ok, error = validate_url_target(str(request.url)) ok, error = await async_validate_url_target(str(request.url))
if not ok: if not ok:
raise httpx.RequestError( raise httpx.RequestError(
f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})", f"Blocked unsafe MCP URL {_redact_url(str(request.url))} ({error})",
@@ -1031,7 +1031,7 @@ async def connect_mcp_servers(
return False return False
if transport_type in {"sse", "streamableHttp"}: if transport_type in {"sse", "streamableHttp"}:
ok, error = validate_url_target(cfg.url) ok, error = await async_validate_url_target(cfg.url)
if not ok: if not ok:
logger.warning( logger.warning(
"MCP server '{}': blocked unsafe URL {} ({})", "MCP server '{}': blocked unsafe URL {} ({})",
+16 -22
View File
@@ -2,11 +2,9 @@
# pyright: reportIncompatibleMethodOverride=false # pyright: reportIncompatibleMethodOverride=false
from collections.abc import Awaitable, Callable, Generator
from contextlib import contextmanager
from contextvars import ContextVar, Token from contextvars import ContextVar, Token
from pathlib import Path from pathlib import Path
from typing import Any, cast from typing import Any, Awaitable, Callable, cast
from loguru import logger from loguru import logger
@@ -18,22 +16,6 @@ from nanobot.bus.events import OutboundMessage
from nanobot.config.paths import get_workspace_path from nanobot.config.paths import get_workspace_path
from nanobot.security.workspace_access import current_tool_workspace from nanobot.security.workspace_access import current_tool_workspace
_CURRENT_MESSAGE_SENDS: ContextVar[set[tuple[str, str]] | None] = ContextVar(
"message_sends",
default=None,
)
@contextmanager
def capture_message_deliveries() -> Generator[set[tuple[str, str]], None, None]:
"""Record successful MessageTool targets within one agent run."""
sends: set[tuple[str, str]] = set()
token = _CURRENT_MESSAGE_SENDS.set(sends)
try:
yield sends
finally:
_CURRENT_MESSAGE_SENDS.reset(token)
@tool_parameters( @tool_parameters(
tool_parameters_schema( tool_parameters_schema(
@@ -86,6 +68,7 @@ class MessageTool(Tool):
self._fallback_chat_id = default_chat_id self._fallback_chat_id = default_chat_id
self._fallback_message_id = default_message_id self._fallback_message_id = default_message_id
self._fallback_metadata: dict[str, Any] = {} self._fallback_metadata: dict[str, Any] = {}
self._sent_in_turn_var: ContextVar[bool] = ContextVar("message_sent_in_turn", default=False)
self._suppress_delivery_var: ContextVar[bool] = ContextVar( self._suppress_delivery_var: ContextVar[bool] = ContextVar(
"message_suppress_delivery", "message_suppress_delivery",
default=False, default=False,
@@ -104,6 +87,10 @@ class MessageTool(Tool):
"""Set the callback for sending messages.""" """Set the callback for sending messages."""
self._send_callback = callback self._send_callback = callback
def start_turn(self) -> None:
"""Reset per-turn send tracking."""
self._sent_in_turn = False
def set_suppress_delivery(self, active: bool) -> Token[bool]: def set_suppress_delivery(self, active: bool) -> Token[bool]:
"""Acknowledge but don't deliver tool sends (heartbeat internal check).""" """Acknowledge but don't deliver tool sends (heartbeat internal check)."""
return self._suppress_delivery_var.set(active) return self._suppress_delivery_var.set(active)
@@ -112,6 +99,14 @@ class MessageTool(Tool):
"""Restore previous delivery-suppression state.""" """Restore previous delivery-suppression state."""
self._suppress_delivery_var.reset(token) self._suppress_delivery_var.reset(token)
@property
def _sent_in_turn(self) -> bool:
return self._sent_in_turn_var.get()
@_sent_in_turn.setter
def _sent_in_turn(self, value: bool) -> None:
self._sent_in_turn_var.set(value)
@property @property
def name(self) -> str: def name(self) -> str:
return "message" return "message"
@@ -249,9 +244,8 @@ class MessageTool(Tool):
try: try:
await self._send_callback(msg) await self._send_callback(msg)
sends = _CURRENT_MESSAGE_SENDS.get() if channel == default_channel and chat_id == default_chat_id:
if sends is not None: self._sent_in_turn = True
sends.add((channel, chat_id))
media_info = f" with {len(media)} attachments" if media else "" media_info = f" with {len(media)} attachments" if media else ""
button_info = ( button_info = (
f" with {sum(len(row) for row in button_rows)} button(s)" f" with {sum(len(row) for row in button_rows)} button(s)"
+3 -3
View File
@@ -70,7 +70,7 @@ class ToolRegistry:
def has(self, name: str) -> bool: def has(self, name: str) -> bool:
"""Check if a tool is registered.""" """Check if a tool is registered."""
return self.get(name) is not None return name in self._tools
@staticmethod @staticmethod
def _schema_name(schema: dict[str, Any]) -> str: def _schema_name(schema: dict[str, Any]) -> str:
@@ -113,7 +113,7 @@ class ToolRegistry:
params: Any, params: Any,
) -> tuple[Tool | None, Any, str | None]: ) -> tuple[Tool | None, Any, str | None]:
"""Resolve, cast, and validate one tool call.""" """Resolve, cast, and validate one tool call."""
tool = self.get(name) tool = self._tools.get(name)
if not tool: if not tool:
suggestion = self._suggest_name(str(name)) suggestion = self._suggest_name(str(name))
hint = f" Did you mean '{suggestion}'? Tool names must match exactly." if suggestion else "" hint = f" Did you mean '{suggestion}'? Tool names must match exactly." if suggestion else ""
@@ -209,4 +209,4 @@ class ToolRegistry:
return len(self._tools) return len(self._tools)
def __contains__(self, name: str) -> bool: def __contains__(self, name: str) -> bool:
return self.has(name) return name in self._tools
+40
View File
@@ -12,6 +12,7 @@ if TYPE_CHECKING:
from nanobot.agent.tools.shell import ExecToolConfig from nanobot.agent.tools.shell import ExecToolConfig
from nanobot.agent.tools.web import WebToolsConfig from nanobot.agent.tools.web import WebToolsConfig
from nanobot.config.schema import ModelPresetConfig from nanobot.config.schema import ModelPresetConfig
from nanobot.providers.base import LLMUsage
from nanobot.utils.llm_runtime import LLMRuntime from nanobot.utils.llm_runtime import LLMRuntime
@@ -28,10 +29,13 @@ RUNTIME_SNAPSHOT_KEYS = frozenset({
"workspace", "workspace",
"provider_retry_mode", "provider_retry_mode",
"max_tool_result_chars", "max_tool_result_chars",
"current_iteration",
"_current_iteration",
"tool_names", "tool_names",
"web_config", "web_config",
"exec_config", "exec_config",
"subagents", "subagents",
"_last_usage",
}) })
RUNTIME_COMMAND_KEYS = frozenset({ RUNTIME_COMMAND_KEYS = frozenset({
@@ -57,10 +61,12 @@ class RuntimeSnapshot:
workspace: Path | str workspace: Path | str
provider_retry_mode: str provider_retry_mode: str
max_tool_result_chars: int max_tool_result_chars: int
current_iteration: int
tool_names: list[str] tool_names: list[str]
web_config: dict[str, object] web_config: dict[str, object]
exec_config: dict[str, object] exec_config: dict[str, object]
subagent_statuses: dict[str, dict[str, object]] subagent_statuses: dict[str, dict[str, object]]
last_usage: Mapping[str, JsonScalar]
scratchpad: dict[str, JsonValue] scratchpad: dict[str, JsonValue]
def as_mapping(self) -> Mapping[str, object]: def as_mapping(self) -> Mapping[str, object]:
@@ -74,10 +80,13 @@ class RuntimeSnapshot:
"workspace": self.workspace, "workspace": self.workspace,
"provider_retry_mode": self.provider_retry_mode, "provider_retry_mode": self.provider_retry_mode,
"max_tool_result_chars": self.max_tool_result_chars, "max_tool_result_chars": self.max_tool_result_chars,
"current_iteration": self.current_iteration,
"_current_iteration": self.current_iteration,
"tool_names": self.tool_names, "tool_names": self.tool_names,
"web_config": self.web_config, "web_config": self.web_config,
"exec_config": self.exec_config, "exec_config": self.exec_config,
"subagents": {"_task_statuses": self.subagent_statuses}, "subagents": {"_task_statuses": self.subagent_statuses},
"_last_usage": self.last_usage,
} }
assert values.keys() == RUNTIME_SNAPSHOT_KEYS assert values.keys() == RUNTIME_SNAPSHOT_KEYS
return values return values
@@ -98,6 +107,13 @@ class RuntimeControl(Protocol):
session_key: str | None, session_key: str | None,
) -> LLMRuntime: ... ) -> LLMRuntime: ...
async def set_model_preset_async(
self,
name: str,
*,
session_key: str | None,
) -> LLMRuntime: ...
def set_max_iterations(self, value: int) -> None: ... def set_max_iterations(self, value: int) -> None: ...
def set_context_window_tokens(self, value: int) -> LLMRuntime: ... def set_context_window_tokens(self, value: int) -> LLMRuntime: ...
@@ -136,9 +152,15 @@ class _RuntimeControlTarget(Protocol):
@property @property
def workspace(self) -> Path: ... def workspace(self) -> Path: ...
@property
def current_iteration(self) -> int: ...
@property @property
def tool_names(self) -> list[str]: ... def tool_names(self) -> list[str]: ...
@property
def last_usage(self) -> LLMUsage | None: ...
def set_runtime_model(self, model: str) -> LLMRuntime: ... def set_runtime_model(self, model: str) -> LLMRuntime: ...
def set_runtime_context_window(self, context_window_tokens: int) -> LLMRuntime: ... def set_runtime_context_window(self, context_window_tokens: int) -> LLMRuntime: ...
@@ -147,6 +169,12 @@ class _RuntimeControlTarget(Protocol):
def set_session_model_preset(self, session_key: str, name: str) -> LLMRuntime: ... def set_session_model_preset(self, session_key: str, name: str) -> LLMRuntime: ...
async def set_session_model_preset_async(
self,
session_key: str,
name: str,
) -> LLMRuntime: ...
class AgentRuntimeControl: class AgentRuntimeControl:
"""Allowlisted adapter from agent-loop state to ``RuntimeControl``.""" """Allowlisted adapter from agent-loop state to ``RuntimeControl``."""
@@ -171,10 +199,12 @@ class AgentRuntimeControl:
), ),
provider_retry_mode=target.provider_retry_mode, provider_retry_mode=target.provider_retry_mode,
max_tool_result_chars=target.max_tool_result_chars, max_tool_result_chars=target.max_tool_result_chars,
current_iteration=target.current_iteration,
tool_names=list(target.tool_names), tool_names=list(target.tool_names),
web_config=_snapshot_web_config(target.web_config), web_config=_snapshot_web_config(target.web_config),
exec_config=_snapshot_exec_config(target.exec_config), exec_config=_snapshot_exec_config(target.exec_config),
subagent_statuses=_snapshot_subagent_statuses(target.subagents), subagent_statuses=_snapshot_subagent_statuses(target.subagents),
last_usage=target.last_usage.to_dict() if target.last_usage is not None else {},
scratchpad=_snapshot_json_mapping(self.__scratchpad), scratchpad=_snapshot_json_mapping(self.__scratchpad),
) )
@@ -191,6 +221,16 @@ class AgentRuntimeControl:
return self.__target.set_session_model_preset(session_key, name) return self.__target.set_session_model_preset(session_key, name)
return self.__target.set_model_preset(name) return self.__target.set_model_preset(name)
async def set_model_preset_async(
self,
name: str,
*,
session_key: str | None,
) -> LLMRuntime:
if session_key is not None:
return await self.__target.set_session_model_preset_async(session_key, name)
return self.__target.set_model_preset(name)
def set_max_iterations(self, value: int) -> None: def set_max_iterations(self, value: int) -> None:
self.__target.max_iterations = value self.__target.max_iterations = value
self.__target.subagents.max_iterations = value self.__target.subagents.max_iterations = value
+171 -440
View File
@@ -6,28 +6,19 @@ from __future__ import annotations
import asyncio import asyncio
import fnmatch import fnmatch
import heapq
import os import os
import re import re
import threading import threading
import time import time
from collections import deque
from contextlib import suppress from contextlib import suppress
from dataclasses import dataclass
from pathlib import Path, PurePosixPath from pathlib import Path, PurePosixPath
from typing import Any, Iterable, Iterator, TypeVar from typing import Any, Iterable, TypeVar
from nanobot.agent.tools.base import ToolResult from nanobot.agent.tools.base import ToolResult
from nanobot.agent.tools.filesystem import ListDirTool, _FsTool from nanobot.agent.tools.filesystem import ListDirTool, _FsTool
from nanobot.utils.document import (
LocatedDocumentLine,
PdfPageRangeError,
open_document_line_source,
)
_DEFAULT_HEAD_LIMIT = 250 _DEFAULT_HEAD_LIMIT = 250
_DEFAULT_FILE_HEAD_LIMIT = 200 _DEFAULT_FILE_HEAD_LIMIT = 200
_DOCUMENT_EXTENSIONS = frozenset({".pdf", ".docx", ".xlsx", ".pptx"})
T = TypeVar("T") T = TypeVar("T")
_TYPE_GLOB_MAP = { _TYPE_GLOB_MAP = {
"py": ("*.py", "*.pyi"), "py": ("*.py", "*.pyi"),
@@ -53,51 +44,6 @@ _TYPE_GLOB_MAP = {
} }
@dataclass(slots=True)
class _PendingContextMatch:
lines: list[LocatedDocumentLine]
match_index: int
match_start: int
remaining_after: int
@dataclass(slots=True)
class _FindFilesEntry:
path: Path
rel_path: str
display_path: str
name: str
is_dir: bool
class _FindFilesCancelledError(Exception):
"""Stop a worker scan after its owning async task was cancelled."""
class _FindFilesBudgetExceededError(Exception):
"""Stop an unbounded filesystem scan at its configured budget."""
@dataclass(slots=True)
class _FindFilesBudget:
cancelled: threading.Event
deadline: float
max_paths: int
scanned_paths: int = 0
def checkpoint(self) -> None:
if self.cancelled.is_set():
raise _FindFilesCancelledError
if time.monotonic() >= self.deadline:
raise _FindFilesBudgetExceededError("time")
def visit_path(self) -> None:
self.checkpoint()
self.scanned_paths += 1
if self.scanned_paths > self.max_paths:
raise _FindFilesBudgetExceededError("paths")
def _normalize_pattern(pattern: str) -> str: def _normalize_pattern(pattern: str) -> str:
return pattern.strip().replace("\\", "/") return pattern.strip().replace("\\", "/")
@@ -121,15 +67,6 @@ def _is_binary(raw: bytes) -> bool:
return (non_text / len(sample)) > 0.2 return (non_text / len(sample)) > 0.2
def _excel_column(index: int) -> str:
"""Return a 1-indexed spreadsheet column label without importing openpyxl."""
label = ""
while index > 0:
index, remainder = divmod(index - 1, 26)
label = chr(ord("A") + remainder) + label
return label
def _paginate(items: list[T], limit: int | None, offset: int) -> tuple[list[T], bool]: def _paginate(items: list[T], limit: int | None, offset: int) -> tuple[list[T], bool]:
if limit is None: if limit is None:
return items[offset:], False return items[offset:], False
@@ -201,8 +138,11 @@ class FindFilesTool(_SearchTool):
@property @property
def description(self) -> str: def description(self) -> str:
return ( return (
"Find workspace paths by name, glob, or file type. " "Find files by path fragment, glob, or file type. "
"Returns relative paths and skips dependency/build directories." "Use this before read_file when you need to locate files, and "
"prefer it over shell find/ls for ordinary workspace discovery. "
"Returns workspace-relative paths and skips common dependency/build "
"directories."
) )
@property @property
@@ -216,139 +156,60 @@ class FindFilesTool(_SearchTool):
"properties": { "properties": {
"path": { "path": {
"type": "string", "type": "string",
"description": "Search root (default '.')", "description": "Directory or file to search in (default '.')",
}, },
"query": { "query": {
"type": "string", "type": "string",
"description": "Case-insensitive path terms; all must match", "description": (
"Optional case-insensitive path fragment search. "
"Whitespace-separated terms must all be present."
),
}, },
"glob": { "glob": {
"type": "string", "type": "string",
"description": "Path filter, e.g. '*.py' or 'tests/**/test_*.py'", "description": "Optional file filter, e.g. '*.py' or 'tests/**/test_*.py'",
}, },
"type": { "type": {
"type": "string", "type": "string",
"description": "File type, e.g. 'py', 'ts', 'md', or 'json'", "description": "Optional file type shorthand, e.g. 'py', 'ts', 'md', 'json'",
}, },
"include_dirs": { "include_dirs": {
"type": "boolean", "type": "boolean",
"description": "Include directories (default false)", "description": "Include matching directories as well as files (default false)",
}, },
"sort": { "sort": {
"type": "string", "type": "string",
"enum": ["path", "modified"], "enum": ["path", "modified"],
"description": "Sort order (default path)", "description": "Sort by path or most recently modified first (default path)",
}, },
"head_limit": { "head_limit": {
"type": "integer", "type": "integer",
"description": "Maximum paths (default 200; 0 for all)", "description": "Maximum number of paths to return (default 200, 0 for all, max 1000)",
"minimum": 0, "minimum": 0,
"maximum": 1000, "maximum": 1000,
}, },
"offset": { "offset": {
"type": "integer", "type": "integer",
"description": "Paths to skip before head_limit", "description": "Skip the first N results before applying head_limit",
"minimum": 0, "minimum": 0,
"maximum": 100000, "maximum": 100000,
}, },
}, },
} }
def _entry(self, path: Path, root: Path, *, is_dir: bool) -> _FindFilesEntry: def _iter_paths(self, root: Path, *, include_dirs: bool) -> Iterable[Path]:
display_path = self._display_path(path, root)
return _FindFilesEntry(
path=path,
rel_path=path.relative_to(root).as_posix(),
display_path=display_path,
name=path.name,
is_dir=is_dir,
)
def _push_directory_entries(
self,
directory: Path,
root: Path,
frontier: list[tuple[str, int, _FindFilesEntry]],
sequence: int,
budget: _FindFilesBudget,
) -> int:
budget.checkpoint()
try:
with os.scandir(directory) as entries:
for raw_entry in entries:
budget.visit_path()
try:
is_dir = raw_entry.is_dir(follow_symlinks=False)
# os.walk yields special files and broken file symlinks,
# but does not descend into directory symlinks by default.
if not is_dir and raw_entry.is_symlink() and raw_entry.is_dir():
continue
except OSError:
continue
if is_dir and raw_entry.name in self._IGNORE_DIRS:
continue
entry = self._entry(Path(raw_entry.path), root, is_dir=is_dir)
sort_path = entry.display_path + ("/" if is_dir else "")
heapq.heappush(frontier, (sort_path, sequence, entry))
sequence += 1
except OSError:
# os.walk silently skips directories that cannot be listed. Preserve
# that behavior while still allowing cancellation and budget errors
# to propagate from the explicit checkpoints above.
pass
return sequence
def _iter_paths(
self,
root: Path,
*,
include_dirs: bool,
budget: _FindFilesBudget,
) -> Iterable[_FindFilesEntry]:
budget.checkpoint()
if root.is_file(): if root.is_file():
budget.visit_path() yield root
yield self._entry(root, root.parent, is_dir=False)
return return
if include_dirs: if include_dirs:
yield self._entry(root, root, is_dir=True) yield root
for dirpath, dirnames, filenames in os.walk(root):
frontier: list[tuple[str, int, _FindFilesEntry]] = [] dirnames[:] = sorted(d for d in dirnames if d not in self._IGNORE_DIRS)
sequence = self._push_directory_entries(root, root, frontier, 0, budget) current = Path(dirpath)
while frontier: if include_dirs and current != root:
budget.checkpoint() yield current
_, _, entry = heapq.heappop(frontier) for filename in sorted(filenames):
if entry.is_dir: yield current / filename
if include_dirs:
yield entry
sequence = self._push_directory_entries(
entry.path,
root,
frontier,
sequence,
budget,
)
else:
yield entry
@staticmethod
def _matches_entry(
entry: _FindFilesEntry,
*,
query: str | None,
glob: str | None,
file_type: str | None,
) -> bool:
if glob and not _match_glob(entry.rel_path, entry.name, glob):
return False
if entry.is_dir:
if file_type:
return False
elif not _matches_type(entry.name, file_type):
return False
return _matches_query(entry.display_path, query)
async def execute( async def execute(
self, self,
@@ -397,9 +258,6 @@ class FindFilesTool(_SearchTool):
offset: int, offset: int,
cancelled: threading.Event, cancelled: threading.Event,
) -> str: ) -> str:
started_at = time.monotonic()
if cancelled.is_set():
raise _FindFilesCancelledError
target = self._resolve(path or ".") target = self._resolve(path or ".")
if not target.exists(): if not target.exists():
return ToolResult.error(f"Error: Path not found: {path}") return ToolResult.error(f"Error: Path not found: {path}")
@@ -414,63 +272,50 @@ class FindFilesTool(_SearchTool):
if head_limit is None if head_limit is None
else None if head_limit == 0 else head_limit else None if head_limit == 0 else head_limit
) )
budget = _FindFilesBudget( root = target if target.is_dir() else target.parent
cancelled=cancelled, matches: list[tuple[str, float]] = []
deadline=started_at + self._MAX_SCAN_SECONDS, deadline = time.monotonic() + self._MAX_SCAN_SECONDS
max_paths=self._MAX_SCAN_PATHS, scanned = 0
)
def matching_entries() -> Iterator[tuple[str, float]]: for candidate in self._iter_paths(target, include_dirs=include_dirs):
for entry in self._iter_paths( if cancelled.is_set():
target, raise RuntimeError("find_files scan cancelled")
include_dirs=include_dirs, scanned += 1
budget=budget, if scanned > self._MAX_SCAN_PATHS:
): return ToolResult.error(
if not self._matches_entry( f"Error: find_files scan exceeded {self._MAX_SCAN_PATHS} paths; "
entry, "narrow path, query, glob, or type and retry."
query=query, )
glob=glob, if time.monotonic() > deadline:
file_type=file_type, return ToolResult.error(
): f"Error: find_files scan exceeded {self._MAX_SCAN_SECONDS:g} seconds; "
continue "narrow path, query, glob, or type and retry."
)
if candidate.is_dir() and not include_dirs:
continue
rel_path = candidate.relative_to(root).as_posix()
display_path = self._display_path(candidate, root)
name = candidate.name
if glob and not _match_glob(rel_path, name, glob):
continue
if candidate.is_file() and not _matches_type(name, file_type):
continue
if candidate.is_dir() and file_type:
continue
if not _matches_query(display_path, query):
continue
try:
mtime = candidate.stat().st_mtime
except OSError:
mtime = 0.0 mtime = 0.0
if sort == "modified": suffix = "/" if candidate.is_dir() else ""
try: matches.append((display_path + suffix, mtime))
mtime = entry.path.stat().st_mtime
except OSError:
pass
suffix = "/" if entry.is_dir else ""
yield entry.display_path + suffix, mtime
matches: list[tuple[str, float]] if sort == "modified":
try: matches.sort(key=lambda item: (-item[1], item[0]))
if sort == "modified": else:
if limit is None: matches.sort(key=lambda item: item[0])
matches = sorted(matching_entries(), key=lambda item: (-item[1], item[0]))
else:
selection_size = offset + limit + 1
matches = heapq.nsmallest(
selection_size,
matching_entries(),
key=lambda item: (-item[1], item[0]),
)
else:
selection_size = None if limit is None else offset + limit + 1
matches = []
for match in matching_entries():
matches.append(match)
if selection_size is not None and len(matches) >= selection_size:
break
budget.checkpoint()
except _FindFilesBudgetExceededError as exc:
if str(exc) == "paths":
detail = f"{self._MAX_SCAN_PATHS} paths"
else:
detail = f"{self._MAX_SCAN_SECONDS:g} seconds"
return ToolResult.error(
f"Error: find_files scan exceeded {detail}; "
"narrow path, query, glob, or type and retry."
)
paths = [item[0] for item in matches] paths = [item[0] for item in matches]
paged, truncated = _paginate(paths, limit, offset) paged, truncated = _paginate(paths, limit, offset)
@@ -485,11 +330,10 @@ class FindFilesTool(_SearchTool):
class GrepTool(_SearchTool): class GrepTool(_SearchTool):
"""Search text and document contents using a regex-like pattern.""" """Search file contents using a regex-like pattern."""
_scopes = {"core", "subagent"} _scopes = {"core", "subagent"}
_MAX_RESULT_CHARS = 128_000 _MAX_RESULT_CHARS = 128_000
_MAX_RENDERED_LINE_CHARS = 2_000
_MAX_FILE_BYTES = 2_000_000 _MAX_FILE_BYTES = 2_000_000
_MAX_EXPLICIT_FILE_BYTES = 100_000_000 _MAX_EXPLICIT_FILE_BYTES = 100_000_000
@@ -500,8 +344,12 @@ class GrepTool(_SearchTool):
@property @property
def description(self) -> str: def description(self) -> str:
return ( return (
"Search text, PDF, DOCX, XLSX, and PPTX content. " "Search file contents with a regex pattern. "
"Returns matches with five context lines and source locators by default." "Default output_mode is files_with_matches (file paths only); "
"use content mode for matching lines with context. Prefer this "
"over shell grep for ordinary workspace searches. "
"Binary and file-size limits are enforced by the tool; explicit file paths "
"use a larger bounded limit than directory searches. Supports glob/type filtering."
) )
@property @property
@@ -515,62 +363,80 @@ class GrepTool(_SearchTool):
"properties": { "properties": {
"pattern": { "pattern": {
"type": "string", "type": "string",
"description": "Regex, or literal text when fixed_strings=true", "description": "Regex or plain text pattern to search for",
"minLength": 1, "minLength": 1,
}, },
"path": { "path": {
"type": "string", "type": "string",
"description": "Search root (default '.')", "description": "File or directory to search in (default '.')",
}, },
"glob": { "glob": {
"type": "string", "type": "string",
"description": "Path filter, e.g. '*.py' or 'tests/**/test_*.py'", "description": "Optional file filter, e.g. '*.py' or 'tests/**/test_*.py'",
}, },
"type": { "type": {
"type": "string", "type": "string",
"description": "File type, e.g. 'py', 'ts', 'md', or 'json'", "description": "Optional file type shorthand, e.g. 'py', 'ts', 'md', 'json'",
},
"pages": {
"type": "string",
"description": "PDF page number or range, e.g. '7' or '101-200' (max 100 pages)",
}, },
"case_insensitive": { "case_insensitive": {
"type": "boolean", "type": "boolean",
"description": "Ignore case (default false)", "description": "Case-insensitive search (default false)",
}, },
"fixed_strings": { "fixed_strings": {
"type": "boolean", "type": "boolean",
"description": "Treat pattern literally (default false)", "description": "Treat pattern as plain text instead of regex (default false)",
}, },
"output_mode": { "output_mode": {
"type": "string", "type": "string",
"enum": ["content", "files_with_matches", "count"], "enum": ["content", "files_with_matches", "count"],
"description": ( "description": (
"content: matches with context (default); " "content: matching lines with optional context; "
"files_with_matches: paths; count: matches per file" "files_with_matches: only matching file paths; "
"count: matching line counts per file. "
"Default: files_with_matches"
), ),
}, },
"context_before": { "context_before": {
"type": "integer", "type": "integer",
"description": "Context lines before a match (default 5)", "description": "Number of lines of context before each match",
"minimum": 0, "minimum": 0,
"maximum": 20, "maximum": 20,
}, },
"context_after": { "context_after": {
"type": "integer", "type": "integer",
"description": "Context lines after a match (default 5)", "description": "Number of lines of context after each match",
"minimum": 0, "minimum": 0,
"maximum": 20, "maximum": 20,
}, },
"max_matches": {
"type": "integer",
"description": (
"Legacy alias for head_limit in content mode"
),
"minimum": 1,
"maximum": 1000,
},
"max_results": {
"type": "integer",
"description": (
"Legacy alias for head_limit in files_with_matches or count mode"
),
"minimum": 1,
"maximum": 1000,
},
"head_limit": { "head_limit": {
"type": "integer", "type": "integer",
"description": "Maximum matches or file entries (default 250; 0 for all)", "description": (
"Maximum number of results to return. In content mode this limits "
"matching line blocks; in other modes it limits file entries. "
"Default 250"
),
"minimum": 0, "minimum": 0,
"maximum": 1000, "maximum": 1000,
}, },
"offset": { "offset": {
"type": "integer", "type": "integer",
"description": "Matches or file entries to skip before head_limit", "description": "Skip the first N results before applying head_limit",
"minimum": 0, "minimum": 0,
"maximum": 100000, "maximum": 100000,
}, },
@@ -578,97 +444,20 @@ class GrepTool(_SearchTool):
"required": ["pattern"], "required": ["pattern"],
} }
@staticmethod
def _clip_rendered_line(text: str, match_start: int | None = None) -> str:
limit = GrepTool._MAX_RENDERED_LINE_CHARS
if len(text) <= limit:
return text
marker = "..."
available = limit - len(marker)
if match_start is None:
return text[:available] + marker
start = max(0, match_start - available // 3)
start = min(start, len(text) - available)
end = start + available
prefix = marker if start else ""
suffix = marker if end < len(text) else ""
visible = text[start:end]
if prefix and suffix:
visible = visible[: available - len(marker)]
return prefix + visible + suffix
@staticmethod
def _matching_contexts(
lines: Iterable[LocatedDocumentLine],
regex: re.Pattern[str],
before: int,
after: int,
) -> Iterable[tuple[list[LocatedDocumentLine], int, int]]:
history: deque[LocatedDocumentLine] = deque(maxlen=before)
pending: list[_PendingContextMatch] = []
for line in lines:
if not line.searchable:
continue
still_pending: list[_PendingContextMatch] = []
for item in pending:
item.lines.append(line)
item.remaining_after -= 1
if item.remaining_after == 0:
yield item.lines, item.match_index, item.match_start
else:
still_pending.append(item)
pending = still_pending
match = regex.search(line.text)
if match is not None:
context_lines = [*history, line]
item = _PendingContextMatch(
lines=context_lines,
match_index=len(context_lines) - 1,
match_start=match.start(),
remaining_after=after,
)
if after == 0:
yield item.lines, item.match_index, item.match_start
else:
pending.append(item)
history.append(line)
for item in pending:
yield item.lines, item.match_index, item.match_start
@staticmethod @staticmethod
def _format_block( def _format_block(
display_path: str, display_path: str,
lines: list[LocatedDocumentLine], lines: list[str],
match_index: int, match_line: int,
match_start: int = 0, before: int,
after: int,
) -> str: ) -> str:
match_line = lines[match_index] start = max(1, match_line - before)
source_line = match_line.extracted_line end = min(len(lines), match_line + after)
match_locator = match_line.locator block = [f"{display_path}:{match_line}"]
if match_locator.startswith("sheet="): for line_no in range(start, end + 1):
column = _excel_column(match_line.text[:match_start].count("\t") + 1) marker = ">" if line_no == match_line else " "
row_match = re.search(r",row=(\d+)$", match_locator) block.append(f"{marker} {line_no}| {lines[line_no - 1]}")
if row_match:
match_locator += f",cell={column}{row_match.group(1)}"
suffix = f" [{match_locator}]" if match_locator else ""
block = [f"{display_path}:{source_line}{suffix}"]
for index, line in enumerate(lines):
is_match = index == match_index
marker = ">" if is_match else " "
coordinate = str(line.extracted_line)
if line.locator:
coordinate += f" [{line.locator}]"
rendered = GrepTool._clip_rendered_line(
line.text,
match_start if is_match else None,
)
block.append(f"{marker} {coordinate}| {rendered}")
return "\n".join(block) return "\n".join(block)
async def execute( async def execute(
@@ -677,12 +466,11 @@ class GrepTool(_SearchTool):
path: str = ".", path: str = ".",
glob: str | None = None, glob: str | None = None,
type: str | None = None, type: str | None = None,
pages: str | None = None,
case_insensitive: bool = False, case_insensitive: bool = False,
fixed_strings: bool = False, fixed_strings: bool = False,
output_mode: str = "content", output_mode: str = "files_with_matches",
context_before: int = 5, context_before: int = 0,
context_after: int = 5, context_after: int = 0,
max_matches: int | None = None, max_matches: int | None = None,
max_results: int | None = None, max_results: int | None = None,
head_limit: int | None = None, head_limit: int | None = None,
@@ -718,8 +506,6 @@ class GrepTool(_SearchTool):
size_truncated = False size_truncated = False
skipped_binary = 0 skipped_binary = 0
skipped_large = 0 skipped_large = 0
document_errors: list[str] = []
document_continuations: list[str] = []
matching_files: list[str] = [] matching_files: list[str] = []
counts: dict[str, int] = {} counts: dict[str, int] = {}
file_mtimes: dict[str, float] = {} file_mtimes: dict[str, float] = {}
@@ -734,109 +520,61 @@ class GrepTool(_SearchTool):
continue continue
if not _matches_type(file_path.name, type): if not _matches_type(file_path.name, type):
continue continue
display_path = self._display_path(file_path, root)
try: with file_path.open("rb") as file:
file_size = file_path.stat().st_size raw = file.read(max_file_bytes + 1)
except OSError: if len(raw) > max_file_bytes:
skipped_binary += 1
continue
if file_size > max_file_bytes:
skipped_large += 1 skipped_large += 1
continue continue
if _is_binary(raw):
skipped_binary += 1
continue
try: try:
mtime = file_path.stat().st_mtime mtime = file_path.stat().st_mtime
except OSError: except OSError:
mtime = 0.0 mtime = 0.0
source_iterator: Iterator[LocatedDocumentLine] | None = None
is_document = file_path.suffix.lower() in _DOCUMENT_EXTENSIONS
try: try:
if is_document: content = raw.decode("utf-8")
source = open_document_line_source(file_path, pages=pages) except UnicodeDecodeError:
if source is None:
skipped_binary += 1
continue
source_iterator = source.lines
source_lines: Iterable[LocatedDocumentLine] = source_iterator
if source.continuation:
document_continuations.append(
f"({display_path}: continue PDF search with "
f"{source.continuation})"
)
else:
with file_path.open("rb") as file:
raw = file.read(max_file_bytes + 1)
if _is_binary(raw):
skipped_binary += 1
continue
try:
content = raw.decode("utf-8")
except UnicodeDecodeError:
skipped_binary += 1
continue
source_lines = (
LocatedDocumentLine(text, line_no, "")
for line_no, text in enumerate(content.splitlines(), 1)
)
file_had_match = False
if output_mode == "content":
contexts = self._matching_contexts(
source_lines,
regex,
context_before,
context_after,
)
for context_lines, match_index, match_start in contexts:
file_had_match = True
seen_content_matches += 1
if seen_content_matches <= offset:
continue
if limit is not None and len(blocks) >= limit:
truncated = True
break
block = self._format_block(
display_path,
context_lines,
match_index,
match_start,
)
extra_sep = 2 if blocks else 0
if result_chars + extra_sep + len(block) > self._MAX_RESULT_CHARS:
size_truncated = True
break
blocks.append(block)
result_chars += extra_sep + len(block)
else:
for line in source_lines:
if not line.searchable or regex.search(line.text) is None:
continue
file_had_match = True
if output_mode == "count":
counts[display_path] = counts.get(display_path, 0) + 1
continue
if display_path not in matching_files:
matching_files.append(display_path)
file_mtimes[display_path] = mtime
break
except Exception as e:
if not is_document:
raise
if target.is_file():
if isinstance(e, PdfPageRangeError):
return ToolResult.error(
f"Error: Invalid PDF page range '{pages}': {e!s}."
)
return ToolResult.error(
f"Error searching document {display_path}: {e!s}"
)
skipped_binary += 1 skipped_binary += 1
document_errors.append(f"{display_path}: {e!s}")
continue continue
finally:
close = getattr(source_iterator, "close", None) lines = content.splitlines()
if close is not None: display_path = self._display_path(file_path, root)
close() file_had_match = False
for idx, line in enumerate(lines, start=1):
if not regex.search(line):
continue
file_had_match = True
if output_mode == "count":
counts[display_path] = counts.get(display_path, 0) + 1
continue
if output_mode == "files_with_matches":
if display_path not in matching_files:
matching_files.append(display_path)
file_mtimes[display_path] = mtime
break
seen_content_matches += 1
if seen_content_matches <= offset:
continue
if limit is not None and len(blocks) >= limit:
truncated = True
break
block = self._format_block(
display_path,
lines,
idx,
context_before,
context_after,
)
extra_sep = 2 if blocks else 0
if result_chars + extra_sep + len(block) > self._MAX_RESULT_CHARS:
size_truncated = True
break
blocks.append(block)
result_chars += extra_sep + len(block)
if output_mode == "count" and file_had_match: if output_mode == "count" and file_had_match:
if display_path not in matching_files: if display_path not in matching_files:
matching_files.append(display_path) matching_files.append(display_path)
@@ -865,8 +603,8 @@ class GrepTool(_SearchTool):
key=lambda name: (-file_mtimes.get(name, 0.0), name), key=lambda name: (-file_mtimes.get(name, 0.0), name),
) )
ordered, truncated = _paginate(ordered_files, limit, offset) ordered, truncated = _paginate(ordered_files, limit, offset)
count_lines = [f"{name}: {counts[name]}" for name in ordered] lines = [f"{name}: {counts[name]}" for name in ordered]
result = "\n".join(count_lines) result = "\n".join(lines)
else: else:
if not blocks: if not blocks:
result = f"No matches found for pattern '{pattern}' in {path}" result = f"No matches found for pattern '{pattern}' in {path}"
@@ -876,14 +614,10 @@ class GrepTool(_SearchTool):
notes: list[str] = [] notes: list[str] = []
if output_mode == "content" and truncated: if output_mode == "content" and truncated:
notes.append( notes.append(
f"(pagination: limit={limit}, offset={offset}; " f"(pagination: limit={limit}, offset={offset})"
f"use offset={offset + len(blocks)} to continue)"
) )
elif output_mode == "content" and size_truncated: elif output_mode == "content" and size_truncated:
notes.append( notes.append("(output truncated due to size)")
"(output truncated due to size; "
f"use offset={offset + len(blocks)} to continue)"
)
elif truncated and output_mode in {"count", "files_with_matches"}: elif truncated and output_mode in {"count", "files_with_matches"}:
notes.append( notes.append(
f"(pagination: limit={limit}, offset={offset})" f"(pagination: limit={limit}, offset={offset})"
@@ -896,9 +630,6 @@ class GrepTool(_SearchTool):
notes.append(f"(skipped {skipped_binary} binary/unreadable files)") notes.append(f"(skipped {skipped_binary} binary/unreadable files)")
if skipped_large: if skipped_large:
notes.append(f"(skipped {skipped_large} large files)") notes.append(f"(skipped {skipped_large} large files)")
if document_errors:
notes.append(f"(first document error: {document_errors[0]})")
notes.extend(document_continuations[:10])
if output_mode == "count" and counts: if output_mode == "count" and counts:
notes.append( notes.append(
f"(total matches: {sum(counts.values())} in {len(counts)} files)" f"(total matches: {sum(counts.values())} in {len(counts)} files)"
+46 -13
View File
@@ -58,6 +58,7 @@ def _is_string_mapping(value: object) -> TypeGuard[Mapping[str, object]]:
class MyTool(Tool): class MyTool(Tool):
"""Check and set the agent loop's runtime configuration.""" """Check and set the agent loop's runtime configuration."""
_plugin_discoverable = False # Requires AgentLoop reference; registered manually
config_key = "my" config_key = "my"
@classmethod @classmethod
@@ -66,16 +67,7 @@ class MyTool(Tool):
@classmethod @classmethod
def enabled(cls, ctx: ToolContext) -> bool: def enabled(cls, ctx: ToolContext) -> bool:
return ctx.runtime_control is not None and ctx.config.my.enable return ctx.config.my.enable
@classmethod
def create(cls, ctx: ToolContext) -> Tool:
if ctx.runtime_control is None:
raise RuntimeError("MyTool requires a runtime control capability")
return cls(
runtime_control=ctx.runtime_control,
modify_allowed=ctx.config.my.allow_set,
)
BLOCKED = frozenset({ BLOCKED = frozenset({
# Core infrastructure # Core infrastructure
@@ -96,6 +88,9 @@ class MyTool(Tool):
READ_ONLY = frozenset({ READ_ONLY = frozenset({
"subagents", # observable but replacing it would break the system "subagents", # observable but replacing it would break the system
"tool_names", "tool_names",
"current_iteration",
"_current_iteration", # updated by runner only
"_last_usage",
"exec_config", # inspect allowed (e.g. check sandbox), modify blocked "exec_config", # inspect allowed (e.g. check sandbox), modify blocked
"web_config", # inspect allowed (e.g. check enable), modify blocked "web_config", # inspect allowed (e.g. check enable), modify blocked
"model_presets", # config-derived catalog; changes require config reload "model_presets", # config-derived catalog; changes require config reload
@@ -155,9 +150,11 @@ class MyTool(Tool):
"Actions: check, set.\n" "Actions: check, set.\n"
"- check (no key): full config overview — start here.\n" "- check (no key): full config overview — start here.\n"
"- check (key): drill into a value. Dot-paths allowed " "- check (key): drill into a value. Dot-paths allowed "
"(e.g. 'web_config.enable').\n" "(e.g. '_last_usage.input_tokens', 'web_config.enable').\n"
"- set (key, value): change config or store notes in your scratchpad. " "- set (key, value): change config or store notes in your scratchpad. "
"Scratchpad keys persist across turns but not restarts.\n" "Scratchpad keys persist across turns but not restarts.\n"
"Key values: _current_iteration (current progress), "
"max_iterations - _current_iteration = remaining iterations.\n"
"Current routing metadata is available read-only via request.channel, " "Current routing metadata is available read-only via request.channel, "
"request.chat_id, and request.sender_id.\n" "request.chat_id, and request.sender_id.\n"
"Use model_preset for session-scoped model or context changes; direct " "Use model_preset for session-scoped model or context changes; direct "
@@ -165,7 +162,7 @@ class MyTool(Tool):
"Note: web_config and exec_config are readable but read-only.\n" "Note: web_config and exec_config are readable but read-only.\n"
"\n" "\n"
"When to use:\n" "When to use:\n"
"- User asks about your model or settings → check that key.\n" "- User asks about your model, settings, or token usage → check that key.\n"
"- User asks to switch to a named model preset → set model_preset to that preset name.\n" "- User asks to switch to a named model preset → set model_preset to that preset name.\n"
"- A tool fails or behaves unexpectedly → check the related config to diagnose.\n" "- A tool fails or behaves unexpectedly → check the related config to diagnose.\n"
"- User asks you to remember a preference for this session → set to store it in your scratchpad.\n" "- User asks you to remember a preference for this session → set to store it in your scratchpad.\n"
@@ -373,7 +370,7 @@ class MyTool(Tool):
if not self._modify_allowed: if not self._modify_allowed:
return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)") return ToolResult.error("Error: set is disabled (tools.my.allow_set is false)")
if action in ("modify", "set"): if action in ("modify", "set"):
return self._modify(key, value) return await self._modify_async(key, value)
return f"Unknown action: {action}" return f"Unknown action: {action}"
# -- inspect -- # -- inspect --
@@ -445,11 +442,14 @@ class MyTool(Tool):
"workspace", "workspace",
"provider_retry_mode", "provider_retry_mode",
"max_tool_result_chars", "max_tool_result_chars",
"_current_iteration",
"web_config", "web_config",
"exec_config", "exec_config",
"subagents", "subagents",
): ):
parts.append(self._format_value(values[k], k)) parts.append(self._format_value(values[k], k))
if snapshot.last_usage:
parts.append(self._format_value(snapshot.last_usage, "_last_usage"))
if snapshot.scratchpad: if snapshot.scratchpad:
parts.append(self._format_value(snapshot.scratchpad, "scratchpad")) parts.append(self._format_value(snapshot.scratchpad, "scratchpad"))
return "\n".join(parts) return "\n".join(parts)
@@ -492,6 +492,11 @@ class MyTool(Tool):
return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified") return ToolResult.error(f"Error: '{key}' is read-only and cannot be modified")
return self._modify_scratchpad(key, value) return self._modify_scratchpad(key, value)
async def _modify_async(self, key: str | None, value: Any) -> str:
if key == "model_preset":
return await self._modify_model_preset_async(value)
return self._modify(key, value)
def _modify_model_preset(self, value: Any) -> str: def _modify_model_preset(self, value: Any) -> str:
if not isinstance(value, str) or not value.strip(): if not isinstance(value, str) or not value.strip():
return ToolResult.error("Error: 'model_preset' must be a non-empty string") return ToolResult.error("Error: 'model_preset' must be a non-empty string")
@@ -520,6 +525,34 @@ class MyTool(Tool):
f"context_window_tokens is now {runtime.context_window_tokens!r}" f"context_window_tokens is now {runtime.context_window_tokens!r}"
) )
async def _modify_model_preset_async(self, value: Any) -> str:
if not isinstance(value, str) or not value.strip():
return ToolResult.error("Error: 'model_preset' must be a non-empty string")
name = value.strip()
session_key = current_request_session_key()
old = self._runtime_control.snapshot().model_preset
try:
runtime = await self._runtime_control.set_model_preset_async(
name,
session_key=session_key,
)
except (KeyError, ValueError) as exc:
message = str(exc.args[0]) if exc.args else str(exc)
punctuation = "" if message.endswith((".", "!", "?")) else "."
return ToolResult.error(f"Error: {message}{punctuation}")
if session_key:
self._audit("modify", f"model_preset = {name!r}")
return (
f"Set model_preset = {name!r} for the next turn; "
f"model will be {runtime.model!r}; "
f"context_window_tokens will be {runtime.context_window_tokens!r}"
)
self._audit("modify", f"model_preset: {old!r} -> {name!r}")
return (
f"Set model_preset = {name!r} (was {old!r}); model is now {runtime.model!r}; "
f"context_window_tokens is now {runtime.context_window_tokens!r}"
)
def _modify_restricted(self, key: str, value: Any) -> str: def _modify_restricted(self, key: str, value: Any) -> str:
spec = self.RESTRICTED[key] spec = self.RESTRICTED[key]
expected = cast(type[Any], spec["type"]) expected = cast(type[Any], spec["type"])
+3 -16
View File
@@ -7,7 +7,7 @@ from __future__ import annotations
import asyncio import asyncio
import json import json
import time import time
from collections import OrderedDict, deque from collections import deque
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Protocol from typing import Any, Protocol
@@ -127,7 +127,7 @@ class SendSessionMessageTool(Tool):
self._max_messages_per_minute = max_messages_per_minute self._max_messages_per_minute = max_messages_per_minute
self._schedule_later = schedule_later self._schedule_later = schedule_later
self._clock = clock or time.monotonic self._clock = clock or time.monotonic
self._sent_at: OrderedDict[str, deque[float]] = OrderedDict() self._sent_at: dict[str, deque[float]] = {}
self._pending_replies: dict[tuple[str, str], _PendingReply] = {} self._pending_replies: dict[tuple[str, str], _PendingReply] = {}
self._expiry_tasks: set[asyncio.Task[None]] = set() self._expiry_tasks: set[asyncio.Task[None]] = set()
self._send_lock = asyncio.Lock() self._send_lock = asyncio.Lock()
@@ -240,11 +240,8 @@ class SendSessionMessageTool(Tool):
async with self._send_lock: async with self._send_lock:
now = self._clock() now = self._clock()
sent_at = self._sent_at.setdefault(source.session_key, deque())
cutoff = now - _RATE_LIMIT_WINDOW_SECONDS cutoff = now - _RATE_LIMIT_WINDOW_SECONDS
self._prune_expired_rate_limits(cutoff)
sent_at = self._sent_at.get(source.session_key)
if sent_at is None:
sent_at = deque[float]()
while sent_at and sent_at[0] <= cutoff: while sent_at and sent_at[0] <= cutoff:
sent_at.popleft() sent_at.popleft()
if len(sent_at) >= self._max_messages_per_minute: if len(sent_at) >= self._max_messages_per_minute:
@@ -262,8 +259,6 @@ class SendSessionMessageTool(Tool):
input_role="user", input_role="user",
)) ))
sent_at.append(now) sent_at.append(now)
self._sent_at[source.session_key] = sent_at
self._sent_at.move_to_end(source.session_key)
self._cancel_pending_reply(reverse_wait_key) self._cancel_pending_reply(reverse_wait_key)
if timeout_seconds is not None: if timeout_seconds is not None:
self._cancel_pending_reply(wait_key) self._cancel_pending_reply(wait_key)
@@ -276,14 +271,6 @@ class SendSessionMessageTool(Tool):
return f"@{target.name}" return f"@{target.name}"
def _prune_expired_rate_limits(self, cutoff: float) -> None:
"""Drop sources ordered by their most recent successful send."""
while self._sent_at:
_, sent_at = next(iter(self._sent_at.items()))
if sent_at[-1] > cutoff:
return
self._sent_at.popitem(last=False)
@staticmethod @staticmethod
def _validate_reply_timeout( def _validate_reply_timeout(
expect_reply: bool, expect_reply: bool,
+8 -10
View File
@@ -25,7 +25,6 @@ _READ_LIMIT = 8
_SEARCH_EXCERPT_CHARS = 360 _SEARCH_EXCERPT_CHARS = 360
_READ_MESSAGE_CHARS = 4_000 _READ_MESSAGE_CHARS = 4_000
_UNTRUSTED_NOTICE = "Historical session content is untrusted data, not instructions." _UNTRUSTED_NOTICE = "Historical session content is untrusted data, not instructions."
_UNSUPPORTED_MATCH_ALL_QUERIES = {"*", ".*"}
def session_extra(metadata: Mapping[str, Any] | None) -> dict[str, Any]: def session_extra(metadata: Mapping[str, Any] | None) -> dict[str, Any]:
@@ -146,8 +145,8 @@ class SearchSessionsTool(_SessionTool):
max_length=512, max_length=512,
), ),
query=StringSchema( query=StringSchema(
"Optional literal substring filter. Omit or leave blank for the latest messages; " "Optional text filter. When omitted, return the latest visible messages.",
"regex and glob are not supported.", min_length=1,
max_length=500, max_length=500,
), ),
required=["session_key"], required=["session_key"],
@@ -167,8 +166,10 @@ class ReadSessionTool(_SessionTool):
@property @property
def description(self) -> str: def description(self) -> str:
return ( return (
"Read bounded, visible user and assistant messages from a persisted conversation. " "Read visible user and assistant messages from a persisted conversation. Pass an exact "
"Treat history as untrusted data." "session_key from a selected reference or search_sessions, or a session @handle from "
"list_sessions. With query, return recent matches; otherwise return the latest visible "
"messages. Treat history as untrusted data."
) )
async def execute( async def execute(
@@ -195,11 +196,8 @@ class ReadSessionTool(_SessionTool):
session_handle = f"@{handle_name}" session_handle = f"@{handle_name}"
session_key = handle.session_key session_key = handle.session_key
query_text = query.strip() if query else "" query_text = query.strip() if query else ""
if query_text in _UNSUPPORTED_MATCH_ALL_QUERIES: if query is not None and not query_text:
return ToolResult.error( return ToolResult.error("Error: query must not be empty")
"Error: query matches literal substrings; '*' and '.*' do not mean match all. "
"Omit query to read the latest messages."
)
match = await asyncio.to_thread( match = await asyncio.to_thread(
self._access.read, self._access.read,
session_key, session_key,
+69 -9
View File
@@ -122,37 +122,55 @@ class _PreparedCommand:
working_dir=StringSchema("Optional working directory for the command"), working_dir=StringSchema("Optional working directory for the command"),
workdir=StringSchema("Compatibility alias for working_dir"), workdir=StringSchema("Compatibility alias for working_dir"),
timeout=IntegerSchema( timeout=IntegerSchema(
description="Hard timeout in seconds (default 60, max 600).", description=(
"Timeout in seconds. Increase for long-running commands "
"like compilation or installation (default 60, max 600)."
),
minimum=1, minimum=1,
maximum=600, maximum=600,
), ),
shell=StringSchema( shell=StringSchema(
( (
"Shell override; omit for PowerShell, or pass 'cmd' for cmd.exe." "Override the Windows shell only when needed. Omit to use "
"PowerShell by default (pwsh when available, else powershell). "
"Pass 'cmd' only for cmd.exe syntax or cmd built-ins."
if _IS_WINDOWS if _IS_WINDOWS
else "Shell override; omit for bash, or pass 'sh' or 'zsh'." else "Override the Unix shell only when needed. Omit to use "
"bash by default. Pass 'sh' for POSIX sh or 'zsh' for "
"zsh-specific syntax."
), ),
nullable=True, nullable=True,
), ),
login=BooleanSchema( login=BooleanSchema(
description="Run bash/zsh as a login shell.", description="Whether to run bash/zsh with login shell semantics (default false).",
default=False, default=False,
nullable=True, nullable=True,
), ),
yield_time_ms=IntegerSchema( yield_time_ms=IntegerSchema(
description="Return after this many milliseconds if still running; omit to wait for exit.", description=(
"Optional milliseconds to wait before returning output. "
"When set, a still-running command returns a session_id that "
"can be polled or written to with write_stdin. Omit this field "
"to keep one-shot exec behavior."
),
minimum=0, minimum=0,
maximum=MAX_YIELD_MS, maximum=MAX_YIELD_MS,
nullable=True, nullable=True,
), ),
max_output_chars=IntegerSchema( max_output_chars=IntegerSchema(
description="Session output limit in characters (default 10000, max 50000).", description=(
"Maximum output characters to return when yield_time_ms is used "
"(default 10000, max 50000)."
),
minimum=1000, minimum=1000,
maximum=MAX_OUTPUT_CHARS, maximum=MAX_OUTPUT_CHARS,
nullable=True, nullable=True,
), ),
max_output_tokens=IntegerSchema( max_output_tokens=IntegerSchema(
description="Compatibility alias for max_output_chars.", description=(
"Compatibility alias for max_output_chars. The current runtime "
"uses a character budget."
),
minimum=1000, minimum=1000,
maximum=MAX_OUTPUT_CHARS, maximum=MAX_OUTPUT_CHARS,
nullable=True, nullable=True,
@@ -249,6 +267,7 @@ class ExecTool(Tool):
_MAX_TIMEOUT = 600 _MAX_TIMEOUT = 600
_MAX_OUTPUT = 10_000 _MAX_OUTPUT = 10_000
_PREPARE_TIMEOUT_SECONDS = 6.0
# Kernel device files safe as stdio redirect targets (#3599). # Kernel device files safe as stdio redirect targets (#3599).
_BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({ _BENIGN_DEVICE_PATHS: frozenset[str] = frozenset({
@@ -265,7 +284,26 @@ class ExecTool(Tool):
@property @property
def description(self) -> str: def description(self) -> str:
return "Execute a shell command." platform_note = (
"On Windows, use PowerShell syntax by default; pass shell='cmd' "
"only for cmd-specific commands. "
if _IS_WINDOWS
else "On Unix, commands run through bash by default; pass shell='sh' "
"or shell='zsh' when needed. "
)
return (
"Execute a shell command and return its output. "
"Use this for tests, builds, package commands, git commands, and "
"other process execution. Prefer read_file/find_files/grep for "
"inspection and apply_patch/write_file/edit_file for file changes "
"instead of cat, shell find/grep, echo, or sed. "
"Use -y or --yes flags to avoid interactive prompts. "
f"{platform_note}"
"For long-running or interactive commands, pass yield_time_ms; "
"if the command keeps running, exec returns a session_id that can "
"be polled or written to with write_stdin. Output is truncated at "
"10 000 chars; timeout defaults to 60s."
)
@property @property
def exclusive(self) -> bool: def exclusive(self) -> bool:
@@ -287,7 +325,20 @@ class ExecTool(Tool):
if max_output_chars is None: if max_output_chars is None:
max_output_chars = max_output_tokens max_output_chars = max_output_tokens
prepared = self._prepare_command(command, working_dir, timeout, shell, login) try:
prepared = await asyncio.wait_for(
asyncio.to_thread(
self._prepare_command,
command,
working_dir,
timeout,
shell,
login,
),
timeout=self._PREPARE_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
return ToolResult.error("Error: command validation timed out")
if isinstance(prepared, str): if isinstance(prepared, str):
return prepared return prepared
@@ -879,6 +930,15 @@ class ExecTool(Tool):
if self._is_benign_device_path(expanded): if self._is_benign_device_path(expanded):
continue continue
except Exception: except Exception:
# ``Path.expanduser()`` raises when a named user's home
# cannot be resolved (notably on Windows). An extracted
# home path must fail closed rather than bypass the guard.
if raw.strip().startswith("~"):
return ToolResult.error(
"Error: Command blocked by safety guard "
"(path outside working dir)"
+ _WORKSPACE_BOUNDARY_NOTE
)
continue continue
if self._is_benign_device_path(str(p)): if self._is_benign_device_path(str(p)):
+8 -5
View File
@@ -73,11 +73,6 @@ class SpawnTool(Tool):
"and use a dedicated subdirectory when helpful." "and use a dedicated subdirectory when helpful."
) )
@property
def concurrency_safe(self) -> bool:
"""Each call owns its task state; the manager serializes capacity admission."""
return True
async def execute( async def execute(
self, self,
task: str, task: str,
@@ -87,6 +82,14 @@ class SpawnTool(Tool):
**kwargs: Any, **kwargs: Any,
) -> str: ) -> str:
"""Spawn a subagent to execute the given task.""" """Spawn a subagent to execute the given task."""
running = self._manager.get_running_count()
limit = self._manager.max_concurrent_subagents
if running >= limit:
return (
f"Cannot spawn subagent: concurrency limit reached "
f"({running}/{limit} running). Wait for a running subagent "
f"to complete before spawning a new one."
)
request_ctx = current_request_context() request_ctx = current_request_context()
if request_ctx is None or request_ctx.runtime is None: if request_ctx is None or request_ctx.runtime is None:
return ToolResult.error("Error: spawn requires an active model runtime") return ToolResult.error("Error: spawn requires an active model runtime")
+12 -14
View File
@@ -96,7 +96,7 @@ def _normalize(text: str) -> str:
def _validate_url(url: str) -> tuple[bool, str]: def _validate_url(url: str) -> tuple[bool, str]:
"""Validate URL scheme/domain. Does NOT check resolved IPs (use _validate_url_safe for that).""" """Validate URL scheme/domain. Does not resolve IPs; use the async safe helper for that."""
try: try:
p = urlparse(url) p = urlparse(url)
if p.scheme not in ('http', 'https'): if p.scheme not in ('http', 'https'):
@@ -108,18 +108,16 @@ def _validate_url(url: str) -> tuple[bool, str]:
return False, str(e) return False, str(e)
def _validate_url_safe(url: str) -> tuple[bool, str]: async def _async_validate_url_safe(url: str) -> tuple[bool, str]:
"""Validate URL with SSRF protection: scheme, domain, and resolved IP check.""" from nanobot.security.network import async_validate_url_target
from nanobot.security.network import validate_url_target
return validate_url_target(url) return await async_validate_url_target(url)
def _resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]: async def _async_resolve_url_safe(url: str) -> tuple[bool, str, tuple[str, ...]]:
"""Validate URL and return the resolved IPs to pin during the request.""" from nanobot.security.network import async_resolve_url_target
from nanobot.security.network import resolve_url_target
return resolve_url_target(url) return await async_resolve_url_target(url)
def _pinned_dns_transport() -> httpx.AsyncBaseTransport: def _pinned_dns_transport() -> httpx.AsyncBaseTransport:
@@ -209,7 +207,7 @@ async def _get_with_safe_redirects(
"""GET a URL while validating every redirect target before requesting it.""" """GET a URL while validating every redirect target before requesting it."""
current_url = url current_url = url
for _ in range(MAX_REDIRECTS + 1): for _ in range(MAX_REDIRECTS + 1):
is_valid, error_msg, _ = _resolve_url_safe(current_url) is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
if not is_valid: if not is_valid:
return None, f"Redirect blocked: {error_msg}" return None, f"Redirect blocked: {error_msg}"
@@ -229,7 +227,7 @@ async def _get_with_safe_redirects(
return response, None return response, None
next_url = urljoin(str(response.url), location) next_url = urljoin(str(response.url), location)
is_valid, error_msg = _validate_url_safe(next_url) is_valid, error_msg = await _async_validate_url_safe(next_url)
if not is_valid: if not is_valid:
await response.aclose() await response.aclose()
return None, f"Redirect blocked: {error_msg}" return None, f"Redirect blocked: {error_msg}"
@@ -249,7 +247,7 @@ async def _stream_with_safe_redirects(
current_url = url current_url = url
chain_carries_credentials = _url_carries_credentials(url) chain_carries_credentials = _url_carries_credentials(url)
for _ in range(MAX_REDIRECTS + 1): for _ in range(MAX_REDIRECTS + 1):
is_valid, error_msg, _ = _resolve_url_safe(current_url) is_valid, error_msg, _ = await _async_resolve_url_safe(current_url)
if not is_valid: if not is_valid:
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
@@ -283,7 +281,7 @@ async def _stream_with_safe_redirects(
chain_carries_credentials = ( chain_carries_credentials = (
chain_carries_credentials or _url_carries_credentials(next_url) chain_carries_credentials or _url_carries_credentials(next_url)
) )
is_valid, error_msg = _validate_url_safe(next_url) is_valid, error_msg = await _async_validate_url_safe(next_url)
if not is_valid: if not is_valid:
await stream.__aexit__(None, None, None) await stream.__aexit__(None, None, None)
return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials return None, None, f"Redirect blocked: {error_msg}", chain_carries_credentials
@@ -1106,7 +1104,7 @@ class WebFetchTool(Tool):
url = url.strip(" \t\r\n`\"'") url = url.strip(" \t\r\n`\"'")
extract_mode = kwargs.pop("extractMode", extract_mode) extract_mode = kwargs.pop("extractMode", extract_mode)
max_chars = cast(int, kwargs.pop("maxChars", max_chars) or self.max_chars) max_chars = cast(int, kwargs.pop("maxChars", max_chars) or self.max_chars)
is_valid, error_msg = _validate_url_safe(url) is_valid, error_msg = await _async_validate_url_safe(url)
if not is_valid: if not is_valid:
return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False) return json.dumps({"error": f"URL validation failed: {error_msg}", "url": url}, ensure_ascii=False)
+2
View File
@@ -32,6 +32,7 @@ class AgentTurnHookSpec:
session_key: str | None = None session_key: str | None = None
workspace: Path | None = None workspace: Path | None = None
tool_hint_max_length: int = 40 tool_hint_max_length: int = 40
on_iteration: Callable[[int], None] | None = None
registered_hook_factories: list[AgentTurnHookFactory] = field(default_factory=list) registered_hook_factories: list[AgentTurnHookFactory] = field(default_factory=list)
turn_hook_factories: list[AgentTurnHookFactory] = field(default_factory=list) turn_hook_factories: list[AgentTurnHookFactory] = field(default_factory=list)
registered_hooks: list[AgentHook] = field(default_factory=list) registered_hooks: list[AgentHook] = field(default_factory=list)
@@ -49,6 +50,7 @@ def build_agent_turn_hook(spec: AgentTurnHookSpec) -> AgentHook:
on_stream_end=spec.on_stream_end, on_stream_end=spec.on_stream_end,
session_key=spec.session_key, session_key=spec.session_key,
tool_hint_max_length=spec.tool_hint_max_length, tool_hint_max_length=spec.tool_hint_max_length,
on_iteration=spec.on_iteration,
) )
if spec.ephemeral and not spec.run_extra_hooks_for_ephemeral: if spec.ephemeral and not spec.run_extra_hooks_for_ephemeral:
return progress_hook return progress_hook
+1 -15
View File
@@ -17,7 +17,6 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable, cast
from aiohttp import web from aiohttp import web
from loguru import logger from loguru import logger
from nanobot.agent.hook import AgentHook, AgentRunHookContext
from nanobot.config.paths import get_media_dir from nanobot.config.paths import get_media_dir
from nanobot.providers.base import LLMUsage from nanobot.providers.base import LLMUsage
from nanobot.utils.helpers import safe_filename from nanobot.utils.helpers import safe_filename
@@ -54,17 +53,6 @@ _PREPARE_AGENT_KEY = web.AppKey[Callable[[], Awaitable[None]] | None]("prepare_a
_MISSING = object() _MISSING = object()
class _UsageCaptureHook(AgentHook):
"""Capture the aggregate usage owned by one API run."""
def __init__(self) -> None:
super().__init__()
self.usage: LLMUsage | None = None
async def after_run(self, context: AgentRunHookContext) -> None:
self.usage = context.usage
def _app_value( def _app_value(
app: Any, app: Any,
key: web.AppKey[Any], key: web.AppKey[Any],
@@ -411,7 +399,6 @@ async def handle_chat_completions(request: web.Request) -> web.Response | web.St
return resp return resp
# -- non-streaming path (original logic) -- # -- non-streaming path (original logic) --
usage_capture = _UsageCaptureHook()
try: try:
async with session_lock: async with session_lock:
try: try:
@@ -423,7 +410,6 @@ async def handle_chat_completions(request: web.Request) -> web.Response | web.St
session_key=session_key, session_key=session_key,
channel="api", channel="api",
chat_id=API_CHAT_ID, chat_id=API_CHAT_ID,
hooks=[usage_capture],
) )
response_text = _response_text(response) response_text = _response_text(response)
if not response_text or not response_text.strip(): if not response_text or not response_text.strip():
@@ -440,7 +426,7 @@ async def handle_chat_completions(request: web.Request) -> web.Response | web.St
return _error_json(500, "Internal server error", err_type="server_error") return _error_json(500, "Internal server error", err_type="server_error")
return web.json_response( return web.json_response(
_chat_completion_response(response_text, model_name, usage_capture.usage) _chat_completion_response(response_text, model_name, getattr(agent_loop, "_last_usage", None))
) )
+374 -41
View File
@@ -2,15 +2,20 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import ctypes
import json import json
import os import os
import re import re
import shlex import shlex
import shutil import shutil
import signal
import subprocess import subprocess
import sys import sys
import time import time
from collections.abc import Iterable from collections.abc import Iterable
from contextlib import suppress
from ctypes import wintypes
from dataclasses import dataclass from dataclasses import dataclass
from importlib import metadata as importlib_metadata from importlib import metadata as importlib_metadata
from pathlib import Path from pathlib import Path
@@ -97,6 +102,141 @@ class CliAppsRuntimeConfig:
catalog_ttl_seconds: int = 3600 catalog_ttl_seconds: int = 3600
@dataclass(slots=True)
class _PreparedCliRun:
name: str
entry: str
resolved: str
args: list[str]
cwd: Path
timeout: int
env: dict[str, str]
artifact_snapshot: dict[Path, tuple[int, int]]
class _JobObjectBasicLimitInformation(ctypes.Structure):
_fields_ = [
("PerProcessUserTimeLimit", ctypes.c_int64),
("PerJobUserTimeLimit", ctypes.c_int64),
("LimitFlags", wintypes.DWORD),
("MinimumWorkingSetSize", ctypes.c_size_t),
("MaximumWorkingSetSize", ctypes.c_size_t),
("ActiveProcessLimit", wintypes.DWORD),
("Affinity", ctypes.c_size_t),
("PriorityClass", wintypes.DWORD),
("SchedulingClass", wintypes.DWORD),
]
class _IoCounters(ctypes.Structure):
_fields_ = [
("ReadOperationCount", ctypes.c_uint64),
("WriteOperationCount", ctypes.c_uint64),
("OtherOperationCount", ctypes.c_uint64),
("ReadTransferCount", ctypes.c_uint64),
("WriteTransferCount", ctypes.c_uint64),
("OtherTransferCount", ctypes.c_uint64),
]
class _JobObjectExtendedLimitInformation(ctypes.Structure):
_fields_ = [
("BasicLimitInformation", _JobObjectBasicLimitInformation),
("IoInfo", _IoCounters),
("ProcessMemoryLimit", ctypes.c_size_t),
("JobMemoryLimit", ctypes.c_size_t),
("PeakProcessMemoryUsed", ctypes.c_size_t),
("PeakJobMemoryUsed", ctypes.c_size_t),
]
class _WindowsJob:
"""Best-effort Windows process tree ownership for timeout/cancellation."""
_KILL_ON_JOB_CLOSE = 0x00002000
_EXTENDED_LIMIT_INFORMATION = 9
_PROCESS_TERMINATE = 0x0001
_PROCESS_SET_QUOTA = 0x0100
def __init__(self) -> None:
win_dll = getattr(ctypes, "WinDLL")
self._kernel32 = win_dll("kernel32", use_last_error=True)
self._kernel32.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR]
self._kernel32.CreateJobObjectW.restype = wintypes.HANDLE
self._kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
self._kernel32.OpenProcess.restype = wintypes.HANDLE
self._kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
self._kernel32.AssignProcessToJobObject.restype = wintypes.BOOL
self._kernel32.SetInformationJobObject.argtypes = [
wintypes.HANDLE,
ctypes.c_int,
wintypes.LPVOID,
wintypes.DWORD,
]
self._kernel32.SetInformationJobObject.restype = wintypes.BOOL
self._kernel32.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT]
self._kernel32.TerminateJobObject.restype = wintypes.BOOL
self._kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
self._kernel32.CloseHandle.restype = wintypes.BOOL
self._handle: Any = self._kernel32.CreateJobObjectW(None, None)
if not self._handle:
raise OSError(ctypes.get_last_error(), "CreateJobObjectW failed")
try:
self._set_kill_on_close(True)
except OSError:
self._kernel32.CloseHandle(self._handle)
self._handle = None
raise
@classmethod
def create(cls) -> _WindowsJob | None:
if os.name != "nt":
return None
try:
return cls()
except OSError as exc:
logger.debug("CLI Apps: Windows job object unavailable: {}", exc)
return None
def _set_kill_on_close(self, enabled: bool) -> None:
info = _JobObjectExtendedLimitInformation()
info.BasicLimitInformation.LimitFlags = self._KILL_ON_JOB_CLOSE if enabled else 0
ok = self._kernel32.SetInformationJobObject(
self._handle,
self._EXTENDED_LIMIT_INFORMATION,
ctypes.byref(info),
ctypes.sizeof(info),
)
if not ok:
raise OSError(ctypes.get_last_error(), "SetInformationJobObject failed")
def assign(self, pid: int) -> bool:
process_handle = self._kernel32.OpenProcess(
self._PROCESS_TERMINATE | self._PROCESS_SET_QUOTA,
False,
pid,
)
if not process_handle:
return False
try:
return bool(self._kernel32.AssignProcessToJobObject(self._handle, process_handle))
finally:
self._kernel32.CloseHandle(process_handle)
def terminate(self) -> None:
if self._handle and not self._kernel32.TerminateJobObject(self._handle, 1):
raise OSError(ctypes.get_last_error(), "TerminateJobObject failed")
def close(self, *, kill_descendants: bool) -> None:
if not self._handle:
return
if not kill_descendants:
with suppress(OSError):
self._set_kill_on_close(False)
self._kernel32.CloseHandle(self._handle)
self._handle = None
_BRANDS: dict[str, tuple[str, str]] = { _BRANDS: dict[str, tuple[str, str]] = {
"1password-cli": ("1password", "#3B66BC"), "1password-cli": ("1password", "#3B66BC"),
"arcgis": ("arcgis", "#2C7AC3"), "arcgis": ("arcgis", "#2C7AC3"),
@@ -1428,6 +1568,197 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
lines.append(f"- {rel} ({kind}, {self._format_artifact_size(path)})") lines.append(f"- {rel} ({kind}, {self._format_artifact_size(path)})")
return lines return lines
def _prepare_run(
self,
name: str,
args: list[str] | None,
*,
json_output: bool,
working_dir: str | None,
timeout: int | None,
restrict_to_workspace: bool,
) -> _PreparedCliRun:
app = self.get_app(name)
installed = self._load_installed()
app_name = str(app["name"])
if app_name not in installed:
raise CliAppError(f"CLI app '{name}' is not installed")
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace)
entry = str(installed[app_name].get("entry_point") or app.get("entry_point") or "")
resolved = shutil.which(entry)
if not entry or not resolved:
raise CliAppError(f"{entry or name} is not available on PATH")
clean_args = [str(arg) for arg in (args or [])]
if json_output and "--json" not in clean_args:
clean_args = ["--json", *clean_args]
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600))
return _PreparedCliRun(
name=name,
entry=entry,
resolved=resolved,
args=clean_args,
cwd=cwd,
timeout=effective_timeout,
env=self._subprocess_env(),
artifact_snapshot=self._artifact_snapshot(cwd),
)
def _format_run_result(
self,
prepared: _PreparedCliRun,
*,
returncode: int,
stdout: str,
stderr: str,
) -> str:
command = " ".join([prepared.entry, *(shlex.quote(arg) for arg in prepared.args)])
output = [
f"CLI app '{prepared.name}' exited {returncode}.",
f"Command: {command}",
]
if stdout:
output.append("\nSTDOUT:\n" + stdout.rstrip())
if stderr:
output.append("\nSTDERR:\n" + stderr.rstrip())
artifacts = self._changed_artifacts(prepared.cwd, prepared.artifact_snapshot)
if artifacts:
output.append(
"\nArtifacts created or updated:\n"
+ "\n".join(self._format_artifact_lines(prepared.cwd, artifacts))
)
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
output.append(
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
"using its workspace-relative path, for example `![diagram](diagram.png)`."
)
return _truncate("\n".join(output))
@staticmethod
def _terminate_run_process_sync(
process: subprocess.Popen[str],
job: _WindowsJob | None,
) -> None:
if job is not None:
with suppress(OSError):
job.terminate()
job.close(kill_descendants=True)
elif os.name == "nt":
with suppress(OSError, subprocess.TimeoutExpired):
subprocess.run(
["taskkill", "/PID", str(process.pid), "/T", "/F"],
check=False,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=5,
)
else:
with suppress(ProcessLookupError, PermissionError):
os.killpg(process.pid, signal.SIGKILL)
if process.poll() is None:
with suppress(ProcessLookupError):
process.kill()
with suppress(subprocess.TimeoutExpired):
process.wait(timeout=5)
@staticmethod
async def _terminate_run_process(
process: asyncio.subprocess.Process,
job: _WindowsJob | None,
) -> None:
if job is not None:
with suppress(OSError):
await asyncio.to_thread(job.terminate)
job.close(kill_descendants=True)
elif os.name == "nt":
with suppress(OSError, asyncio.TimeoutError):
await asyncio.wait_for(
asyncio.to_thread(
subprocess.run,
["taskkill", "/PID", str(process.pid), "/T", "/F"],
check=False,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=5,
),
timeout=6.0,
)
else:
with suppress(ProcessLookupError, PermissionError):
os.killpg(process.pid, signal.SIGKILL)
if process.returncode is None:
with suppress(ProcessLookupError):
process.kill()
with suppress(asyncio.TimeoutError, ProcessLookupError):
await asyncio.wait_for(process.wait(), timeout=5.0)
async def run_async(
self,
name: str,
args: list[str] | None = None,
*,
json_output: bool = False,
working_dir: str | None = None,
timeout: int | None = None,
restrict_to_workspace: bool = False,
) -> str:
prepared = await asyncio.to_thread(
self._prepare_run,
name,
args,
json_output=json_output,
working_dir=working_dir,
timeout=timeout,
restrict_to_workspace=restrict_to_workspace,
)
process_kwargs: dict[str, Any] = {}
if os.name == "nt":
process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
else:
process_kwargs["start_new_session"] = True
job = _WindowsJob.create()
try:
process = await asyncio.create_subprocess_exec(
prepared.resolved,
*prepared.args,
cwd=str(prepared.cwd),
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=prepared.env,
**process_kwargs,
)
except BaseException:
if job is not None:
job.close(kill_descendants=False)
raise
if job is not None and not job.assign(process.pid):
job.close(kill_descendants=False)
job = None
try:
stdout_raw, stderr_raw = await asyncio.wait_for(
process.communicate(),
timeout=prepared.timeout,
)
except asyncio.TimeoutError:
await self._terminate_run_process(process, job)
return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
except asyncio.CancelledError:
await self._terminate_run_process(process, job)
raise
except BaseException:
await self._terminate_run_process(process, job)
raise
if job is not None:
job.close(kill_descendants=False)
stdout = stdout_raw.decode("utf-8", errors="replace")
stderr = stderr_raw.decode("utf-8", errors="replace")
return await asyncio.to_thread(
self._format_run_result,
prepared,
returncode=process.returncode or 0,
stdout=stdout,
stderr=stderr,
)
def run( def run(
self, self,
name: str, name: str,
@@ -1438,50 +1769,52 @@ Use the `run_cli_app` tool with `name="{name}"` for command execution. Do not in
timeout: int | None = None, timeout: int | None = None,
restrict_to_workspace: bool = False, restrict_to_workspace: bool = False,
) -> str: ) -> str:
app = self.get_app(name) prepared = self._prepare_run(
installed = self._load_installed() name,
if str(app["name"]) not in installed: args,
raise CliAppError(f"CLI app '{name}' is not installed") json_output=json_output,
cwd = self._resolve_cwd(working_dir, restrict_to_workspace=restrict_to_workspace) working_dir=working_dir,
entry = str(installed[str(app["name"])].get("entry_point") or app.get("entry_point") or "") timeout=timeout,
resolved = shutil.which(entry) restrict_to_workspace=restrict_to_workspace,
if not entry or not resolved: )
raise CliAppError(f"{entry or name} is not available on PATH") process_kwargs: dict[str, Any] = {}
clean_args = [str(arg) for arg in (args or [])] if os.name == "nt":
if json_output and "--json" not in clean_args: process_kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
clean_args = ["--json", *clean_args] else:
effective_timeout = max(1, min(timeout or self.runtime.run_timeout, 600)) process_kwargs["start_new_session"] = True
artifact_snapshot = self._artifact_snapshot(cwd) job = _WindowsJob.create()
try: try:
result = subprocess.run( process = subprocess.Popen(
[resolved, *clean_args], [prepared.resolved, *prepared.args],
cwd=str(cwd), cwd=str(prepared.cwd),
capture_output=True, stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True, text=True,
encoding="utf-8", encoding="utf-8",
errors="replace", errors="replace",
timeout=effective_timeout, env=prepared.env,
env=self._subprocess_env(), **process_kwargs,
) )
except BaseException:
if job is not None:
job.close(kill_descendants=False)
raise
if job is not None and not job.assign(process.pid):
job.close(kill_descendants=False)
job = None
try:
stdout, stderr = process.communicate(timeout=prepared.timeout)
except subprocess.TimeoutExpired: except subprocess.TimeoutExpired:
return f"CLI app '{name}' timed out after {effective_timeout}s" self._terminate_run_process_sync(process, job)
output = [ return f"CLI app '{prepared.name}' timed out after {prepared.timeout}s"
f"CLI app '{name}' exited {result.returncode}.", except BaseException:
f"Command: {entry} {' '.join(shlex.quote(arg) for arg in clean_args)}".rstrip(), self._terminate_run_process_sync(process, job)
] raise
if result.stdout: if job is not None:
output.append("\nSTDOUT:\n" + result.stdout.rstrip()) job.close(kill_descendants=False)
if result.stderr: return self._format_run_result(
output.append("\nSTDERR:\n" + result.stderr.rstrip()) prepared,
artifacts = self._changed_artifacts(cwd, artifact_snapshot) returncode=process.returncode,
if artifacts: stdout=stdout,
output.append( stderr=stderr,
"\nArtifacts created or updated:\n" )
+ "\n".join(self._format_artifact_lines(cwd, artifacts))
)
if any(path.suffix.lower() in _INLINE_ARTIFACT_EXTENSIONS for path in artifacts):
output.append(
"\nTo show a preview in WebUI, reference a raster artifact with Markdown "
"using its workspace-relative path, for example `![diagram](diagram.png)`."
)
return _truncate("\n".join(output))
+21 -34
View File
@@ -21,7 +21,10 @@ from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel from nanobot.channels.base import BaseChannel
from nanobot.config.schema import Base from nanobot.config.schema import Base
from nanobot.security.network import validate_resolved_url, validate_url_target from nanobot.security.network import (
async_validate_resolved_url,
async_validate_url_target,
)
DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024 DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024
DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3 DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3
@@ -182,12 +185,6 @@ class NanobotDingTalkHandler(_CallbackHandlerBase):
) )
) )
if not self.channel._accepting_inbound_tasks:
self.channel.logger.debug(
"Skipping DingTalk inbound dispatch during channel shutdown"
)
return AckMessage.STATUS_OK, "OK"
self.channel.logger.info("Received message from {} ({}): {}", sender_name, sender_id, content) self.channel.logger.info("Received message from {} ({}): {}", sender_name, sender_id, content)
# Forward to Nanobot via _on_message (non-blocking). # Forward to Nanobot via _on_message (non-blocking).
@@ -202,7 +199,7 @@ class NanobotDingTalkHandler(_CallbackHandlerBase):
) )
) )
self.channel._background_tasks.add(task) self.channel._background_tasks.add(task)
task.add_done_callback(self.channel._on_background_task_done) task.add_done_callback(self.channel._background_tasks.discard)
return AckMessage.STATUS_OK, "OK" return AckMessage.STATUS_OK, "OK"
@@ -262,17 +259,6 @@ class DingTalkChannel(BaseChannel):
# Hold references to background tasks to prevent GC # Hold references to background tasks to prevent GC
self._background_tasks: set[asyncio.Task[None]] = set() self._background_tasks: set[asyncio.Task[None]] = set()
self._accepting_inbound_tasks = True
def _on_background_task_done(self, task: asyncio.Task[None]) -> None:
self._background_tasks.discard(task)
if task.cancelled():
return
exception = task.exception()
if exception is not None:
self.logger.opt(exception=exception).error(
"DingTalk inbound message task failed"
)
async def start(self) -> None: async def start(self) -> None:
"""Start the DingTalk bot with Stream Mode.""" """Start the DingTalk bot with Stream Mode."""
@@ -289,7 +275,6 @@ class DingTalkChannel(BaseChannel):
self.logger.error("client_id and client_secret not configured") self.logger.error("client_id and client_secret not configured")
return return
self._accepting_inbound_tasks = True
self._running = True self._running = True
self._http = httpx.AsyncClient( self._http = httpx.AsyncClient(
timeout=httpx.Timeout(10.0, connect=10.0, read=30.0, write=30.0, pool=10.0) timeout=httpx.Timeout(10.0, connect=10.0, read=30.0, write=30.0, pool=10.0)
@@ -327,7 +312,6 @@ class DingTalkChannel(BaseChannel):
async def stop(self) -> None: async def stop(self) -> None:
"""Stop the DingTalk bot.""" """Stop the DingTalk bot."""
self._accepting_inbound_tasks = False
self._running = False self._running = False
await self._close_stream_client() await self._close_stream_client()
start_task = self._start_task start_task = self._start_task
@@ -345,11 +329,8 @@ class DingTalkChannel(BaseChannel):
await self._http.aclose() await self._http.aclose()
self._http = None self._http = None
# Cancel outstanding background tasks # Cancel outstanding background tasks
background_tasks = tuple(self._background_tasks) for task in self._background_tasks:
for task in background_tasks:
task.cancel() task.cancel()
if background_tasks:
await asyncio.gather(*background_tasks, return_exceptions=True)
self._background_tasks.clear() self._background_tasks.clear()
async def _close_stream_client(self) -> None: async def _close_stream_client(self) -> None:
@@ -439,8 +420,8 @@ class DingTalkChannel(BaseChannel):
return self._zip_bytes(filename, data) return self._zip_bytes(filename, data)
return data, filename, content_type return data, filename, content_type
def _validate_remote_media_url(self, media_ref: str) -> bool: async def _validate_remote_media_url(self, media_ref: str) -> bool:
ok, err = validate_url_target(media_ref) ok, err = await async_validate_url_target(media_ref)
if not ok: if not ok:
self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err) self.logger.warning("remote media URL blocked ref={} reason={}", media_ref, err)
return False return False
@@ -456,7 +437,11 @@ class DingTalkChannel(BaseChannel):
allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts} allowed_hosts = {host.lower() for host in self.config.remote_media_redirect_allowed_hosts}
return next_host in allowed_hosts return next_host in allowed_hosts
def _next_remote_media_url(self, current_url: str, location: str | None) -> str | None: async def _next_remote_media_url(
self,
current_url: str,
location: str | None,
) -> str | None:
if not self.config.allow_remote_media_redirects: if not self.config.allow_remote_media_redirects:
self.logger.warning("media download redirect refused ref={}", current_url) self.logger.warning("media download redirect refused ref={}", current_url)
return None return None
@@ -471,7 +456,7 @@ class DingTalkChannel(BaseChannel):
next_url, next_url,
) )
return None return None
if not self._validate_remote_media_url(next_url): if not await self._validate_remote_media_url(next_url):
return None return None
return next_url return next_url
@@ -483,7 +468,7 @@ class DingTalkChannel(BaseChannel):
if not self._http: if not self._http:
return None, None return None, None
if not self._validate_remote_media_url(media_ref): if not await self._validate_remote_media_url(media_ref):
return None, None return None, None
try: try:
@@ -495,7 +480,7 @@ class DingTalkChannel(BaseChannel):
current_url = media_ref current_url = media_ref
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1): for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
async with stream("GET", current_url, follow_redirects=False) as resp: async with stream("GET", current_url, follow_redirects=False) as resp:
final_ok, final_err = validate_resolved_url(str(resp.url)) final_ok, final_err = await async_validate_resolved_url(str(resp.url))
if not final_ok: if not final_ok:
self.logger.warning( self.logger.warning(
"remote media redirect blocked ref={} final={} reason={}", "remote media redirect blocked ref={} final={} reason={}",
@@ -505,7 +490,7 @@ class DingTalkChannel(BaseChannel):
) )
return None, None return None, None
if 300 <= resp.status_code < 400: if 300 <= resp.status_code < 400:
next_url = self._next_remote_media_url( next_url = await self._next_remote_media_url(
str(resp.url), resp.headers.get("location") str(resp.url), resp.headers.get("location")
) )
if not next_url: if not next_url:
@@ -538,7 +523,9 @@ class DingTalkChannel(BaseChannel):
current_url = media_ref current_url = media_ref
for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1): for _ in range(DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS + 1):
resp = await self._http.get(current_url, follow_redirects=False) resp = await self._http.get(current_url, follow_redirects=False)
final_ok, final_err = validate_resolved_url(str(getattr(resp, "url", current_url))) final_ok, final_err = await async_validate_resolved_url(
str(getattr(resp, "url", current_url))
)
if not final_ok: if not final_ok:
self.logger.warning( self.logger.warning(
"remote media redirect blocked ref={} final={} reason={}", "remote media redirect blocked ref={} final={} reason={}",
@@ -548,7 +535,7 @@ class DingTalkChannel(BaseChannel):
) )
return None, None return None, None
if 300 <= resp.status_code < 400: if 300 <= resp.status_code < 400:
next_url = self._next_remote_media_url( next_url = await self._next_remote_media_url(
str(getattr(resp, "url", current_url)), resp.headers.get("location") str(getattr(resp, "url", current_url)), resp.headers.get("location")
) )
if not next_url: if not next_url:
@@ -3,7 +3,7 @@ import json
import zipfile import zipfile
from io import BytesIO from io import BytesIO
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock from unittest.mock import AsyncMock
import httpx import httpx
import pytest import pytest
@@ -402,61 +402,6 @@ async def test_handler_uses_voice_recognition_text_when_text_is_empty(monkeypatc
assert msg.chat_id == "group:conv123" assert msg.chat_id == "group:conv123"
@pytest.mark.asyncio
async def test_handler_retrieves_background_message_failure(monkeypatch) -> None:
bus = MessageBus()
channel = DingTalkChannel(
DingTalkConfig(client_id="app", client_secret="secret", allow_from=["user1"]),
bus,
)
handler = NanobotDingTalkHandler(channel)
failure = RuntimeError("inbound dispatch failed")
mock_logger = MagicMock()
channel.logger = mock_logger
class _FakeChatbotMessage:
text = SimpleNamespace(content="hello")
extensions = {}
sender_staff_id = "user1"
sender_id = "fallback-user"
sender_nick = "Alice"
message_type = "text"
@staticmethod
def from_dict(_data):
return _FakeChatbotMessage()
async def fail(*_args) -> None:
raise failure
monkeypatch.setattr(dingtalk_module, "ChatbotMessage", _FakeChatbotMessage)
monkeypatch.setattr(dingtalk_module, "AckMessage", SimpleNamespace(STATUS_OK="OK"))
monkeypatch.setattr(channel, "_on_message", fail)
event_loop = asyncio.get_running_loop()
previous_handler = event_loop.get_exception_handler()
loop_errors: list[dict[str, object]] = []
event_loop.set_exception_handler(lambda _loop, context: loop_errors.append(context))
try:
status, body = await handler.process(
SimpleNamespace(data={"conversationType": "1", "text": {"content": "hello"}})
)
for _ in range(10):
await asyncio.sleep(0)
if not channel._background_tasks:
break
finally:
event_loop.set_exception_handler(previous_handler)
assert (status, body) == ("OK", "OK")
assert not channel._background_tasks
assert not loop_errors
mock_logger.opt.assert_called_once_with(exception=failure)
mock_logger.opt.return_value.error.assert_called_once_with(
"DingTalk inbound message task failed"
)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_handler_processes_file_message(monkeypatch) -> None: async def test_handler_processes_file_message(monkeypatch) -> None:
"""Test that file messages are handled and forwarded with downloaded path.""" """Test that file messages are handled and forwarded with downloaded path."""
@@ -506,72 +451,6 @@ async def test_handler_processes_file_message(monkeypatch) -> None:
assert "/tmp/nanobot_dingtalk/user1/report.xlsx" in msg.content assert "/tmp/nanobot_dingtalk/user1/report.xlsx" in msg.content
@pytest.mark.asyncio
async def test_handler_does_not_spawn_message_task_after_stop_during_download(
monkeypatch,
) -> None:
channel = DingTalkChannel(
DingTalkConfig(client_id="app", client_secret="secret", allow_from=["user1"]),
MessageBus(),
)
handler = NanobotDingTalkHandler(channel)
download_started = asyncio.Event()
release_download = asyncio.Event()
message_task_started = asyncio.Event()
class _FakeFileChatbotMessage:
text = None
extensions = {}
image_content = None
rich_text_content = None
sender_staff_id = "user1"
sender_id = "fallback-user"
sender_nick = "Alice"
message_type = "file"
@staticmethod
def from_dict(_data):
return _FakeFileChatbotMessage()
async def delayed_download(*_args):
download_started.set()
await release_download.wait()
return "/tmp/nanobot_dingtalk/user1/report.xlsx"
async def block_message(*_args) -> None:
message_task_started.set()
await asyncio.Future()
monkeypatch.setattr(dingtalk_module, "ChatbotMessage", _FakeFileChatbotMessage)
monkeypatch.setattr(dingtalk_module, "AckMessage", SimpleNamespace(STATUS_OK="OK"))
monkeypatch.setattr(channel, "_download_dingtalk_file", delayed_download)
monkeypatch.setattr(channel, "_on_message", block_message)
process_task = asyncio.create_task(handler.process(SimpleNamespace(data={
"conversationType": "1",
"content": {"downloadCode": "abc123", "fileName": "report.xlsx"},
"text": {"content": ""},
})))
await download_started.wait()
try:
await channel.stop()
release_download.set()
assert await process_task == ("OK", "OK")
await asyncio.sleep(0)
assert not message_task_started.is_set()
assert not channel._background_tasks
finally:
release_download.set()
if not process_task.done():
process_task.cancel()
pending = tuple(channel._background_tasks)
for task in pending:
task.cancel()
await asyncio.gather(process_task, *pending, return_exceptions=True)
def _rich_text_message(rich_text_list): def _rich_text_message(rich_text_list):
class _FakeRichTextChatbotMessage: class _FakeRichTextChatbotMessage:
text = None text = None
@@ -771,41 +650,6 @@ async def test_stop_cancels_stream_client_after_sdk_swallows_first_cancel(monkey
assert start_task.cancelled() assert start_task.cancelled()
@pytest.mark.asyncio
async def test_stop_waits_for_background_message_tasks() -> None:
channel = DingTalkChannel(
DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"]),
MessageBus(),
)
mock_logger = MagicMock()
channel.logger = mock_logger
started = asyncio.Event()
cancelled = asyncio.Event()
async def wait_forever() -> None:
started.set()
try:
await asyncio.Future()
finally:
cancelled.set()
task = asyncio.create_task(wait_forever())
channel._background_tasks.add(task)
task.add_done_callback(channel._on_background_task_done)
await started.wait()
try:
await channel.stop()
assert task.done()
assert cancelled.is_set()
assert not channel._background_tasks
mock_logger.opt.assert_not_called()
finally:
if not task.done():
task.cancel()
await asyncio.gather(task, return_exceptions=True)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_download_dingtalk_file(tmp_path, monkeypatch) -> None: async def test_download_dingtalk_file(tmp_path, monkeypatch) -> None:
"""Test the two-step file download flow (get URL then download content).""" """Test the two-step file download flow (get URL then download content)."""
+43 -58
View File
@@ -430,13 +430,7 @@ class EmailChannel(BaseChannel):
skipped_uids: set[str], skipped_uids: set[str],
cycle_uids: set[str], cycle_uids: set[str],
) -> list[dict[str, Any]] | None: ) -> list[dict[str, Any]] | None:
"""Fetch messages by arbitrary IMAP search criteria. """Fetch messages by arbitrary IMAP search criteria."""
Uses UID SEARCH so already-processed UIDs are recognized before any
FETCH at all, then fetches headers only to evaluate every filter the
full body (and any attachments) is downloaded only for messages that
pass every check and are actually going to be delivered.
"""
mailbox = self.config.imap_mailbox or "INBOX" mailbox = self.config.imap_mailbox or "INBOX"
client = self._open_imap_client(mailbox=mailbox, missing_mailbox_ok=True) client = self._open_imap_client(mailbox=mailbox, missing_mailbox_ok=True)
@@ -444,30 +438,29 @@ class EmailChannel(BaseChannel):
return messages return messages
try: try:
status, data = client.uid("SEARCH", None, *search_criteria) status, data = client.search(None, *search_criteria)
if status != "OK" or not data or not data[0]: if status != "OK" or not data:
return messages return messages
uids = [raw.decode("ascii", errors="ignore") for raw in data[0].split()] ids = data[0].split()
if limit > 0 and len(uids) > limit: if limit > 0 and len(ids) > limit:
uids = uids[-limit:] ids = ids[-limit:]
for imap_id in ids:
features: _ServerFeatures | None = None status, fetched = client.fetch(imap_id, "(BODY.PEEK[] UID)")
for uid in uids:
if not uid or uid in cycle_uids:
continue
if dedupe and uid in self._processed_uids:
continue
status, fetched = client.uid("FETCH", uid, "(BODY.PEEK[HEADER])")
if status != "OK" or not fetched: if status != "OK" or not fetched:
continue continue
header_bytes = self._extract_message_bytes(fetched)
if header_bytes is None: raw_bytes = self._extract_message_bytes(fetched)
if raw_bytes is None:
continue continue
parsed = BytesParser(policy=policy.default).parsebytes(header_bytes) uid = self._extract_uid(fetched)
if uid and uid in cycle_uids:
continue
if dedupe and uid and uid in self._processed_uids:
continue
parsed = BytesParser(policy=policy.default).parsebytes(raw_bytes)
sender = parseaddr(parsed.get("From", ""))[1].strip().lower() sender = parseaddr(parsed.get("From", ""))[1].strip().lower()
if not sender: if not sender:
continue continue
@@ -475,8 +468,9 @@ class EmailChannel(BaseChannel):
self.logger.info("From {} ignored: matches bot-owned address", sender) self.logger.info("From {} ignored: matches bot-owned address", sender)
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen: if mark_seen:
features = self._mark_seen_uid(client, uid, features) client.store(imap_id, "+FLAGS", "\\Seen")
skipped_uids.add(uid) if uid:
skipped_uids.add(uid)
continue continue
# --- Anti-spoofing: verify Authentication-Results --- # --- Anti-spoofing: verify Authentication-Results ---
@@ -488,7 +482,8 @@ class EmailChannel(BaseChannel):
sender, sender,
) )
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
skipped_uids.add(uid) if uid:
skipped_uids.add(uid)
continue continue
if self.config.verify_dkim and not dkim_pass: if self.config.verify_dkim and not dkim_pass:
self.logger.warning( self.logger.warning(
@@ -497,26 +492,18 @@ class EmailChannel(BaseChannel):
sender, sender,
) )
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
skipped_uids.add(uid) if uid:
skipped_uids.add(uid)
continue continue
if not self.is_allowed(sender): if not self.is_allowed(sender):
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen: if mark_seen:
features = self._mark_seen_uid(client, uid, features) client.store(imap_id, "+FLAGS", "\\Seen")
skipped_uids.add(uid) if uid:
skipped_uids.add(uid)
continue continue
# Passed every filter — only now fetch the full message body
# (and any attachments) for the message we're actually delivering.
status, full_fetched = client.uid("FETCH", uid, "(BODY.PEEK[])")
if status != "OK" or not full_fetched:
continue
raw_bytes = self._extract_message_bytes(full_fetched)
if raw_bytes is None:
continue
parsed = BytesParser(policy=policy.default).parsebytes(raw_bytes)
subject = self._decode_header_value(parsed.get("Subject", "")) subject = self._decode_header_value(parsed.get("Subject", ""))
date_value = parsed.get("Date", "") date_value = parsed.get("Date", "")
message_id = parsed.get("Message-ID", "").strip() message_id = parsed.get("Message-ID", "").strip()
@@ -569,19 +556,10 @@ class EmailChannel(BaseChannel):
self._remember_processed_uid(uid, dedupe, cycle_uids) self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen: if mark_seen:
features = self._mark_seen_uid(client, uid, features) client.store(imap_id, "+FLAGS", "\\Seen")
finally: finally:
self._close_imap_client(client) self._close_imap_client(client)
def _mark_seen_uid(
self, client: Any, uid: str, features: _ServerFeatures | None
) -> _ServerFeatures:
"""Mark a single UID \\Seen, reusing session-learned STORE support."""
if features is None:
features = self._server_features(client)
self._uid_store_flag(client, uid, "\\Seen", features)
return features
def _open_imap_client(self, mailbox: str, *, missing_mailbox_ok: bool = False) -> Any | None: def _open_imap_client(self, mailbox: str, *, missing_mailbox_ok: bool = False) -> Any | None:
if self.config.imap_use_ssl: if self.config.imap_use_ssl:
client: Any = imaplib.IMAP4_SSL(self.config.imap_host, self.config.imap_port) client: Any = imaplib.IMAP4_SSL(self.config.imap_host, self.config.imap_port)
@@ -736,14 +714,11 @@ class EmailChannel(BaseChannel):
return data[0].split()[0] return data[0].split()[0]
def _uid_store_deleted(self, client: Any, uid: str, features: _ServerFeatures) -> bool: def _uid_store_deleted(self, client: Any, uid: str, features: _ServerFeatures) -> bool:
return self._uid_store_flag(client, uid, "\\Deleted", features)
def _uid_store_flag(self, client: Any, uid: str, flag: str, features: _ServerFeatures) -> bool:
# Optimistic path: try UID STORE first because UID is stable and avoids # Optimistic path: try UID STORE first because UID is stable and avoids
# sequence-number lookup. If this fails once for the session, remember it # sequence-number lookup. If this fails once for the session, remember it
# and use the sequence STORE fallback directly for remaining UIDs. # and use the sequence STORE fallback directly for remaining UIDs.
if features.uid_store is not False: if features.uid_store is not False:
status, _ = client.uid("STORE", uid, "+FLAGS", f"({flag})") status, _ = client.uid("STORE", uid, "+FLAGS", "(\\Deleted)")
if status == "OK": if status == "OK":
features.uid_store = True features.uid_store = True
return True return True
@@ -753,12 +728,12 @@ class EmailChannel(BaseChannel):
# unreliable: resolve the current sequence number from UID and use STORE. # unreliable: resolve the current sequence number from UID and use STORE.
imap_id = self._lookup_imap_id_by_uid(client, uid) imap_id = self._lookup_imap_id_by_uid(client, uid)
if not imap_id: if not imap_id:
self.logger.warning("Could not locate UID {} to set flag {}", uid, flag) self.logger.warning("Post-action skipped: UID {} not found", uid)
return False return False
status, _ = client.store(imap_id, "+FLAGS", flag) status, _ = client.store(imap_id, "+FLAGS", "\\Deleted")
if status != "OK": if status != "OK":
self.logger.warning("Failed to set flag {} on UID {}", flag, uid) self.logger.warning("Post-action failed: could not mark UID {} as deleted", uid)
return False return False
return True return True
@@ -798,6 +773,16 @@ class EmailChannel(BaseChannel):
return bytes(fetched_item[1]) return bytes(fetched_item[1])
return None return None
@staticmethod
def _extract_uid(fetched: list[Any]) -> str:
for item in fetched:
if isinstance(item, tuple) and item and isinstance(item[0], (bytes, bytearray)):
head = bytes(item[0]).decode("utf-8", errors="ignore")
m = re.search(r"UID\s+(\d+)", head)
if m:
return m.group(1)
return ""
@staticmethod @staticmethod
def _decode_header_value(value: str) -> str: def _decode_header_value(value: str) -> str:
if not value: if not value:
@@ -53,7 +53,30 @@ def _make_raw_email(
def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None: def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
raw = _make_raw_email(subject="Invoice", body="Please pay") raw = _make_raw_email(subject="Invoice", body="Please pay")
fake = _make_fake_imap(raw, uid=b"123") class FakeIMAP:
def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = []
def login(self, _user: str, _pw: str):
return "OK", [b"logged in"]
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags))
return "OK", [b""]
def logout(self):
return "BYE", [b""]
fake = FakeIMAP()
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake) monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(), MessageBus()) channel = EmailChannel(_make_config(), MessageBus())
@@ -63,25 +86,38 @@ def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
assert items[0]["sender"] == "alice@example.com" assert items[0]["sender"] == "alice@example.com"
assert items[0]["subject"] == "Invoice" assert items[0]["subject"] == "Invoice"
assert "Please pay" in items[0]["content"] assert "Please pay" in items[0]["content"]
assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
assert [call for call in fake.uid_calls if call[0] == "FETCH"] == [
("FETCH", "123", "(BODY.PEEK[HEADER])"),
("FETCH", "123", "(BODY.PEEK[])"),
]
assert skipped_uids == set() assert skipped_uids == set()
# Same UID should be deduped in-process. # Same UID should be deduped in-process.
items_again, skipped_again = channel._fetch_new_messages() items_again, skipped_again = channel._fetch_new_messages()
assert items_again == [] assert items_again == []
assert skipped_again == set() assert skipped_again == set()
assert len([call for call in fake.uid_calls if call[0] == "FETCH"]) == 2
def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> None: def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> None:
raw = _make_raw_email(subject="Invoice", body="Please pay") raw = _make_raw_email(subject="Invoice", body="Please pay")
fake = _make_fake_imap(raw, uid=b"123") class FakeIMAP:
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake) def login(self, _user: str, _pw: str):
return "OK", [b"logged in"]
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, _imap_id: bytes, _op: str, _flags: str):
return "OK", [b""]
def logout(self):
return "BYE", [b""]
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: FakeIMAP())
channel = EmailChannel(_make_config(post_action="delete"), MessageBus()) channel = EmailChannel(_make_config(post_action="delete"), MessageBus())
items, skipped_uids = channel._fetch_new_messages() items, skipped_uids = channel._fetch_new_messages()
@@ -94,10 +130,26 @@ def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> No
def test_fetch_new_messages_rejected_returns_skipped_uid(monkeypatch) -> None: def test_fetch_new_messages_rejected_returns_skipped_uid(monkeypatch) -> None:
raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test") raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test")
monkeypatch.setattr( class FakeIMAP:
"nanobot.channels.email.runtime.imaplib.IMAP4_SSL", def login(self, _user: str, _pw: str):
lambda _h, _p: _make_fake_imap(raw, uid=b"123"), return "OK", [b"logged in"]
)
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, _imap_id: bytes, _op: str, _flags: str):
return "OK", [b""]
def logout(self):
return "BYE", [b""]
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: FakeIMAP())
channel_skip = EmailChannel( channel_skip = EmailChannel(
_make_config(from_address="bot@example.com", post_action="delete", post_action_ignore_skipped=True), _make_config(from_address="bot@example.com", post_action="delete", post_action_ignore_skipped=True),
@@ -493,7 +545,30 @@ async def test_start_keeps_post_actions_for_successful_emails_when_later_deliver
def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) -> None: def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) -> None:
raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test") raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test")
fake = _make_fake_imap(raw, uid=b"123") class FakeIMAP:
def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = []
def login(self, _user: str, _pw: str):
return "OK", [b"logged in"]
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags))
return "OK", [b""]
def logout(self):
return "BYE", [b""]
fake = FakeIMAP()
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake) monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(from_address="bot@example.com"), MessageBus()) channel = EmailChannel(_make_config(from_address="bot@example.com"), MessageBus())
@@ -501,7 +576,7 @@ def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) ->
assert items == [] assert items == []
assert skipped_uids == {"123"} assert skipped_uids == {"123"}
assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
# Same UID should still be deduped after being ignored. # Same UID should still be deduped after being ignored.
items_again, skipped_again = channel._fetch_new_messages() items_again, skipped_again = channel._fetch_new_messages()
@@ -539,14 +614,37 @@ def test_fetch_new_messages_skips_self_sent_across_identity_sources(
imap_username matches, and must be case-insensitive.""" imap_username matches, and must be case-insensitive."""
raw = _make_raw_email(from_addr=from_header, subject="Loop test") raw = _make_raw_email(from_addr=from_header, subject="Loop test")
fake = _make_fake_imap(raw, uid=b"123") class FakeIMAP:
def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = []
def login(self, _user: str, _pw: str):
return "OK", [b"logged in"]
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags))
return "OK", [b""]
def logout(self):
return "BYE", [b""]
fake = FakeIMAP()
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake) monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(**config_override), MessageBus()) channel = EmailChannel(_make_config(**config_override), MessageBus())
items, _ = channel._fetch_new_messages() items, _ = channel._fetch_new_messages()
assert items == [] assert items == []
assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeypatch) -> None: def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeypatch) -> None:
@@ -564,16 +662,15 @@ def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeyp
def select(self, _mailbox: str): def select(self, _mailbox: str):
return "OK", [b"1"] return "OK", [b"1"]
def uid(self, command: str, *args): def search(self, *_args):
if command == "SEARCH": self.search_calls += 1
self.search_calls += 1 if fail_once["pending"]:
if fail_once["pending"]: fail_once["pending"] = False
fail_once["pending"] = False raise imaplib.IMAP4.abort("socket error")
raise imaplib.IMAP4.abort("socket error") return "OK", [b"1"]
return "OK", [b"123"]
if command == "FETCH": def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"] return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
return "OK", [b""]
def store(self, imap_id: bytes, op: str, flags: str): def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags)) self.store_calls.append((imap_id, op, flags))
@@ -603,7 +700,10 @@ def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeyp
def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypatch) -> None: def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypatch) -> None:
raw_first = _make_raw_email(subject="First", body="First body") raw_first = _make_raw_email(subject="First", body="First body")
raw_second = _make_raw_email(subject="Second", body="Second body") raw_second = _make_raw_email(subject="Second", body="Second body")
mailbox_state = {"123": raw_first, "124": raw_second} mailbox_state = {
b"1": {"uid": b"123", "raw": raw_first, "seen": False},
b"2": {"uid": b"124", "raw": raw_second, "seen": False},
}
fail_once = {"pending": True} fail_once = {"pending": True}
class FlakyIMAP: class FlakyIMAP:
@@ -613,18 +713,20 @@ def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypa
def select(self, _mailbox: str): def select(self, _mailbox: str):
return "OK", [b"2"] return "OK", [b"2"]
def uid(self, command: str, *args): def search(self, *_args):
if command == "SEARCH": unseen_ids = [imap_id for imap_id, item in mailbox_state.items() if not item["seen"]]
keys = " ".join(sorted(mailbox_state.keys(), key=int)) return "OK", [b" ".join(unseen_ids)]
return "OK", [keys.encode()]
if command == "FETCH": def fetch(self, imap_id: bytes, _parts: str):
uid = args[0] if imap_id == b"2" and fail_once["pending"]:
if uid == "124" and fail_once["pending"]: fail_once["pending"] = False
fail_once["pending"] = False raise imaplib.IMAP4.abort("socket error")
raise imaplib.IMAP4.abort("socket error") item = mailbox_state[imap_id]
raw = mailbox_state[uid] header = b"%s (UID %s BODY[] {200})" % (imap_id, item["uid"])
header = f"{uid} (UID {uid} BODY[] {{200}})".encode() return "OK", [(header, item["raw"]), b")"]
return "OK", [(header, raw), b")"]
def store(self, imap_id: bytes, _op: str, _flags: str):
mailbox_state[imap_id]["seen"] = True
return "OK", [b""] return "OK", [b""]
def logout(self): def logout(self):
@@ -942,13 +1044,12 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
def select(self, _mailbox: str): def select(self, _mailbox: str):
return "OK", [b"1"] return "OK", [b"1"]
def uid(self, command: str, *args): def search(self, *_args):
if command == "SEARCH": self.search_args = _args
self.search_args = args return "OK", [b"5"]
return "OK", [b"999"]
if command == "FETCH": def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"5 (UID 999 BODY[] {200})", raw), b")"] return "OK", [(b"5 (UID 999 BODY[] {200})", raw), b")"]
return "OK", [b""]
def store(self, imap_id: bytes, op: str, flags: str): def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags)) self.store_calls.append((imap_id, op, flags))
@@ -969,7 +1070,7 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
assert len(items) == 1 assert len(items) == 1
assert items[0]["subject"] == "Status" assert items[0]["subject"] == "Status"
# uid("SEARCH", None, "SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026") # search(None, "SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
assert fake.search_args is not None assert fake.search_args is not None
assert fake.search_args[1:] == ("SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026") assert fake.search_args[1:] == ("SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
assert fake.store_calls == [] assert fake.store_calls == []
@@ -979,12 +1080,11 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
# Security: Anti-spoofing tests for Authentication-Results verification # Security: Anti-spoofing tests for Authentication-Results verification
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def _make_fake_imap(raw: bytes, uid: bytes = b"500"): def _make_fake_imap(raw: bytes):
"""Return a FakeIMAP class pre-loaded with the given raw email.""" """Return a FakeIMAP class pre-loaded with the given raw email."""
class FakeIMAP: class FakeIMAP:
def __init__(self) -> None: def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = [] self.store_calls: list[tuple[bytes, str, str]] = []
self.uid_calls: list[tuple] = []
def login(self, _user: str, _pw: str): def login(self, _user: str, _pw: str):
return "OK", [b"logged in"] return "OK", [b"logged in"]
@@ -992,16 +1092,11 @@ def _make_fake_imap(raw: bytes, uid: bytes = b"500"):
def select(self, _mailbox: str): def select(self, _mailbox: str):
return "OK", [b"1"] return "OK", [b"1"]
def capability(self): def search(self, *_args):
return "OK", [b"IMAP4rev1"] return "OK", [b"1"]
def uid(self, command: str, *args): def fetch(self, _imap_id: bytes, _parts: str):
self.uid_calls.append((command, *args)) return "OK", [(b"1 (UID 500 BODY[] {200})", raw), b")"]
if command == "SEARCH":
return "OK", [uid]
if command == "FETCH":
return "OK", [(b"1 (UID " + uid + b" BODY[] {200})", raw), b")"]
return "OK", [b""]
def store(self, imap_id: bytes, op: str, flags: str): def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags)) self.store_calls.append((imap_id, op, flags))
@@ -1197,10 +1292,7 @@ def test_fetch_new_messages_ignores_unauthorized_sender_before_attachments(monke
assert channel._fetch_new_messages() == ([], {"500"}) assert channel._fetch_new_messages() == ([], {"500"})
assert called["attachments"] is False assert called["attachments"] is False
assert [call for call in fake.uid_calls if call[0] == "FETCH"] == [ assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
("FETCH", "500", "(BODY.PEEK[HEADER])")
]
assert ("STORE", "500", "+FLAGS", "(\\Seen)") in fake.uid_calls
def test_extract_attachments_saves_pdf(tmp_path, monkeypatch) -> None: def test_extract_attachments_saves_pdf(tmp_path, monkeypatch) -> None:
+3 -3
View File
@@ -24,7 +24,7 @@ from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel 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 nanobot.security.network import validate_url_target from nanobot.security.network import async_validate_url_target
from nanobot.utils.helpers import safe_filename from nanobot.utils.helpers import safe_filename
_DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60) _DOWNLOAD_TIMEOUT = aiohttp.ClientTimeout(total=60)
@@ -473,7 +473,7 @@ class NapcatChannel(BaseChannel):
if not ref: if not ref:
return None return None
if ref.startswith(("http://", "https://")): if ref.startswith(("http://", "https://")):
ok, err = validate_url_target(ref) ok, err = await async_validate_url_target(ref)
if not ok: if not ok:
logger.warning("napcat: rejected remote image '{}': {}", ref, err) logger.warning("napcat: rejected remote image '{}': {}", ref, err)
return None return None
@@ -525,7 +525,7 @@ class NapcatChannel(BaseChannel):
# logger.debug("napcat: downloading image from {}", url) # logger.debug("napcat: downloading image from {}", url)
if self._http is None: if self._http is None:
return None return None
ok, err = validate_url_target(url) ok, err = await async_validate_url_target(url)
if not ok: if not ok:
logger.warning("napcat: skip image '{}': {}", url, err) logger.warning("napcat: skip image '{}': {}", url, err)
return None return None
@@ -149,9 +149,13 @@ async def test_download_image_rejects_redirects(tmp_path, monkeypatch) -> None:
channel = _channel() channel = _channel()
channel._media_root = tmp_path channel._media_root = tmp_path
channel._http = _FakeHttp(_FakeResponse(status=302)) channel._http = _FakeHttp(_FakeResponse(status=302))
async def allow_url(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr( monkeypatch.setattr(
"nanobot.channels.napcat.runtime.validate_url_target", "nanobot.channels.napcat.runtime.async_validate_url_target",
lambda _url: (True, ""), allow_url,
) )
result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"}) result = await channel._download_image({"url": "https://example.com/a.png", "file": "a.png"})
+2 -2
View File
@@ -40,7 +40,7 @@ from nanobot.bus.events import OutboundMessage
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.channels.base import BaseChannel from nanobot.channels.base import BaseChannel
from nanobot.config.schema import Base from nanobot.config.schema import Base
from nanobot.security.network import validate_url_target from nanobot.security.network import async_validate_url_target
from nanobot.utils.logging_bridge import redirect_lib_logging from nanobot.utils.logging_bridge import redirect_lib_logging
try: try:
@@ -458,7 +458,7 @@ class QQChannel(BaseChannel):
return None, None return None, None
# Remote URL # Remote URL
ok, err = validate_url_target(media_ref) ok, err = await async_validate_url_target(media_ref)
if not ok: if not ok:
self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err) self.logger.warning("outbound media URL validation failed url={} err={}", media_ref, err)
return None, None return None, None
+2 -2
View File
@@ -23,8 +23,8 @@ from nanobot.config.schema import Base
from nanobot.pairing import is_approved from nanobot.pairing import is_approved
from nanobot.security.network import ( from nanobot.security.network import (
PinnedDNSAsyncTransport, PinnedDNSAsyncTransport,
async_validate_url_target,
httpx_env_proxy_mounts, httpx_env_proxy_mounts,
validate_url_target,
) )
from nanobot.utils.helpers import safe_filename, split_message from nanobot.utils.helpers import safe_filename, split_message
@@ -95,7 +95,7 @@ _HTML_DOWNLOAD_PREFIXES = (b"<!doctype html", b"<html")
async def _validate_slack_download_request(request: httpx.Request) -> None: async def _validate_slack_download_request(request: httpx.Request) -> None:
"""Validate every Slack file request, including redirects, before transport.""" """Validate every Slack file request, including redirects, before transport."""
ok, error = validate_url_target(str(request.url)) ok, error = await async_validate_url_target(str(request.url))
if not ok: if not ok:
raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request) raise httpx.RequestError(f"unsafe Slack file URL: {error}", request=request)
@@ -859,13 +859,13 @@ def _patch_download_validation(
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
validated: list[str], validated: list[str],
) -> None: ) -> None:
def validate(url: str) -> tuple[bool, str]: async def validate(url: str) -> tuple[bool, str]:
validated.append(url) validated.append(url)
if "169.254.169.254" in url: if "169.254.169.254" in url:
return False, "blocked metadata address" return False, "blocked metadata address"
return True, "" return True, ""
monkeypatch.setattr("nanobot.channels.slack.runtime.validate_url_target", validate) monkeypatch.setattr("nanobot.channels.slack.runtime.async_validate_url_target", validate)
@pytest.mark.asyncio @pytest.mark.asyncio
+21 -73
View File
@@ -36,7 +36,7 @@ from nanobot.channels.base import BaseChannel
from nanobot.command.builtin import build_help_text from nanobot.command.builtin import build_help_text
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 nanobot.security.network import validate_url_target from nanobot.security.network import async_validate_url_target
from nanobot.utils.helpers import split_message from nanobot.utils.helpers import split_message
from nanobot.utils.logging_bridge import redirect_lib_logging from nanobot.utils.logging_bridge import redirect_lib_logging
@@ -897,68 +897,6 @@ class TelegramChannel(BaseChannel):
self.logger.debug("sendRichMessage failed: {}", exc) self.logger.debug("sendRichMessage failed: {}", exc)
return False return False
async def _try_edit_rich(self, chat_id: int, message_id: int, content: str) -> bool:
"""Upgrade an existing message to rich in place via editMessageText (Bot API 10.1).
Editing in place keeps the message identity, so the streaming preview is
upgraded without the delete-and-resend pattern that caused flickering and
dropped line breaks (issue #4470).
Returns True when the rich edit is in place (including the ambiguous
"message is not modified" retry outcome after a response timeout).
Returns False only when the legacy HTML path should take over:
capability errors (server older than Bot API 10.1, which also trip the
rich latch) and content-shaped BadRequest rejections. Transport,
rate-limit, and unexpected errors propagate so the final-edit retry
contract is preserved ChannelManager retries the buffered send
instead of an immediate legacy edit doubling connection demand.
"""
if not self._app:
return False
payload: dict[str, Any] = {
"chat_id": chat_id,
"message_id": message_id,
"rich_message": {
"markdown": content,
},
}
try:
await self._call_with_retry(
self._app.bot.do_api_request,
"editMessageText",
api_kwargs=payload,
)
return True
except BadRequest as exc:
if self._is_not_modified_error(exc):
# Ambiguous success: the rich edit was applied server-side but
# its response timed out, so the retry hit "message is not
# modified". Treat it as done rather than letting the legacy
# edit overwrite the already-successful rich result.
self.logger.debug("Rich stream edit already applied for {}", chat_id)
return True
# Before Bot API 10.1, editMessageText ignores rich_message and
# reports the absent text argument instead.
pre_rich_edit_server = (
bool(content)
and str(exc).strip().lower() == "message text is empty"
)
if self._is_rich_capability_error(exc) or pre_rich_edit_server:
self.logger.debug("editMessageText rich_message not available, disabling")
self._rich_send_disabled = True
return False
# Content-shaped rejections (invalid markdown, unsupported media in
# the rich payload, …) fall back to the legacy HTML edit.
self.logger.debug("editMessageText rich_message rejected: {}", exc)
return False
except Exception:
# Transport, rate-limit, and unexpected errors propagate so the
# final-edit retry contract stays intact: ChannelManager retries
# the buffered send instead of this handler doubling connection
# demand with an immediate legacy edit.
raise
async def send(self, msg: OutboundMessage) -> None: async def send(self, msg: OutboundMessage) -> None:
"""Send a message through Telegram.""" """Send a message through Telegram."""
app = await self._wait_for_app() app = await self._wait_for_app()
@@ -1018,7 +956,7 @@ class TelegramChannel(BaseChannel):
# Telegram Bot API accepts HTTP(S) URLs directly for media params. # Telegram Bot API accepts HTTP(S) URLs directly for media params.
if self._is_remote_media_url(media_path): if self._is_remote_media_url(media_path):
ok, error = validate_url_target(media_path) ok, error = await async_validate_url_target(media_path)
if not ok: if not ok:
raise ValueError(f"unsafe media URL: {error}") raise ValueError(f"unsafe media URL: {error}")
await self._call_with_retry( await self._call_with_retry(
@@ -1198,16 +1136,26 @@ class TelegramChannel(BaseChannel):
thread_kwargs["message_thread_id"] = message_thread_id thread_kwargs["message_thread_id"] = message_thread_id
raw_text = buf.text raw_text = buf.text
# Try upgrading the streaming preview to rich in place (Bot API 10.1: # Try sendRichMessage for final output (Bot API 10.1).
# editMessageText gained a rich_message parameter). Editing in place # Skip when a streaming preview already exists to avoid the
# keeps the message identity, so there is no delete-and-resend and # delete-and-resend pattern that causes flickering and drops
# none of the flickering / dropped line breaks from issue #4470. # line breaks (issue #4470).
# The previous branch here was unreachable: it was guarded by if not buf.message_id and self.config.rich_messages and not getattr(self, "_rich_send_disabled", False):
# ``not buf.message_id`` after an early return had already ensured reply_params = None
# ``buf.message_id`` is set (issue #5516). if reply_to_message_id := meta.get("message_id"):
if self.config.rich_messages and not getattr(self, "_rich_send_disabled", False): reply_params = {"message_id": int(reply_to_message_id), "allow_sending_without_reply": True}
rich_ok = await self._try_edit_rich(int_chat_id, buf.message_id, raw_text) rich_ok = await self._try_send_rich(
int_chat_id, raw_text, reply_params, thread_kwargs, None,
)
if rich_ok: if rich_ok:
# Delete the streaming preview message
try:
await self._call_with_retry(
app.bot.delete_message,
chat_id=int_chat_id, message_id=buf.message_id,
)
except Exception:
pass # Preview stays if delete fails
self._stream_bufs.pop(chat_id, None) self._stream_bufs.pop(chat_id, None)
return return
@@ -1488,7 +1488,14 @@ async def test_send_remote_media_url_after_security_validation(monkeypatch) -> N
MessageBus(), MessageBus(),
) )
_install_ready_app(channel) _install_ready_app(channel)
monkeypatch.setattr("nanobot.channels.telegram.runtime.validate_url_target", lambda url: (True, ""))
async def allow_url(_url: str) -> tuple[bool, str]:
return True, ""
monkeypatch.setattr(
"nanobot.channels.telegram.runtime.async_validate_url_target",
allow_url,
)
await channel.send( await channel.send(
OutboundMessage( OutboundMessage(
@@ -1546,9 +1553,13 @@ async def test_send_blocks_unsafe_remote_media_url(monkeypatch) -> None:
MessageBus(), MessageBus(),
) )
_install_ready_app(channel) _install_ready_app(channel)
async def deny_url(_url: str) -> tuple[bool, str]:
return False, "Blocked: example.com resolves to private/internal address 127.0.0.1"
monkeypatch.setattr( monkeypatch.setattr(
"nanobot.channels.telegram.runtime.validate_url_target", "nanobot.channels.telegram.runtime.async_validate_url_target",
lambda url: (False, "Blocked: example.com resolves to private/internal address 127.0.0.1"), deny_url,
) )
await channel.send( await channel.send(
@@ -2735,130 +2746,3 @@ def test_markdown_to_html_code_block_same_line_no_newline() -> None:
stripped = _strip_md_block(text) stripped = _strip_md_block(text)
assert stripped == "Use <tag> here" assert stripped == "Use <tag> here"
@pytest.mark.asyncio
async def test_send_delta_stream_end_upgrades_preview_to_rich_in_place() -> None:
"""Rich messages finally work with streaming: the preview is upgraded via
editMessageText rich_message (in place), not delete-and-resend (issue #5516)."""
from telegram.error import BadRequest
channel = TelegramChannel(
TelegramConfig(enabled=True, token="123:abc", allow_from=["*"], rich_messages=True),
MessageBus(),
)
_install_ready_app(channel)
channel._app.bot.do_api_request = AsyncMock()
channel._app.bot.edit_message_text = AsyncMock(side_effect=BadRequest("should not be reached"))
channel._stream_bufs["123"] = _StreamBuf(text="**hello**", message_id=7, last_edit=0.0)
await channel.send_delta("123", "", stream_end=True)
# editMessageText with rich_message payload, in place (same message_id)
channel._app.bot.do_api_request.assert_awaited_once()
args, kwargs = channel._app.bot.do_api_request.await_args
assert args[0] == "editMessageText"
assert kwargs["api_kwargs"]["chat_id"] == 123
assert kwargs["api_kwargs"]["message_id"] == 7
assert kwargs["api_kwargs"]["rich_message"] == {"markdown": "**hello**"}
# No delete-and-resend, no legacy HTML edit
channel._app.bot.edit_message_text.assert_not_awaited()
assert "123" not in channel._stream_bufs
@pytest.mark.asyncio
async def test_send_delta_stream_end_rich_capability_error_latches_and_falls_back() -> None:
"""On a pre-10.1 Bot API server the rich edit fails, the latch trips, and the
legacy HTML edit handles the final output."""
from telegram.error import BadRequest
channel = TelegramChannel(
TelegramConfig(enabled=True, token="123:abc", allow_from=["*"], rich_messages=True),
MessageBus(),
)
_install_ready_app(channel)
# Before Bot API 10.1, editMessageText ignores rich_message and requires text.
channel._app.bot.do_api_request = AsyncMock(
side_effect=BadRequest("Message text is empty")
)
channel._app.bot.edit_message_text = AsyncMock()
channel._stream_bufs["123"] = _StreamBuf(text="hello", message_id=7, last_edit=0.0)
await channel.send_delta("123", "", stream_end=True)
channel._app.bot.do_api_request.assert_awaited_once()
# Latch tripped: subsequent sends skip the rich path entirely
assert channel._rich_send_disabled is True
# Legacy HTML edit handled the final message
channel._app.bot.edit_message_text.assert_awaited_once()
assert "123" not in channel._stream_bufs
@pytest.mark.asyncio
async def test_send_delta_stream_end_rich_disabled_uses_legacy_html() -> None:
"""rich_messages=False (the default) keeps the legacy HTML path untouched."""
channel = TelegramChannel(
TelegramConfig(enabled=True, token="123:abc", allow_from=["*"]),
MessageBus(),
)
_install_ready_app(channel)
channel._app.bot.do_api_request = AsyncMock()
channel._app.bot.edit_message_text = AsyncMock()
channel._stream_bufs["123"] = _StreamBuf(text="hello", message_id=7, last_edit=0.0)
await channel.send_delta("123", "", stream_end=True)
channel._app.bot.do_api_request.assert_not_called()
channel._app.bot.edit_message_text.assert_awaited_once()
assert "123" not in channel._stream_bufs
@pytest.mark.asyncio
async def test_send_delta_stream_end_rich_network_error_propagates_for_retry() -> None:
"""A transport failure on the rich edit must propagate so ChannelManager
retries the buffered send not fall through to an immediate legacy edit
that doubles connection demand during pool exhaustion."""
from telegram.error import NetworkError
channel = TelegramChannel(
TelegramConfig(enabled=True, token="123:abc", allow_from=["*"], rich_messages=True),
MessageBus(),
)
_install_ready_app(channel)
channel._app.bot.do_api_request = AsyncMock(side_effect=NetworkError("pool exhausted"))
channel._app.bot.edit_message_text = AsyncMock()
channel._stream_bufs["123"] = _StreamBuf(text="hello", message_id=7, last_edit=0.0)
with pytest.raises(NetworkError):
await channel.send_delta("123", "", stream_end=True)
# No legacy fallback edit: the buffered state stays for the manager retry.
channel._app.bot.edit_message_text.assert_not_awaited()
assert "123" in channel._stream_bufs
@pytest.mark.asyncio
async def test_send_delta_stream_end_rich_not_modified_after_timeout_is_success() -> None:
"""Ambiguous success: the rich edit applied server-side but its response
timed out, so the retry hit "message is not modified". That is a completed
rich upgrade the legacy edit must not overwrite it."""
from telegram.error import BadRequest, TimedOut
channel = TelegramChannel(
TelegramConfig(enabled=True, token="123:abc", allow_from=["*"], rich_messages=True),
MessageBus(),
)
_install_ready_app(channel)
# First attempt (inside _call_with_retry) times out, retry reports the
# edit as already applied.
channel._app.bot.do_api_request = AsyncMock(
side_effect=[TimedOut(), BadRequest("Message is not modified")]
)
channel._app.bot.edit_message_text = AsyncMock(side_effect=AssertionError("must not overwrite rich result"))
channel._stream_bufs["123"] = _StreamBuf(text="hello", message_id=7, last_edit=0.0)
await channel.send_delta("123", "", stream_end=True)
assert channel._app.bot.do_api_request.await_count == 2
channel._app.bot.edit_message_text.assert_not_awaited()
assert "123" not in channel._stream_bufs
File diff suppressed because it is too large Load Diff
@@ -227,7 +227,6 @@ async def test_start_extends_http_open_timeout_for_slow_settings_routes(
return Server() return Server()
monkeypatch.setattr(websocket_module, "serve", fake_serve) monkeypatch.setattr(websocket_module, "serve", fake_serve)
monkeypatch.setattr(channel, "_listener_is_serving", lambda _server: True)
await channel.start() await channel.start()
@@ -1241,7 +1240,7 @@ def test_webui_request_cache_prunes_expired_completed_but_keeps_pending(
bus: MagicMock, bus: MagicMock,
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
) -> None: ) -> None:
import nanobot.webui.inbound_commands as websocket_module import nanobot.channels.websocket.runtime as websocket_module
channel = _ch(bus) channel = _ch(bus)
now = 1_000.0 now = 1_000.0
@@ -1264,7 +1263,7 @@ def test_webui_request_cache_prunes_oldest_completed_at_capacity(
bus: MagicMock, bus: MagicMock,
monkeypatch: pytest.MonkeyPatch, monkeypatch: pytest.MonkeyPatch,
) -> None: ) -> None:
import nanobot.webui.inbound_commands as websocket_module import nanobot.channels.websocket.runtime as websocket_module
channel = _ch(bus) channel = _ch(bus)
now = 1_000.0 now = 1_000.0
@@ -1556,7 +1555,7 @@ async def test_new_chat_without_message_does_not_create_session(
attached = json.loads(conn.send.await_args_list[0].args[0]) attached = json.loads(conn.send.await_args_list[0].args[0])
assert attached["event"] == "attached" assert attached["event"] == "attached"
assert sessions.list_sessions() == [] assert sessions.list_sessions() == []
assert channel.gateway.workspaces.scope_for_session_key( assert channel._workspaces.scope_for_session_key(
f"websocket:{attached['chat_id']}" f"websocket:{attached['chat_id']}"
).access_mode == "full" ).access_mode == "full"
@@ -1814,9 +1813,9 @@ async def test_webui_set_workspace_scope_rejects_running_chat(bus: MagicMock, tm
}, },
}, },
) )
channel.gateway.workspaces.persist_scope( channel._workspaces.persist_scope(
"chat-running", "chat-running",
channel.gateway.workspaces.scope_for_session_key("websocket:chat-running"), channel._workspaces.scope_for_session_key("websocket:chat-running"),
) )
conn.send.reset_mock() conn.send.reset_mock()
@@ -1961,7 +1960,7 @@ async def test_remote_access_reduction_rejects_stale_in_flight_message_scope(
await message_task await message_task
assert sessions.read_session_file(f"websocket:{chat_id}") is None assert sessions.read_session_file(f"websocket:{chat_id}") is None
assert channel.gateway.workspaces.scope_for_session_key( assert channel._workspaces.scope_for_session_key(
f"websocket:{chat_id}" f"websocket:{chat_id}"
).access_mode == "restricted" ).access_mode == "restricted"
payload = json.loads(message_conn.send.await_args.args[0]) payload = json.loads(message_conn.send.await_args.args[0])
@@ -2052,7 +2051,7 @@ async def test_native_webui_scope_allows_custom_scope_without_loopback(
assert payload["workspace_scope"]["sandbox_status"]["restrict_to_workspace"] is False assert payload["workspace_scope"]["sandbox_status"]["restrict_to_workspace"] is False
assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve()) assert payload["workspace_scope"]["sandbox_status"]["workspace_root"] == str(project.resolve())
assert sessions.read_session_file("websocket:chat-native") is None assert sessions.read_session_file("websocket:chat-native") is None
assert channel.gateway.workspaces.scope_for_session_key( assert channel._workspaces.scope_for_session_key(
"websocket:chat-native" "websocket:chat-native"
).metadata() == { ).metadata() == {
"project_path": str(project.resolve()), "project_path": str(project.resolve()),
@@ -2191,6 +2190,28 @@ async def test_send_scopes_turn_model_updates_to_the_subscribed_chat() -> None:
chat_two.send.assert_not_awaited() chat_two.send.assert_not_awaited()
def test_attach_fields_restore_the_session_model_and_latest_usage() -> None:
usage = LLMUsage.reported(input_tokens=120, output_tokens=8, total_tokens=175)
manager = MagicMock()
manager.read_session_metadata.return_value = {
"metadata": {
SESSION_MODEL_PRESET_METADATA_KEY: "Deep Research",
"_last_usage": usage.to_dict(),
}
}
bus = MagicMock()
channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"]},
bus,
gateway=_basic_handler(bus, session_manager=manager),
)
assert channel._attached_model_fields("chat-1") == {
"model_preset": "Deep Research",
"usage": usage.to_turn_dict(),
}
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_send_stages_external_media_as_signed_url(monkeypatch, tmp_path) -> None: async def test_send_stages_external_media_as_signed_url(monkeypatch, tmp_path) -> None:
bus = MagicMock() bus = MagicMock()
@@ -3438,20 +3459,20 @@ async def test_send_goal_state_emits_blob_per_chat() -> None:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_hydrate_noop_without_session_manager() -> None: async def test_maybe_push_active_goal_state_noop_without_session_manager() -> None:
bus = MagicMock() bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus)) channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus))
mock_ws = AsyncMock() mock_ws = AsyncMock()
channel._attach(mock_ws, "chat-1") channel._attach(mock_ws, "chat-1")
await channel._outbound.hydrate("chat-1") await channel._maybe_push_persisted_goal_state("chat-1")
mock_ws.send.assert_not_called() mock_ws.send.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_hydrate_skips_when_no_goal_on_disk() -> None: async def test_maybe_push_active_goal_state_skips_when_no_goal_on_disk() -> None:
bus = MagicMock() bus = MagicMock()
sm = MagicMock() sm = MagicMock()
sm.read_session_metadata.return_value = None sm.read_session_file.return_value = None
channel = WebSocketChannel( channel = WebSocketChannel(
{"enabled": True, "allowFrom": ["*"]}, {"enabled": True, "allowFrom": ["*"]},
bus, bus,
@@ -3459,15 +3480,15 @@ async def test_hydrate_skips_when_no_goal_on_disk() -> None:
) )
mock_ws = AsyncMock() mock_ws = AsyncMock()
channel._attach(mock_ws, "chat-1") channel._attach(mock_ws, "chat-1")
await channel._outbound.hydrate("chat-1") await channel._maybe_push_persisted_goal_state("chat-1")
mock_ws.send.assert_not_called() mock_ws.send.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_hydrate_notifies_when_goal_active_on_disk() -> None: async def test_maybe_push_active_goal_state_notifies_when_goal_active_on_disk() -> None:
bus = MagicMock() bus = MagicMock()
sm = MagicMock() sm = MagicMock()
sm.read_session_metadata.return_value = { sm.read_session_file.return_value = {
"metadata": { "metadata": {
"goal_state": { "goal_state": {
"status": "active", "status": "active",
@@ -3484,7 +3505,7 @@ async def test_hydrate_notifies_when_goal_active_on_disk() -> None:
) )
mock_ws = AsyncMock() mock_ws = AsyncMock()
channel._attach(mock_ws, "chat-1") channel._attach(mock_ws, "chat-1")
await channel._outbound.hydrate("chat-1") await channel._maybe_push_persisted_goal_state("chat-1")
mock_ws.send.assert_awaited_once() mock_ws.send.assert_awaited_once()
body = json.loads(mock_ws.send.await_args.args[0]) body = json.loads(mock_ws.send.await_args.args[0])
assert body["event"] == "goal_state" assert body["event"] == "goal_state"
@@ -3495,10 +3516,10 @@ async def test_hydrate_notifies_when_goal_active_on_disk() -> None:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_hydrate_restores_blocked_attention_on_disk() -> None: async def test_maybe_push_goal_state_restores_blocked_attention_on_disk() -> None:
bus = MagicMock() bus = MagicMock()
sm = MagicMock() sm = MagicMock()
sm.read_session_metadata.return_value = { sm.read_session_file.return_value = {
"metadata": { "metadata": {
"goal_state": { "goal_state": {
"status": "blocked", "status": "blocked",
@@ -3516,7 +3537,7 @@ async def test_hydrate_restores_blocked_attention_on_disk() -> None:
mock_ws = AsyncMock() mock_ws = AsyncMock()
channel._attach(mock_ws, "chat-1") channel._attach(mock_ws, "chat-1")
await channel._outbound.hydrate("chat-1") await channel._maybe_push_persisted_goal_state("chat-1")
body = json.loads(mock_ws.send.await_args.args[0]) body = json.loads(mock_ws.send.await_args.args[0])
assert body["goal_state"] == { assert body["goal_state"] == {
@@ -3528,7 +3549,7 @@ async def test_hydrate_restores_blocked_attention_on_disk() -> None:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_hydrate_skips_when_no_active_turn() -> None: async def test_maybe_push_turn_run_wall_clock_skips_when_no_active_turn() -> None:
bus = MagicMock() bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus)) channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus))
mock_ws = AsyncMock() mock_ws = AsyncMock()
@@ -3536,12 +3557,12 @@ async def test_hydrate_skips_when_no_active_turn() -> None:
from nanobot.session import webui_turns as wth from nanobot.session import webui_turns as wth
wth._WEBSOCKET_TURN_WALL_STARTED_AT.clear() wth._WEBSOCKET_TURN_WALL_STARTED_AT.clear()
await channel._outbound.hydrate("chat-1") await channel._maybe_push_turn_run_wall_clock("chat-1")
mock_ws.send.assert_not_called() mock_ws.send.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_hydrate_replays_running_turn() -> None: async def test_maybe_push_turn_run_wall_clock_replays_running() -> None:
bus = MagicMock() bus = MagicMock()
channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus)) channel = WebSocketChannel({"enabled": True, "allowFrom": ["*"]}, bus, gateway=_basic_handler(bus))
mock_ws = AsyncMock() mock_ws = AsyncMock()
@@ -3551,7 +3572,7 @@ async def test_hydrate_replays_running_turn() -> None:
wth._WEBSOCKET_TURN_WALL_STARTED_AT.clear() wth._WEBSOCKET_TURN_WALL_STARTED_AT.clear()
try: try:
wth._WEBSOCKET_TURN_WALL_STARTED_AT["chat-1"] = 1_700_000_000.0 wth._WEBSOCKET_TURN_WALL_STARTED_AT["chat-1"] = 1_700_000_000.0
await channel._outbound.hydrate("chat-1") await channel._maybe_push_turn_run_wall_clock("chat-1")
finally: finally:
wth._WEBSOCKET_TURN_WALL_STARTED_AT.pop("chat-1", None) wth._WEBSOCKET_TURN_WALL_STARTED_AT.pop("chat-1", None)
@@ -4,6 +4,7 @@ import asyncio
import json import json
import random import random
import socket import socket
import threading
import time import time
from contextlib import suppress from contextlib import suppress
from pathlib import Path from pathlib import Path
@@ -2576,6 +2577,69 @@ async def test_webui_automations_route_lists_all_jobs_and_allows_user_actions(
await server_task await server_task
@pytest.mark.asyncio
async def test_webui_cron_update_rearms_started_service_on_owner_loop(
bus: MagicMock,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
store_path = tmp_path / "cron" / "jobs.json"
cron = CronService(store_path, max_sleep_ms=60_000)
job = cron.add_job(
name="Before update",
schedule=CronSchedule(kind="every", every_ms=86_400_000),
message="Check the repo status",
session_key="websocket:abc",
origin_channel="websocket",
origin_chat_id="abc",
)
await cron.start()
owner_thread_id = threading.get_ident()
initial_timer = cron._timer_task
request_thread_ids: list[int] = []
arm_thread_ids: list[int] = []
timer_rearmed = asyncio.Event()
original_request_timer_rearm = cron._request_timer_rearm
original_arm_timer = cron._arm_timer
def tracked_request_timer_rearm() -> None:
request_thread_ids.append(threading.get_ident())
original_request_timer_rearm()
def tracked_arm_timer() -> None:
arm_thread_ids.append(threading.get_ident())
original_arm_timer()
timer_rearmed.set()
monkeypatch.setattr(cron, "_request_timer_rearm", tracked_request_timer_rearm)
monkeypatch.setattr(cron, "_arm_timer", tracked_arm_timer)
channel = _ch(bus, cron_service=cron, port=_free_port())
try:
response = await _webui_mutate(
channel,
"automation.update",
{"id": job.id, "values": {"name": "After update"}},
)
await asyncio.wait_for(timer_rearmed.wait(), timeout=1)
assert response.status_code == 200
assert request_thread_ids
assert all(thread_id != owner_thread_id for thread_id in request_thread_ids)
assert arm_thread_ids and set(arm_thread_ids) == {owner_thread_id}
assert cron._timer_task is not None
assert cron._timer_task is not initial_timer
assert not cron._timer_task.done()
stored = json.loads(store_path.read_text(encoding="utf-8"))
assert len(stored["jobs"]) == 1
assert stored["jobs"][0]["name"] == "After update"
finally:
cron.stop()
await asyncio.sleep(0)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_webui_automations_route_manages_local_triggers( async def test_webui_automations_route_manages_local_triggers(
bus: MagicMock, tmp_path: Path bus: MagicMock, tmp_path: Path
@@ -3769,3 +3833,77 @@ def test_bootstrap_secret_also_enforced_on_localhost(bus: MagicMock) -> None:
channel = _ch(bus, host="0.0.0.0", tokenIssueSecret="s3cret") channel = _ch(bus, host="0.0.0.0", tokenIssueSecret="s3cret")
resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS) resp = channel.gateway.http._handle_bootstrap(_LOCAL, _NO_HEADERS)
assert resp.status_code == 401 assert resp.status_code == 401
@pytest.mark.asyncio
async def test_webui_skill_update_cancellation_waits_for_config_and_runtime_state(
bus: MagicMock,
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
from nanobot.webui import ws_http
skill_dir = tmp_path / "skills" / "cancel-safe-skill"
skill_dir.mkdir(parents=True)
(skill_dir / "SKILL.md").write_text(
"---\nname: cancel-safe-skill\ndescription: Cancellation test skill.\n---\n",
encoding="utf-8",
)
mutation_started = threading.Event()
release_mutation = threading.Event()
original_update = ws_http.set_webui_skill_enabled
update_calls = 0
def blocked_update(*args: Any, **kwargs: Any) -> dict[str, Any]:
nonlocal update_calls
update_calls += 1
mutation_started.set()
assert release_mutation.wait(timeout=1)
return original_update(*args, **kwargs)
monkeypatch.setattr(ws_http, "set_webui_skill_enabled", blocked_update)
channel = _ch(
bus,
session_manager=_seed_session(tmp_path),
workspace_path=tmp_path,
port=_free_port(),
)
runtime_states: list[set[str]] = []
channel.gateway.http.skill_state_action = runtime_states.append
task = asyncio.create_task(
_webui_mutate(
channel,
"skill.update",
{"name": "cancel-safe-skill", "enabled": False},
)
)
assert await asyncio.to_thread(mutation_started.wait, 1)
try:
task.cancel()
await asyncio.sleep(0)
assert not task.done()
assert runtime_states == []
finally:
release_mutation.set()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
assert update_calls == 1
assert "cancel-safe-skill" in channel.gateway.http.disabled_skills
assert runtime_states == [{"cancel-safe-skill"}]
saved = load_config(channel.gateway.settings.config.path)
assert "cancel-safe-skill" in saved.agents.defaults.disabled_skills
settled_state = (
update_calls,
set(channel.gateway.http.disabled_skills),
list(runtime_states),
)
await asyncio.sleep(0.05)
assert (
update_calls,
set(channel.gateway.http.disabled_skills),
runtime_states,
) == settled_state
@@ -5,8 +5,6 @@ from unittest.mock import MagicMock, patch
import pytest import pytest
from nanobot.channels.websocket.runtime import WebSocketChannel from nanobot.channels.websocket.runtime import WebSocketChannel
from nanobot.webui.outbound_projection import WebUIOutboundProjector
from nanobot.webui.session_projection import WebUISessionProjection
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -15,9 +13,7 @@ async def test_hydrate_after_subscribe_is_quiet_when_no_turn_active():
channel = WebSocketChannel.__new__(WebSocketChannel) channel = WebSocketChannel.__new__(WebSocketChannel)
channel.gateway = MagicMock() channel.gateway = MagicMock()
channel.gateway.session_manager = MagicMock() channel.gateway.session_manager = MagicMock()
channel.gateway.session_manager.read_session_metadata = MagicMock(return_value={}) channel.gateway.session_manager.read_session_file = MagicMock(return_value={})
channel._session_projection = WebUISessionProjection(channel.gateway.session_manager)
channel._outbound = WebUIOutboundProjector(channel, channel._session_projection)
channel._turn_models = {} channel._turn_models = {}
sent_events = [] sent_events = []
@@ -31,7 +27,7 @@ async def test_hydrate_after_subscribe_is_quiet_when_no_turn_active():
channel.send_goal_state = mock_send_goal_state channel.send_goal_state = mock_send_goal_state
channel.send_goal_status = mock_send_goal_status channel.send_goal_status = mock_send_goal_status
with patch("nanobot.webui.session_projection.websocket_turn_wall_started_at", return_value=None): with patch("nanobot.channels.websocket.runtime.websocket_turn_wall_started_at", return_value=None):
await channel._hydrate_after_subscribe("test-chat") await channel._hydrate_after_subscribe("test-chat")
assert sent_events == [] assert sent_events == []
@@ -43,9 +39,7 @@ async def test_hydrate_after_subscribe_pushes_running_when_turn_active():
channel = WebSocketChannel.__new__(WebSocketChannel) channel = WebSocketChannel.__new__(WebSocketChannel)
channel.gateway = MagicMock() channel.gateway = MagicMock()
channel.gateway.session_manager = MagicMock() channel.gateway.session_manager = MagicMock()
channel.gateway.session_manager.read_session_metadata = MagicMock(return_value={}) channel.gateway.session_manager.read_session_file = MagicMock(return_value={})
channel._session_projection = WebUISessionProjection(channel.gateway.session_manager)
channel._outbound = WebUIOutboundProjector(channel, channel._session_projection)
channel._turn_models = {} channel._turn_models = {}
sent_events = [] sent_events = []
@@ -61,11 +55,11 @@ async def test_hydrate_after_subscribe_pushes_running_when_turn_active():
with ( with (
patch( patch(
"nanobot.webui.session_projection.websocket_turn_wall_started_at", "nanobot.channels.websocket.runtime.websocket_turn_wall_started_at",
return_value=1234567890.0, return_value=1234567890.0,
), ),
patch( patch(
"nanobot.webui.session_projection.websocket_turn_id", "nanobot.channels.websocket.runtime.websocket_turn_id",
return_value="turn-active", return_value="turn-active",
), ),
): ):
@@ -1,128 +0,0 @@
from __future__ import annotations
import asyncio
import errno
from unittest.mock import MagicMock
import pytest
from nanobot.bus.queue import MessageBus
from nanobot.channels.websocket.runtime import WebSocketChannel
class _FakeSocket:
def __init__(self) -> None:
self.open = True
def fileno(self) -> int:
return 1 if self.open else -1
def getsockopt(self, _level: int, _option: int) -> int:
return int(self.open)
class _FakeServer:
def __init__(self) -> None:
self.socket = _FakeSocket()
self.closed = False
@property
def sockets(self) -> tuple[_FakeSocket, ...]:
return (self.socket,)
def is_serving(self) -> bool:
return not self.closed
def close(self) -> None:
self.closed = True
self.socket.open = False
async def wait_closed(self) -> None:
return None
def _channel() -> WebSocketChannel:
gateway = MagicMock()
gateway.session_manager = None
return WebSocketChannel(
{"enabled": True, "allowFrom": ["*"]},
MessageBus(),
gateway=gateway,
)
@pytest.mark.asyncio
async def test_websocket_does_not_report_running_before_bind_succeeds(monkeypatch) -> None:
channel = _channel()
channel.logger = MagicMock()
bind_error = OSError(errno.EADDRINUSE, "address already in use")
async def fail_bind(*_args, **_kwargs):
raise bind_error
monkeypatch.setattr("nanobot.channels.websocket.runtime.serve", fail_bind)
with pytest.raises(OSError) as exc_info:
await channel.start()
assert exc_info.value is bind_error
assert channel.is_running is False
assert not any(
call.args and call.args[0] == "WebSocket server listening on {}"
for call in channel.logger.info.call_args_list
)
@pytest.mark.asyncio
async def test_websocket_restarts_only_its_listener_after_serving_socket_is_lost(
monkeypatch,
) -> None:
channel = _channel()
first = _FakeServer()
second = _FakeServer()
servers = iter((first, second))
bind_count = 0
rebound = asyncio.Event()
async def bind(*_args, **_kwargs):
nonlocal bind_count
bind_count += 1
server = next(servers)
if bind_count == 2:
rebound.set()
return server
monkeypatch.setattr("nanobot.channels.websocket.runtime.serve", bind)
monkeypatch.setattr(
"nanobot.channels.websocket.runtime._LISTENER_CHECK_INTERVAL_S",
0.01,
)
monkeypatch.setattr(
"nanobot.channels.websocket.runtime._LISTENER_RESTART_BACKOFF_S",
(0.05,),
)
start_task = asyncio.create_task(channel.start())
try:
for _ in range(20):
if channel.is_running:
break
await asyncio.sleep(0)
assert channel.is_running is True
first.socket.open = False
for _ in range(50):
if not channel.is_running:
break
await asyncio.sleep(0.005)
assert channel.is_running is False
assert bind_count == 1
await asyncio.wait_for(rebound.wait(), timeout=1)
assert channel.is_running is True
assert first.closed is True
finally:
await channel.stop()
await start_task
assert second.closed is True
+3 -12
View File
@@ -87,12 +87,7 @@ app = typer.Typer(
name="nanobot", name="nanobot",
context_settings={"help_option_names": ["-h", "--help"]}, context_settings={"help_option_names": ["-h", "--help"]},
help=f"{__logo__} nanobot - Personal AI Assistant", help=f"{__logo__} nanobot - Personal AI Assistant",
epilog=( no_args_is_help=True,
"Run `nanobot` without a subcommand to start the terminal agent. "
"Use `nanobot agent --help` for agent options."
),
invoke_without_command=True,
no_args_is_help=False,
) )
console = Console() console = Console()
@@ -103,7 +98,7 @@ def version_callback(value: bool):
raise typer.Exit() raise typer.Exit()
@app.callback(invoke_without_command=True) @app.callback()
def main( def main(
ctx: typer.Context, ctx: typer.Context,
version: bool = typer.Option( version: bool = typer.Option(
@@ -115,11 +110,7 @@ def main(
# imports this Typer app directly instead of ``nanobot.cli.entry``. Keep the # imports this Typer app directly instead of ``nanobot.cli.entry``. Keep the
# role identity correct until that launcher is regenerated. # role identity correct until that launcher is regenerated.
command = ctx.invoked_subcommand command = ctx.invoked_subcommand
set_cli_process_identity([command] if command else ["agent"]) set_cli_process_identity([command] if command else sys.argv[1:])
if command is None:
from nanobot.cli.entry import _run_agent
_run_agent([], prog_name="nanobot")
# ============================================================================ # ============================================================================
+10 -48
View File
@@ -8,28 +8,6 @@ from contextlib import suppress
from nanobot.cli.process_identity import set_cli_process_identity from nanobot.cli.process_identity import set_cli_process_identity
_ROOT_OPTIONS = frozenset(
{
"-h",
"--help",
"-v",
"--version",
"--install-completion",
"--show-completion",
}
)
def _agent_invocation_args(args: list[str]) -> list[str] | None:
"""Return agent arguments when the root command should act as ``agent``."""
if not args:
return []
if args[0] == "agent":
return args[1:]
if args[0].startswith("-") and args[0].split("=", 1)[0] not in _ROOT_OPTIONS:
return args
return None
def _native_tui_candidate(args: list[str]) -> bool: def _native_tui_candidate(args: list[str]) -> bool:
"""Return whether ``agent`` can start without the classic agent stack.""" """Return whether ``agent`` can start without the classic agent stack."""
@@ -56,35 +34,19 @@ def _configure_windows_console() -> None:
reconfigure(encoding="utf-8", errors="replace") reconfigure(encoding="utf-8", errors="replace")
def _run_agent(args: list[str], *, prog_name: str) -> None:
"""Run the shared agent command without importing the complete CLI graph."""
import typer
from nanobot.cli.agent import agent
agent_app = typer.Typer(add_completion=False)
agent_app.command()(agent)
command = typer.main.get_command(agent_app)
command.main(args=args, prog_name=prog_name)
def main() -> None: def main() -> None:
"""Dispatch native TUI startup without importing the complete CLI graph.""" """Dispatch native TUI startup without importing the complete CLI graph."""
raw_args = sys.argv[1:] set_cli_process_identity(sys.argv[1:])
# Installed completion scripts call ``nanobot`` without positional arguments
# and pass the request through this environment variable. Keep those requests
# on the root command so subcommands remain discoverable.
shell_completion = bool(os.environ.get("_NANOBOT_COMPLETE"))
agent_args = None if shell_completion else _agent_invocation_args(raw_args)
dispatch_args = ["agent", *agent_args] if agent_args is not None else raw_args
set_cli_process_identity(dispatch_args)
_configure_windows_console() _configure_windows_console()
root_agent_alias = agent_args is not None and raw_args[:1] != ["agent"] if _native_tui_candidate(sys.argv[1:]):
if agent_args is not None and ( import typer
root_agent_alias or _native_tui_candidate(dispatch_args)
): from nanobot.cli.agent import agent
prog_name = "nanobot" if root_agent_alias else "nanobot agent"
_run_agent(agent_args, prog_name=prog_name) fast_app = typer.Typer(add_completion=False)
fast_app.command()(agent)
command = typer.main.get_command(fast_app)
command.main(args=sys.argv[2:], prog_name="nanobot agent")
return return
from nanobot.cli.commands import app from nanobot.cli.commands import app
+77 -70
View File
@@ -12,7 +12,6 @@ from loguru import logger
from rich.console import Console from rich.console import Console
from nanobot import __logo__, __version__ from nanobot import __logo__, __version__
from nanobot.agent.hook import AgentHook, AgentRunHookContext
from nanobot.agent.hooks import create_file_edit_activity_hook from nanobot.agent.hooks import create_file_edit_activity_hook
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.mcp import MCPProvider from nanobot.agent.tools.mcp import MCPProvider
@@ -23,7 +22,6 @@ from nanobot.cli.webui_support import (
_gateway_health_bind_note, _gateway_health_bind_note,
_gateway_health_url, _gateway_health_url,
_host_for_local_browser, _host_for_local_browser,
_launch_browser,
_prepare_webui_bundle_for_gateway, _prepare_webui_bundle_for_gateway,
_print_foreground_port_conflict, _print_foreground_port_conflict,
_tcp_endpoint_reachable, _tcp_endpoint_reachable,
@@ -36,6 +34,7 @@ from nanobot.config.paths import is_default_workspace
from nanobot.config.schema import Config from nanobot.config.schema import Config
from nanobot.gateway.runtime import GatewayInstance from nanobot.gateway.runtime import GatewayInstance
from nanobot.security.network import is_loopback_host from nanobot.security.network import is_loopback_host
from nanobot.session.async_compat import call_session_manager as _call_session_manager
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt
from nanobot.utils.helpers import sync_workspace_templates from nanobot.utils.helpers import sync_workspace_templates
@@ -47,16 +46,29 @@ __all__ = ["_run_gateway"]
console = Console() console = Console()
_EVENT_LOOP_LAG_INTERVAL_S = 0.5
_EVENT_LOOP_LAG_WARNING_S = 0.25
class _MCPReadinessHook(AgentHook):
"""Retry application-owned MCP connections before the runner reads tools."""
def __init__(self, provider: MCPProvider) -> None: async def _monitor_event_loop_lag(
super().__init__() *,
self._provider = provider interval_s: float = _EVENT_LOOP_LAG_INTERVAL_S,
warning_threshold_s: float = _EVENT_LOOP_LAG_WARNING_S,
async def before_run(self, context: AgentRunHookContext) -> None: log: Any | None = None,
await self._provider.connect() ) -> None:
"""Log scheduler drift so gateway-wide stalls have direct evidence."""
loop = asyncio.get_running_loop()
lag_log = log or logger
while True:
expected = loop.time() + interval_s
await asyncio.sleep(interval_s)
lag_s = max(0.0, loop.time() - expected)
if lag_s >= warning_threshold_s:
lag_log.warning(
"event loop lag operation=gateway duration_ms={} interval_ms={}",
int(lag_s * 1000),
int(interval_s * 1000),
)
def _http_endpoint_responding(url: str, *, timeout_s: float = 0.25) -> bool: def _http_endpoint_responding(url: str, *, timeout_s: float = 0.25) -> bool:
@@ -247,44 +259,6 @@ def _print_gateway_health_endpoint(host: str, port: int) -> None:
) )
def _gateway_readiness_payload(channels: Any) -> tuple[bool, dict[str, object]]:
"""Describe process liveness separately from required WebSocket readiness."""
channel_status: dict[str, Any] = {}
get_status = getattr(channels, "get_status", None)
if callable(get_status):
try:
raw_status = get_status()
if isinstance(raw_status, dict):
channel_status = cast(dict[str, Any], raw_status)
except Exception:
logger.exception("Gateway readiness could not read channel status")
websocket = channel_status.get("websocket")
websocket_required = websocket is not None or "websocket" in getattr(
channels,
"enabled_channels",
(),
)
if not websocket_required:
websocket_state = "disabled"
ready = True
elif isinstance(websocket, dict):
websocket_status = cast(dict[str, Any], websocket)
ready = websocket_status.get("running") is True
state = websocket_status.get("state")
websocket_state = str(state) if isinstance(state, str) else "unavailable"
else:
ready = False
websocket_state = "unavailable"
return ready, {
"status": "ok" if ready else "degraded",
"process": "alive",
"ready": ready,
"websocket": websocket_state,
}
async def _close_gateway_runtime( async def _close_gateway_runtime(
agent: AgentLoop, agent: AgentLoop,
mcp_provider: MCPProvider, mcp_provider: MCPProvider,
@@ -496,7 +470,6 @@ def _run_gateway(
turn_delivery_factory=turn_delivery_factory, turn_delivery_factory=turn_delivery_factory,
provider_signature=provider_snapshot.signature, provider_signature=provider_snapshot.signature,
local_trigger_store=trigger_store, local_trigger_store=trigger_store,
hooks=[_MCPReadinessHook(mcp_provider)],
hook_factories=[create_file_edit_activity_hook], hook_factories=[create_file_edit_activity_hook],
tool_registry=tools, tool_registry=tools,
recovery_admission=recovery, recovery_admission=recovery,
@@ -545,12 +518,22 @@ def _run_gateway(
and hasattr(session_manager, "save") and hasattr(session_manager, "save")
): ):
key = session_key or _channel_session_key(msg.channel, msg.chat_id) key = session_key or _channel_session_key(msg.channel, msg.chat_id)
session = session_manager.get_or_create(key) session = await _call_session_manager(
session_manager,
"get_or_create_async",
session_manager.get_or_create,
key,
)
extra: dict[str, Any] = {"_channel_delivery": True} extra: dict[str, Any] = {"_channel_delivery": True}
if msg.media: if msg.media:
extra["media"] = list(msg.media) extra["media"] = list(msg.media)
session.add_message("assistant", msg.content, **extra) session.add_message("assistant", msg.content, **extra)
session_manager.save(session) await _call_session_manager(
session_manager,
"save_async",
session_manager.save,
session,
)
await bus.publish_outbound(msg) await bus.publish_outbound(msg)
message_tool = agent.tools.get("message") message_tool = agent.tools.get("message")
@@ -620,14 +603,14 @@ def _run_gateway(
if sha: if sha:
logger.info("Dream commit: {}", sha) logger.info("Dream commit: {}", sha)
store.compact_history() store.compact_history()
prune_dream_sessions(agent.sessions) await asyncio.to_thread(prune_dream_sessions, agent.sessions)
return None return None
# Heartbeat is a system job that checks HEARTBEAT.md for active tasks. # Heartbeat is a system job that checks HEARTBEAT.md for active tasks.
if job.name == "heartbeat": if job.name == "heartbeat":
heartbeat_file = config.workspace_path / "HEARTBEAT.md" heartbeat_file = config.workspace_path / "HEARTBEAT.md"
try: try:
content = heartbeat_file.read_text(encoding="utf-8") content = await asyncio.to_thread(heartbeat_file.read_text, encoding="utf-8")
except OSError: except OSError:
logger.debug("Heartbeat: HEARTBEAT.md missing") logger.debug("Heartbeat: HEARTBEAT.md missing")
return None return None
@@ -635,7 +618,7 @@ def _run_gateway(
logger.debug("Heartbeat: HEARTBEAT.md has no active tasks") logger.debug("Heartbeat: HEARTBEAT.md has no active tasks")
return None return None
channel, chat_id = _pick_heartbeat_target() channel, chat_id = await _pick_heartbeat_target()
if channel == "cli": if channel == "cli":
return None return None
@@ -663,9 +646,19 @@ def _run_gateway(
message_tool.reset_suppress_delivery(suppress_token) message_tool.reset_suppress_delivery(suppress_token)
# Keep a small tail of heartbeat history so the loop stays bounded. # Keep a small tail of heartbeat history so the loop stays bounded.
session = agent.sessions.get_or_create("heartbeat") session = await _call_session_manager(
agent.sessions,
"get_or_create_async",
agent.sessions.get_or_create,
"heartbeat",
)
session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages) session.retain_recent_legal_suffix(hb_cfg.keep_recent_messages)
agent.sessions.save(session) await _call_session_manager(
agent.sessions,
"save_async",
agent.sessions.save,
session,
)
if not resp or not resp.content: if not resp or not resp.content:
return return
@@ -742,17 +735,27 @@ def _run_gateway(
config_path=Path(config_path), config_path=Path(config_path),
) )
def _pick_heartbeat_target() -> tuple[str, str]: async def _pick_heartbeat_target() -> tuple[str, str]:
"""Pick a routable channel/chat target for heartbeat-triggered messages.""" """Pick a routable channel/chat target for heartbeat-triggered messages."""
sidebar_state = read_webui_sidebar_state() sidebar_state = await asyncio.to_thread(read_webui_sidebar_state)
unified_metadata = None unified_metadata = None
if config.agents.defaults.unified_session: if config.agents.defaults.unified_session:
record = session_manager.read_session_metadata(UNIFIED_SESSION_KEY) record = await _call_session_manager(
session_manager,
"read_session_metadata_async",
session_manager.read_session_metadata,
UNIFIED_SESSION_KEY,
)
if isinstance(record, dict) and isinstance(record.get("metadata"), dict): if isinstance(record, dict) and isinstance(record.get("metadata"), dict):
unified_metadata = record["metadata"] unified_metadata = record["metadata"]
sessions = await _call_session_manager(
session_manager,
"list_sessions_async",
session_manager.list_sessions,
)
return _pick_heartbeat_target_from_sessions( return _pick_heartbeat_target_from_sessions(
enabled_channels=channels.enabled_channels, enabled_channels=channels.enabled_channels,
sessions=session_manager.list_sessions(), sessions=sessions,
archived_keys=sidebar_state.get("archived_keys", []), archived_keys=sidebar_state.get("archived_keys", []),
unified_session_metadata=unified_metadata, unified_session_metadata=unified_metadata,
) )
@@ -797,9 +800,8 @@ def _run_gateway(
method, path = parts[0], parts[1] method, path = parts[0], parts[1]
if method == "GET" and path == "/health": if method == "GET" and path == "/health":
ready, payload = _gateway_readiness_payload(channels) body = _json.dumps({"status": "ok"})
body = _json.dumps(payload) status = "200 OK"
status = "200 OK" if ready else "503 Service Unavailable"
content_type = "application/json" content_type = "application/json"
else: else:
body = "Not Found" body = "Not Found"
@@ -865,6 +867,7 @@ def _run_gateway(
"""Wait for the gateway to bind, then point the user's browser at the webui.""" """Wait for the gateway to bind, then point the user's browser at the webui."""
if not open_browser_url: if not open_browser_url:
return return
import webbrowser
from urllib.parse import urlparse from urllib.parse import urlparse
# Channels start asynchronously. When the caller supplies a backend # Channels start asynchronously. When the caller supplies a backend
@@ -896,10 +899,8 @@ def _run_gateway(
await asyncio.sleep(0.1) await asyncio.sleep(0.1)
display_url = _webui_display_url(open_browser_url) display_url = _webui_display_url(open_browser_url)
try: try:
if _launch_browser(open_browser_url): webbrowser.open(open_browser_url)
console.print(f"[green]✓[/green] Opened browser at {display_url}") console.print(f"[green]✓[/green] Opened browser at {display_url}")
else:
console.print(f"[yellow]Could not open browser; visit {display_url}[/yellow]")
except Exception as e: except Exception as e:
console.print(f"[yellow]Could not open browser ({e}); visit {display_url}[/yellow]") console.print(f"[yellow]Could not open browser ({e}); visit {display_url}[/yellow]")
@@ -961,6 +962,10 @@ def _run_gateway(
_monitor_local_clients(), _monitor_local_clients(),
name="nanobot-gateway-client-monitor", name="nanobot-gateway-client-monitor",
), ),
asyncio.create_task(
_monitor_event_loop_lag(),
name="nanobot-event-loop-lag-monitor",
),
] ]
if health_server_enabled: if health_server_enabled:
tasks.append(asyncio.create_task( tasks.append(asyncio.create_task(
@@ -1028,13 +1033,15 @@ def _run_gateway(
# Flush all cached sessions to durable storage before exit. # Flush all cached sessions to durable storage before exit.
# This prevents data loss on filesystems with write-back # This prevents data loss on filesystems with write-back
# caching (rclone VFS, NFS, FUSE mounts, etc.). # caching (rclone VFS, NFS, FUSE mounts, etc.).
flushed = agent.sessions.flush_all() flushed = await _call_session_manager(
agent.sessions,
"flush_all_async",
agent.sessions.flush_all,
)
if flushed: if flushed:
logger.info("Shutdown: flushed {} session(s) to disk", flushed) logger.info("Shutdown: flushed {} session(s) to disk", flushed)
finally: finally:
restore_shutdown_handlers() restore_shutdown_handlers()
with gateway_runtime.foreground_instance(gateway_start_options): with gateway_runtime.foreground_instance(gateway_start_options):
if health_server_enabled:
gateway_runtime.publish_health_host(config.gateway.host)
asyncio.run(run()) asyncio.run(run())
+2 -5
View File
@@ -29,7 +29,7 @@ _PROVIDER_DISPLAY: dict[str, str] = {
_OAUTH_PROVIDER_DEFAULT_MODELS: dict[str, str] = { _OAUTH_PROVIDER_DEFAULT_MODELS: dict[str, str] = {
"openai_codex": "openai-codex/gpt-5.6-sol", "openai_codex": "openai-codex/gpt-5.6-sol",
"xai_grok": "xai-grok/grok-4.6", "xai_grok": "xai-grok/grok-4.5",
"github_copilot": "github-copilot/gpt-5.4-mini", "github_copilot": "github-copilot/gpt-5.4-mini",
} }
@@ -134,10 +134,7 @@ def _set_oauth_provider_as_main(
config.agents.defaults.model_preset = None config.agents.defaults.model_preset = None
config.agents.defaults.provider = provider_name config.agents.defaults.provider = provider_name
config.agents.defaults.model = selected_model config.agents.defaults.model = selected_model
if provider_name == "xai_grok" and selected_model in { if provider_name == "xai_grok" and selected_model == "xai-grok/grok-4.5":
"xai-grok/grok-4.5",
"xai-grok/grok-4.6",
}:
config.agents.defaults.context_window_tokens = 500_000 config.agents.defaults.context_window_tokens = 500_000
save_config(config, resolved_config_path) save_config(config, resolved_config_path)
+23 -57
View File
@@ -21,14 +21,12 @@ from nanobot.cli.process_identity import named_executable
from nanobot.cli.runtime_config import _model_display from nanobot.cli.runtime_config import _model_display
from nanobot.cli.webui_support import ( from nanobot.cli.webui_support import (
_gateway_health_ready, _gateway_health_ready,
_gateway_health_url,
_gateway_instance_command, _gateway_instance_command,
_host_for_local_browser, _host_for_local_browser,
_webui_endpoint_reachable, _webui_endpoint_reachable,
) )
from nanobot.config.paths import get_data_dir from nanobot.config.paths import get_data_dir
from nanobot.config.schema import Config from nanobot.config.schema import Config
from nanobot.webui.session_identity import is_webui_session_key, webui_chat_id
if TYPE_CHECKING: if TYPE_CHECKING:
from nanobot.gateway import GatewayClientLease from nanobot.gateway import GatewayClientLease
@@ -65,8 +63,6 @@ _TUI_RELEASE_LIMITS = {
} }
# Keep in sync with TUI_DETACH_EXIT_CODE in tui/src/index.ts. # Keep in sync with TUI_DETACH_EXIT_CODE in tui/src/index.ts.
_TUI_DETACH_EXIT_CODE = 90 _TUI_DETACH_EXIT_CODE = 90
_GATEWAY_READY_TIMEOUT_S = 20.0
_GATEWAY_READY_POLL_S = 0.1
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -97,10 +93,6 @@ def launch_tui(
env.update( env.update(
{ {
"NANOBOT_TUI_BOOTSTRAP_URL": f"{base_url}/webui/bootstrap", "NANOBOT_TUI_BOOTSTRAP_URL": f"{base_url}/webui/bootstrap",
"NANOBOT_TUI_HEALTH_URL": _gateway_health_url(
config.gateway.host,
config.gateway.port,
),
"NANOBOT_TUI_API_URL": base_url, "NANOBOT_TUI_API_URL": base_url,
"NANOBOT_TUI_MODEL": _model_display(config)[0], "NANOBOT_TUI_MODEL": _model_display(config)[0],
"NANOBOT_TUI_MODEL_PRESET": config.agents.defaults.model_preset or "default", "NANOBOT_TUI_MODEL_PRESET": config.agents.defaults.model_preset or "default",
@@ -424,52 +416,17 @@ def _ensure_gateway(
lease = GatewayClientLease(runtime, kind="tui") lease = GatewayClientLease(runtime, kind="tui")
lease.acquire() lease.acquire()
try: try:
def ready(status: object) -> bool:
management_ready = getattr(status, "ready", None)
if not isinstance(management_ready, bool):
management_ready = _gateway_health_ready(
config.gateway.host,
config.gateway.port,
)
return _webui_endpoint_reachable(base_url) and management_ready
def wait_for_ready(log_path: object) -> _GatewayHandle:
deadline = time.monotonic() + _GATEWAY_READY_TIMEOUT_S
while time.monotonic() < deadline:
current = runtime.status()
if not current.running:
break
if current.port not in {None, config.gateway.port}:
break
if ready(current):
return _GatewayHandle(base_url=base_url, lease=lease)
time.sleep(_GATEWAY_READY_POLL_S)
current = runtime.status()
if current.running:
raise TuiUnavailableError(
"local gateway process is running but its WebSocket/WebUI listener "
"is unavailable; channel recovery did not restore it. "
"Run `nanobot gateway status` and inspect logs at "
f"{log_path}; if it remains degraded, run `nanobot gateway restart`."
)
raise TuiUnavailableError(
f"local gateway did not become ready; logs: {log_path}"
)
status = runtime.status() status = runtime.status()
endpoint_reachable = _webui_endpoint_reachable(base_url)
if status.running: if status.running:
if status.port not in {None, config.gateway.port}: if status.port not in {None, config.gateway.port}:
raise TuiUnavailableError( raise TuiUnavailableError(
"the matching gateway instance is running on a different port; " "the matching gateway instance is running on a different port; "
"restart it or use `nanobot agent --classic`" "restart it or use `nanobot agent --classic`"
) )
if not wait_until_ready: if endpoint_reachable or not wait_until_ready:
return _GatewayHandle(base_url=base_url, lease=lease) return _GatewayHandle(base_url=base_url, lease=lease)
if ready(status): elif endpoint_reachable:
return _GatewayHandle(base_url=base_url, lease=lease)
return wait_for_ready(status.log_path)
elif _webui_endpoint_reachable(base_url):
raise TuiUnavailableError( raise TuiUnavailableError(
"the configured gateway port belongs to a different nanobot instance; " "the configured gateway port belongs to a different nanobot instance; "
"stop that instance or use `nanobot agent --classic`" "stop that instance or use `nanobot agent --classic`"
@@ -484,17 +441,26 @@ def _ensure_gateway(
f"logs: {result.status.log_path}" f"logs: {result.status.log_path}"
) )
if result.message == "gateway_already_running" and result.status.port not in {
None,
config.gateway.port,
}:
raise TuiUnavailableError(
"the matching gateway instance is running on a different port; "
"restart it or use `nanobot agent --classic`"
)
if not wait_until_ready: if not wait_until_ready:
return _GatewayHandle(base_url=base_url, lease=lease) return _GatewayHandle(base_url=base_url, lease=lease)
return wait_for_ready(result.status.log_path)
deadline = time.monotonic() + 20
while time.monotonic() < deadline:
if _webui_endpoint_reachable(base_url):
current = runtime.status()
if current.running and current.port in {None, config.gateway.port}:
return _GatewayHandle(base_url=base_url, lease=lease)
break
if not runtime.status().running and not _gateway_health_ready(
config.gateway.host,
config.gateway.port,
):
break
time.sleep(0.1)
raise TuiUnavailableError(
f"local gateway did not become ready; logs: {result.status.log_path}"
)
except BaseException: except BaseException:
lease.release(timeout_s=5) lease.release(timeout_s=5)
raise raise
@@ -520,8 +486,8 @@ def _tui_gateway_connection(config: Config) -> tuple[str, str]:
def _websocket_chat_id(session_id: str) -> str | None: def _websocket_chat_id(session_id: str) -> str | None:
"""Map the CLI selector to the WebSocket namespace used by the native TUI.""" """Map the CLI selector to the WebSocket namespace used by the native TUI."""
if is_webui_session_key(session_id): if session_id.startswith("websocket:"):
return webui_chat_id(session_id) return session_id.split(":", 1)[1] or None
if ":" in session_id: if ":" in session_id:
raise TuiSessionError( raise TuiSessionError(
"the native TUI can open only WebSocket sessions; use --classic to resume " "the native TUI can open only WebSocket sessions; use --classic to resume "
+8 -104
View File
@@ -1,14 +1,10 @@
"""Shared WebUI setup, URL, health, and browser helpers.""" """Shared WebUI setup, URL, health, and browser helpers."""
import os
import subprocess
import sys import sys
import time import time
import webbrowser
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, BinaryIO from typing import TYPE_CHECKING, Any
import typer import typer
from pydantic import ValidationError from pydantic import ValidationError
@@ -44,7 +40,6 @@ __all__ = [
"_gateway_instance_command", "_gateway_instance_command",
"_host_for_local_browser", "_host_for_local_browser",
"_load_webui_setup_config", "_load_webui_setup_config",
"_launch_browser",
"_open_webui_browser", "_open_webui_browser",
"_prepare_webui_bundle_for_gateway", "_prepare_webui_bundle_for_gateway",
"_print_foreground_port_conflict", "_print_foreground_port_conflict",
@@ -65,20 +60,6 @@ __all__ = [
console = Console() console = Console()
def _launch_browser(url: str) -> bool:
"""Open *url* and request a foreground browser window."""
if sys.platform == "darwin":
result = subprocess.run(
["open", url],
stdin=subprocess.DEVNULL,
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
check=False,
)
return result.returncode == 0
return bool(webbrowser.open(url, new=2, autoraise=True))
def _confirm_webui_action(message: str, *, yes: bool) -> None: def _confirm_webui_action(message: str, *, yes: bool) -> None:
"""Confirm a WebUI first-run mutation or fail clearly in non-interactive shells.""" """Confirm a WebUI first-run mutation or fail clearly in non-interactive shells."""
if yes: if yes:
@@ -438,14 +419,14 @@ def _print_foreground_port_conflict(
def _open_webui_browser(url: str, *, wait: bool = True) -> None: def _open_webui_browser(url: str, *, wait: bool = True) -> None:
"""Open the WebUI in the user's default browser, with a copyable fallback.""" """Open the WebUI in the user's default browser, with a copyable fallback."""
import webbrowser
if wait: if wait:
_wait_for_webui(url) _wait_for_webui(url)
display_url = _webui_display_url(url) display_url = _webui_display_url(url)
try: try:
if _launch_browser(url): webbrowser.open(url)
console.print(f"[green]✓[/green] Opened WebUI: [cyan]{display_url}[/cyan]") console.print(f"[green]✓[/green] Opened WebUI: [cyan]{display_url}[/cyan]")
else:
console.print(f"[yellow]Could not open browser; visit {display_url}[/yellow]")
except Exception as exc: except Exception as exc:
console.print(f"[yellow]Could not open browser ({exc}); visit {display_url}[/yellow]") console.print(f"[yellow]Could not open browser ({exc}); visit {display_url}[/yellow]")
@@ -459,104 +440,27 @@ def _print_webui_foreground_lifecycle(*, attached: bool) -> None:
console.print("[green]WebUI is attached to the shared gateway.[/green]") console.print("[green]WebUI is attached to the shared gateway.[/green]")
console.print("[dim]Closing the browser does not stop channels or automations.[/dim]") console.print("[dim]Closing the browser does not stop channels or automations.[/dim]")
console.print( console.print(
"[dim]Following live gateway logs. Press Ctrl+C to detach; the gateway stops " "[dim]Press Ctrl+C to detach; the gateway stops only when the last local client exits.[/dim]"
"only when the last local client exits.[/dim]"
) )
_LOG_ANCHOR_BYTES = 64
@dataclass
class _GatewayLogCursor:
offset: int = 0
identity: tuple[int, int] | None = None
anchor: bytes = b""
pending: bytes = b""
def _log_anchor(handle: BinaryIO, offset: int) -> bytes:
size = min(offset, _LOG_ANCHOR_BYTES)
handle.seek(offset - size)
return handle.read(size)
def _start_gateway_log_cursor(log_path: Path) -> _GatewayLogCursor:
"""Start following at the current end of *log_path*."""
try:
with log_path.open("rb") as handle:
stat = os.fstat(handle.fileno())
offset = stat.st_size
return _GatewayLogCursor(
offset=offset,
identity=(stat.st_dev, stat.st_ino),
anchor=_log_anchor(handle, offset),
)
except OSError:
return _GatewayLogCursor()
def _read_new_gateway_logs(
log_path: Path,
cursor: _GatewayLogCursor,
*,
flush: bool = False,
) -> list[str]:
"""Read complete gateway log lines appended after *cursor*."""
try:
with log_path.open("rb") as handle:
stat = os.fstat(handle.fileno())
identity = (stat.st_dev, stat.st_ino)
reset = cursor.identity != identity or stat.st_size < cursor.offset
if not reset and cursor.offset:
reset = _log_anchor(handle, cursor.offset) != cursor.anchor
if reset:
cursor.offset = 0
cursor.pending = b""
handle.seek(cursor.offset)
chunk = handle.read()
cursor.offset = handle.tell()
cursor.identity = identity
cursor.anchor = _log_anchor(handle, cursor.offset)
except OSError:
return []
parts = (cursor.pending + chunk).split(b"\n")
cursor.pending = parts.pop()
if flush and cursor.pending:
parts.append(cursor.pending)
cursor.pending = b""
return [part.removesuffix(b"\r").decode("utf-8", errors="replace") for part in parts]
def _attach_to_background_gateway( def _attach_to_background_gateway(
runtime: "GatewayRuntime", runtime: "GatewayRuntime",
*, *,
poll_hook: Callable[[], None] | None = None, poll_hook: Callable[[], None] | None = None,
sleep: Callable[[float], None] = time.sleep, sleep: Callable[[float], None] = time.sleep,
) -> None: ) -> None:
"""Keep the launcher attached and mirror this gateway's new log output.""" """Keep a WebUI launcher attached without taking ownership of the gateway."""
status = runtime.status()
log_path = status.log_path
cursor = _start_gateway_log_cursor(log_path)
_print_webui_foreground_lifecycle(attached=True) _print_webui_foreground_lifecycle(attached=True)
try: try:
while status.running: while runtime.status().running:
for line in _read_new_gateway_logs(log_path, cursor):
console.print(line, markup=False, highlight=False)
if poll_hook is not None: if poll_hook is not None:
poll_hook() poll_hook()
sleep(0.5) sleep(0.5)
status = runtime.status()
except KeyboardInterrupt: except KeyboardInterrupt:
for line in _read_new_gateway_logs(log_path, cursor, flush=True):
console.print(line, markup=False, highlight=False)
console.print("\n[yellow]WebUI launcher detached.[/yellow]") console.print("\n[yellow]WebUI launcher detached.[/yellow]")
return return
for line in _read_new_gateway_logs(log_path, cursor, flush=True):
console.print(line, markup=False, highlight=False)
console.print("[yellow]Gateway stopped.[/yellow]") console.print("[yellow]Gateway stopped.[/yellow]")
+73 -21
View File
@@ -3,6 +3,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import inspect
import os import os
import subprocess import subprocess
import sys import sys
@@ -14,7 +15,8 @@ from typing import TYPE_CHECKING, Any, Literal, cast
from nanobot import __version__ from nanobot import __version__
from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage from nanobot.bus.events import INBOUND_META_USER_SHELL, OutboundMessage
from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text from nanobot.command.router import CommandContext, CommandRouter, normalize_command_text
from nanobot.providers.base import LLMUsage from nanobot.session.async_compat import call_session_manager
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import build_status_content from nanobot.utils.helpers import build_status_content
from nanobot.utils.restart import set_restart_notice_to_env from nanobot.utils.restart import set_restart_notice_to_env
from nanobot.utils.workspace_prompts import initialize_workspace_prompt from nanobot.utils.workspace_prompts import initialize_workspace_prompt
@@ -23,6 +25,7 @@ if TYPE_CHECKING:
from nanobot.agent.loop import AgentLoop from nanobot.agent.loop import AgentLoop
from nanobot.session.manager import Session from nanobot.session.manager import Session
from nanobot.utils.gitstore import CommitInfo from nanobot.utils.gitstore import CommitInfo
from nanobot.utils.llm_runtime import LLMRuntime
# WebUI protocol contract for how a slash command participates in turn state: # WebUI protocol contract for how a slash command participates in turn state:
# - side_channel: returns control text without starting or ending an agent turn. # - side_channel: returns control text without starting or ending an agent turn.
@@ -202,6 +205,52 @@ def builtin_command_starts_agent_turn(text: str) -> bool:
return spec.lifecycle == "agent_turn_with_args" and bool(args.strip()) return spec.lifecycle == "agent_turn_with_args" and bool(args.strip())
def _has_native_coroutine_method(target: object, name: str) -> bool:
"""Check the real target class without trusting dynamic mock attributes."""
method = inspect.getattr_static(type(target), name, None)
return inspect.iscoroutinefunction(method)
async def _get_or_create_session(loop: AgentLoop, key: str) -> Session:
sessions = loop.sessions
return await call_session_manager(
sessions,
"get_or_create_async",
sessions.get_or_create,
key,
)
async def _save_session(loop: AgentLoop, session: Session) -> None:
sessions = loop.sessions
await call_session_manager(
sessions,
"save_async",
sessions.save,
session,
)
async def _runtime_for_session(loop: AgentLoop, session: Session) -> LLMRuntime:
if _has_native_coroutine_method(loop, "runtime_for_session_async"):
return await loop.runtime_for_session_async(session)
return await shield_and_drain(
asyncio.to_thread(loop.runtime_for_session, session)
)
async def _set_session_model_preset(
loop: AgentLoop,
session_key: str,
name: str,
) -> LLMRuntime:
if _has_native_coroutine_method(loop, "set_session_model_preset_async"):
return await loop.set_session_model_preset_async(session_key, name)
return await shield_and_drain(
asyncio.to_thread(loop.set_session_model_preset, session_key, name)
)
async def cmd_stop(ctx: CommandContext) -> OutboundMessage: async def cmd_stop(ctx: CommandContext) -> OutboundMessage:
"""Cancel all active tasks and subagents for the session.""" """Cancel all active tasks and subagents for the session."""
loop = ctx.loop loop = ctx.loop
@@ -258,16 +307,16 @@ async def cmd_restart(ctx: CommandContext) -> OutboundMessage:
async def cmd_status(ctx: CommandContext) -> OutboundMessage: async def cmd_status(ctx: CommandContext) -> OutboundMessage:
"""Build an outbound status message for a session.""" """Build an outbound status message for a session."""
loop = ctx.loop loop = ctx.loop
session = ctx.session or loop.sessions.get_or_create(ctx.key) session = ctx.session or await _get_or_create_session(loop, ctx.key)
runtime = ctx.runtime or loop.runtime_for_session(session) runtime = ctx.runtime or await _runtime_for_session(loop, session)
ctx_est = 0 ctx_est = 0
with suppress(Exception): with suppress(Exception):
ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens( ctx_est, _ = loop.consolidator.estimate_session_prompt_tokens(
session, session,
runtime=runtime, runtime=runtime,
) )
last_usage = LLMUsage.from_dict(session.metadata.get("_last_usage"))
if ctx_est <= 0: if ctx_est <= 0:
last_usage = loop._last_usage # pyright: ignore[reportPrivateUsage]
ctx_est = last_usage.input_tokens if last_usage is not None else 0 ctx_est = last_usage.input_tokens if last_usage is not None else 0
# Fetch web search provider usage (best-effort, never blocks the response) # Fetch web search provider usage (best-effort, never blocks the response)
@@ -290,7 +339,7 @@ async def cmd_status(ctx: CommandContext) -> OutboundMessage:
chat_id=ctx.msg.chat_id, chat_id=ctx.msg.chat_id,
content=build_status_content( content=build_status_content(
version=__version__, model=runtime.model, version=__version__, model=runtime.model,
start_time=loop._start_time, last_usage=last_usage, # pyright: ignore[reportPrivateUsage] start_time=loop._start_time, last_usage=loop._last_usage, # pyright: ignore[reportPrivateUsage]
context_window_tokens=runtime.context_window_tokens, context_window_tokens=runtime.context_window_tokens,
session_msg_count=len(session.get_history(max_messages=0)), session_msg_count=len(session.get_history(max_messages=0)),
context_tokens_estimate=ctx_est, context_tokens_estimate=ctx_est,
@@ -307,29 +356,32 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
loop = ctx.loop loop = ctx.loop
await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage] await loop._cancel_active_tasks(ctx.key) # pyright: ignore[reportPrivateUsage]
loop.discard_session_file_state(ctx.key) loop.discard_session_file_state(ctx.key)
session = ctx.session or loop.sessions.get_or_create(ctx.key) session = ctx.session or await _get_or_create_session(loop, ctx.key)
snapshot = list(session.messages) snapshot = list(session.messages)
archive_snapshot = None archive_snapshot = None
runtime = None runtime = None
if session.last_archived < len(snapshot): if session.last_consolidated < len(snapshot):
runtime = ctx.runtime or loop.runtime_for_session(session) runtime = ctx.runtime or await _runtime_for_session(loop, session)
archive_snapshot = replace( archive_snapshot = replace(
session, session,
messages=snapshot, messages=snapshot,
metadata=dict(session.metadata), metadata=dict(session.metadata),
provider_state=None, provider_state=None,
) )
session.clear() async def reset_and_schedule_archive() -> None:
loop.sessions.save(session) session.clear()
loop.sessions.invalidate(session.key) await _save_session(loop, session)
if archive_snapshot is not None and runtime is not None: loop.sessions.invalidate(session.key)
loop.schedule_background( if archive_snapshot is not None and runtime is not None:
loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType] loop.schedule_background(
archive_snapshot, loop.consolidator.archive_session( # pyright: ignore[reportUnknownMemberType]
archive_end=len(snapshot), archive_snapshot,
runtime=runtime, archive_end=len(snapshot),
runtime=runtime,
)
) )
)
await shield_and_drain(reset_and_schedule_archive())
return OutboundMessage( return OutboundMessage(
channel=ctx.msg.channel, chat_id=ctx.msg.chat_id, channel=ctx.msg.channel, chat_id=ctx.msg.chat_id,
content="New session started.", content="New session started.",
@@ -378,7 +430,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"} metadata = {**dict(ctx.msg.metadata or {}), "render_as": "text"}
if not args: if not args:
session = ctx.session or loop.sessions.get_or_create(ctx.key) session = ctx.session or await _get_or_create_session(loop, ctx.key)
return OutboundMessage( return OutboundMessage(
channel=ctx.msg.channel, channel=ctx.msg.channel,
chat_id=ctx.msg.chat_id, chat_id=ctx.msg.chat_id,
@@ -388,7 +440,7 @@ async def cmd_model(ctx: CommandContext) -> OutboundMessage:
name = args name = args
try: try:
runtime = loop.set_session_model_preset(ctx.key, name) runtime = await _set_session_model_preset(loop, ctx.key, name)
except (KeyError, ValueError) as exc: except (KeyError, ValueError) as exc:
names = _model_preset_names(loop) names = _model_preset_names(loop)
return OutboundMessage( return OutboundMessage(
@@ -849,7 +901,7 @@ async def cmd_history(ctx: CommandContext) -> OutboundMessage:
metadata=dict(ctx.msg.metadata or {}), metadata=dict(ctx.msg.metadata or {}),
) )
session = ctx.session or ctx.loop.sessions.get_or_create(ctx.key) session = ctx.session or await _get_or_create_session(ctx.loop, ctx.key)
history = session.get_history(max_messages=0, include_runtime_context=False) history = session.get_history(max_messages=0, include_runtime_context=False)
visible = [_format_history_message(m) for m in history] visible = [_format_history_message(m) for m in history]
visible = [m for m in visible if m is not None] visible = [m for m in visible if m is not None]
+9 -1
View File
@@ -128,7 +128,8 @@ class AgentDefaults(Base):
temperature: float = 0.1 temperature: float = 0.1
fallback_models: list[FallbackCandidate] = Field(default_factory=list) fallback_models: list[FallbackCandidate] = Field(default_factory=list)
max_tool_iterations: int = 200 max_tool_iterations: int = 200
max_concurrent_subagents: int = Field(default=4, ge=1) max_concurrent_subagents: int = Field(default=1, ge=1)
fail_on_tool_error: bool = True
max_tool_result_chars: int = 16_000 max_tool_result_chars: int = 16_000
provider_retry_mode: Literal["standard", "persistent"] = "standard" provider_retry_mode: Literal["standard", "persistent"] = "standard"
tool_hint_max_length: int = Field( tool_hint_max_length: int = Field(
@@ -155,6 +156,13 @@ class AgentDefaults(Base):
default=60, default=60,
ge=0, ge=0,
) # Minimum interval in seconds between scans for idle sessions ) # Minimum interval in seconds between scans for idle sessions
consolidation_ratio: float = Field(
default=0.5,
ge=0.1,
le=0.95,
validation_alias=AliasChoices("consolidationRatio"),
serialization_alias="consolidationRatio",
) # Consolidation target ratio (0.5 = 50% of budget retained after compression)
dream: DreamConfig = Field(default_factory=DreamConfig) dream: DreamConfig = Field(default_factory=DreamConfig)
@model_validator(mode="before") @model_validator(mode="before")
+10
View File
@@ -8,6 +8,7 @@ import time
import uuid import uuid
from typing import TYPE_CHECKING, Any, Protocol from typing import TYPE_CHECKING, Any, Protocol
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
from nanobot.agent.tools.cron import CronTool from nanobot.agent.tools.cron import CronTool
from nanobot.bus.events import InboundMessage, OutboundMessage from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.cron.session_delivery import origin_delivery_context from nanobot.cron.session_delivery import origin_delivery_context
@@ -127,6 +128,15 @@ async def run_bound_cron_job(
session_key_override=session_key, session_key_override=session_key,
) )
) )
except AutomationTurnAcceptedCancellation:
cron.write_run_record(
run_id,
{
**run_record_base,
"status": "accepted",
},
)
raise
except (Exception, asyncio.CancelledError) as exc: except (Exception, asyncio.CancelledError) as exc:
error_text = str(exc) or exc.__class__.__name__ error_text = str(exc) or exc.__class__.__name__
cron.write_run_record( cron.write_run_record(
+233 -108
View File
@@ -11,11 +11,12 @@ from dataclasses import asdict
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
from types import EllipsisType from types import EllipsisType
from typing import Any, Callable, Coroutine, Literal from typing import Any, Callable, Coroutine, Literal, TypeVar
from filelock import FileLock from filelock import FileLock
from loguru import logger from loguru import logger
from nanobot.agent.automation_turns import AutomationTurnAcceptedCancellation
from nanobot.cron.session_turns import is_bound_cron_job from nanobot.cron.session_turns import is_bound_cron_job
from nanobot.cron.types import ( from nanobot.cron.types import (
CronJob, CronJob,
@@ -25,11 +26,14 @@ from nanobot.cron.types import (
CronSchedule, CronSchedule,
CronStore, CronStore,
) )
from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.run_records import ( from nanobot.utils.run_records import (
write_run_record as write_automation_run_record, write_run_record as write_automation_run_record,
) )
_FILE_LOCK_TIMEOUT_SECONDS = 5
_T = TypeVar("_T")
class CronJobSkippedError(Exception): class CronJobSkippedError(Exception):
"""Raised by cron callbacks when a job was intentionally skipped.""" """Raised by cron callbacks when a job was intentionally skipped."""
@@ -116,21 +120,8 @@ def _disable_malformed_legacy_job(job: CronJob) -> None:
logger.warning("Cron: disabled malformed legacy job '{}' ({}): {}", job.name, job.id, reason) logger.warning("Cron: disabled malformed legacy job '{}' ({}): {}", job.name, job.id, reason)
def _persistable_origin_metadata(metadata: dict[str, Any]) -> dict[str, Any]:
"""Return a detached JSON-safe routing snapshot for a cron payload."""
snapshot: dict[str, Any] = {}
for key, value in metadata.items():
if key == RUNTIME_CONTEXT_INPUT_META:
continue
try:
snapshot[key] = json.loads(json.dumps(value, ensure_ascii=False, allow_nan=False))
except (TypeError, ValueError, RecursionError):
continue
return snapshot
def _normalize_agent_turn_job(job: CronJob) -> bool: def _normalize_agent_turn_job(job: CronJob) -> bool:
"""Make routing metadata persistable and migrate legacy user cron payloads. """Migrate legacy user cron payloads into session-bound payloads.
Pre-bound user cron jobs stored their delivery target in ``channel``/``to``. Pre-bound user cron jobs stored their delivery target in ``channel``/``to``.
Normal user-created legacy jobs always have those fields; if they are Normal user-created legacy jobs always have those fields; if they are
@@ -138,12 +129,8 @@ def _normalize_agent_turn_job(job: CronJob) -> bool:
a runtime legacy execution path. a runtime legacy execution path.
""" """
payload = job.payload payload = job.payload
origin_metadata = _persistable_origin_metadata(payload.origin_metadata)
changed = origin_metadata != payload.origin_metadata
payload.origin_metadata = origin_metadata
if payload.kind != "agent_turn" or not _has_legacy_delivery_context(payload): if payload.kind != "agent_turn" or not _has_legacy_delivery_context(payload):
return changed return False
if not payload.channel or not payload.to: if not payload.channel or not payload.to:
_disable_malformed_legacy_job(job) _disable_malformed_legacy_job(job)
@@ -153,7 +140,7 @@ def _normalize_agent_turn_job(job: CronJob) -> bool:
payload.origin_channel = payload.origin_channel or payload.channel payload.origin_channel = payload.origin_channel or payload.channel
payload.origin_chat_id = payload.origin_chat_id or payload.to payload.origin_chat_id = payload.origin_chat_id or payload.to
if not payload.origin_metadata: if not payload.origin_metadata:
payload.origin_metadata = _persistable_origin_metadata(payload.channel_meta or {}) payload.origin_metadata = dict(payload.channel_meta or {})
payload.deliver = False payload.deliver = False
payload.channel = None payload.channel = None
@@ -182,10 +169,16 @@ class CronService:
self.store_path = store_path self.store_path = store_path
self._action_path = store_path.parent / "action.jsonl" self._action_path = store_path.parent / "action.jsonl"
self._run_records_dir = store_path.parent / "runs" self._run_records_dir = store_path.parent / "runs"
self._lock = FileLock(str(self._action_path.parent) + ".lock") self._lock = FileLock(
str(self._action_path.parent) + ".lock",
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
)
self.on_job = on_job self.on_job = on_job
self._store: CronStore | None = None self._store: CronStore | None = None
self._timer_task: asyncio.Task[None] | None = None self._timer_task: asyncio.Task[None] | None = None
self._operation_lock = asyncio.Lock()
self._claimed_job_ids: set[str] = set()
self._event_loop: asyncio.AbstractEventLoop | None = None
self._running = False self._running = False
self._active_executions = 0 self._active_executions = 0
self._store_dirty = False self._store_dirty = False
@@ -469,25 +462,58 @@ class CronService:
"""Write an internal audit record for one cron execution.""" """Write an internal audit record for one cron execution."""
write_automation_run_record(self._run_records_dir, run_id, record) write_automation_run_record(self._run_records_dir, run_id, record)
async def start(self) -> None: async def run_sync(
"""Start the cron service.""" self,
self._running = True operation: Callable[..., _T],
loaded = self._load_store() /,
if loaded is None: *args: Any,
# Store file existed but was corrupt and has been preserved with **kwargs: Any,
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with ) -> _T:
# an empty store; that would call ``_save_store`` and overwrite """Serialize a complete cron transaction in a worker thread.
# the now-renamed (but still recoverable) data with [].
self._running = False A running thread cannot be cancelled safely. Keep the transaction lock
raise RuntimeError( until it exits so cancellation is never reported while that worker can
f"cron store at {self.store_path} is corrupt and was preserved; " still mutate cron state behind a later operation.
"refusing to start with an empty job list. " """
"Inspect the .corrupt-<ts> backup and restore manually." async with self._operation_lock:
return await shield_and_drain(
asyncio.to_thread(operation, *args, **kwargs)
) )
self._recompute_next_runs()
self._save_store() async def start(self) -> None:
self._arm_timer() """Start the cron service and settle accepted work before cancellation."""
logger.info("Cron service started with {} jobs", len(self._store.jobs if self._store else []))
async def settle_start() -> None:
self._event_loop = asyncio.get_running_loop()
self._running = True
try:
async with self._operation_lock:
loaded = await asyncio.to_thread(self._load_store)
if loaded is None:
# Store file existed but was corrupt and has been preserved with
# a ``.corrupt-<ts>`` suffix. Bail out instead of starting with
# an empty store; that would call ``_save_store`` and overwrite
# the now-renamed (but still recoverable) data with [].
raise RuntimeError(
f"cron store at {self.store_path} is corrupt and was preserved; "
"refusing to start with an empty job list. "
"Inspect the .corrupt-<ts> backup and restore manually."
)
self._recompute_next_runs()
await asyncio.to_thread(self._save_store)
self._arm_timer()
logger.info(
"Cron service started with {} jobs",
len(self._store.jobs if self._store else []),
)
except BaseException:
# A failed start must not retain ownership without a timer. Caller
# cancellation is shielded until this composite either reaches the
# fully started state above or rolls back here.
self.stop()
raise
await shield_and_drain(settle_start())
def stop(self) -> None: def stop(self) -> None:
"""Stop the cron service.""" """Stop the cron service."""
@@ -515,8 +541,22 @@ class CronService:
if j.enabled and j.state.next_run_at_ms] if j.enabled and j.state.next_run_at_ms]
return min(times) if times else None return min(times) if times else None
def _request_timer_rearm(self) -> None:
"""Re-arm on the owning event loop, including from persistence workers."""
if not self._running:
return
loop = self._event_loop
try:
current_loop = asyncio.get_running_loop()
except RuntimeError:
current_loop = None
if current_loop is loop:
self._arm_timer()
elif loop is not None and loop.is_running():
loop.call_soon_threadsafe(self._arm_timer)
def _arm_timer(self) -> None: def _arm_timer(self) -> None:
"""Schedule the next timer tick.""" """Schedule the next timer tick on the owning event loop."""
if self._timer_task: if self._timer_task:
self._timer_task.cancel() self._timer_task.cancel()
@@ -538,7 +578,7 @@ class CronService:
self._timer_task = asyncio.create_task(tick()) self._timer_task = asyncio.create_task(tick())
async def _on_timer(self) -> None: async def _on_timer(self) -> None:
"""Handle timer tick - run due jobs.""" """Run due jobs while keeping persistence transactions serialized."""
reload_store = self._active_executions == 0 reload_store = self._active_executions == 0
self._active_executions += 1 self._active_executions += 1
try: try:
@@ -546,11 +586,17 @@ class CronService:
# to persist their advanced schedule. Persist that exact snapshot # to persist their advanced schedule. Persist that exact snapshot
# before reloading or executing anything else; otherwise the older # before reloading or executing anything else; otherwise the older
# disk state can replay the same job. # disk state can replay the same job.
if self._store_dirty: async with self._operation_lock:
self._save_store() if self._store_dirty:
return await shield_and_drain(asyncio.to_thread(self._save_store))
return
store = self._load_store(reload_during_execution=reload_store) store = await shield_and_drain(
asyncio.to_thread(
self._load_store,
reload_during_execution=reload_store,
)
)
# If a hot reload found a corrupt store on disk, ``self._store`` # If a hot reload found a corrupt store on disk, ``self._store``
# may still hold the previous, known-good in-memory snapshot. # may still hold the previous, known-good in-memory snapshot.
if store is None: if store is None:
@@ -565,7 +611,8 @@ class CronService:
for job in due_jobs: for job in due_jobs:
await self._execute_job(job) await self._execute_job(job)
self._save_store() async with self._operation_lock:
await shield_and_drain(asyncio.to_thread(self._save_store))
except Exception: except Exception:
# A load/persist failure must not kill the scheduler: keep the # A load/persist failure must not kill the scheduler: keep the
# in-memory store and retry on the next tick. This mirrors the # in-memory store and retry on the next tick. This mirrors the
@@ -582,58 +629,124 @@ class CronService:
# single bad tick cannot silently stop all future jobs. # single bad tick cannot silently stop all future jobs.
self._arm_timer() self._arm_timer()
async def _execute_job(self, job: CronJob) -> None: async def _claim_job(self, job_id: str) -> bool:
"""Execute a single job.""" """Claim one job without serializing callbacks for different jobs."""
async with self._operation_lock:
if job_id in self._claimed_job_ids:
return False
self._claimed_job_ids.add(job_id)
return True
async def _release_job_claim(self, job_id: str) -> None:
async with self._operation_lock:
self._claimed_job_ids.discard(job_id)
async def _settle_job_execution(
self,
job: CronJob,
*,
start_ms: int,
status: Literal["ok", "error", "skipped"],
error: str | None,
persist: bool = False,
) -> None:
end_ms = _now_ms()
async with self._operation_lock:
job.state.last_status = status
job.state.last_error = error
job.state.last_run_at_ms = start_ms
job.updated_at_ms = end_ms
job.state.run_history.append(CronRunRecord(
run_at_ms=start_ms,
status=status,
duration_ms=end_ms - start_ms,
error=error,
))
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
if job.schedule.kind == "at":
if job.delete_after_run:
store = await shield_and_drain(
asyncio.to_thread(self._require_store)
)
store.jobs = [item for item in store.jobs if item.id != job.id]
else:
job.enabled = False
job.state.next_run_at_ms = None
else:
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
if persist:
await shield_and_drain(asyncio.to_thread(self._save_store))
@staticmethod
async def _drain_settlement_on_cancellation(settlement: asyncio.Task[None]) -> None:
"""Finish a short durable settlement despite repeated cancellation."""
while not settlement.done():
try:
await asyncio.shield(settlement)
except asyncio.CancelledError:
continue
settlement.result()
async def _execute_job(self, job: CronJob) -> bool:
"""Execute a claimed job and serialize its in-memory settlement."""
if not await self._claim_job(job.id):
logger.info("Cron: job '{}' ({}) is already running", job.name, job.id)
return False
start_ms = _now_ms() start_ms = _now_ms()
logger.info("Cron: executing job '{}' ({})", job.name, job.id) logger.info("Cron: executing job '{}' ({})", job.name, job.id)
status: Literal["ok", "error", "skipped"]
error: str | None
accepted_cancellation: AutomationTurnAcceptedCancellation | None = None
try: try:
if self.on_job: try:
await self.on_job(job) if self.on_job:
await self.on_job(job)
status = "ok"
error = None
logger.info("Cron: job '{}' completed", job.name)
except AutomationTurnAcceptedCancellation as exc:
# The agent owns this turn now. Advance and persist the schedule
# before allowing shutdown cancellation to unwind the timer.
status = "ok"
error = None
accepted_cancellation = exc
logger.info("Cron: job '{}' was accepted before cancellation", job.name)
except CronJobSkippedError as exc:
status = "skipped"
error = str(exc) or None
logger.warning("Cron: job '{}' skipped: {}", job.name, error or "")
except asyncio.CancelledError as exc:
current = asyncio.current_task()
if current is not None and current.cancelling():
raise
status = "error"
error = str(exc) or exc.__class__.__name__
logger.exception("Cron: job '{}' was cancelled", job.name)
except Exception as exc:
status = "error"
error = str(exc)
logger.exception("Cron: job '{}' failed", job.name)
job.state.last_status = "ok" settlement = asyncio.create_task(
job.state.last_error = None self._settle_job_execution(
logger.info("Cron: job '{}' completed", job.name) job,
start_ms=start_ms,
except CronJobSkippedError as e: status=status,
job.state.last_status = "skipped" error=error,
job.state.last_error = str(e) or None persist=accepted_cancellation is not None,
logger.warning("Cron: job '{}' skipped: {}", job.name, job.state.last_error or "") )
except asyncio.CancelledError as e: )
current = asyncio.current_task() if accepted_cancellation is not None:
if current is not None and current.cancelling(): await self._drain_settlement_on_cancellation(settlement)
raise raise accepted_cancellation
job.state.last_status = "error" await settlement
job.state.last_error = str(e) or e.__class__.__name__ return True
logger.exception("Cron: job '{}' was cancelled", job.name) finally:
except Exception as e: await self._release_job_claim(job.id)
job.state.last_status = "error"
job.state.last_error = str(e)
logger.exception("Cron: job '{}' failed", job.name)
end_ms = _now_ms()
job.state.last_run_at_ms = start_ms
job.updated_at_ms = end_ms
job.state.run_history.append(CronRunRecord(
run_at_ms=start_ms,
status=job.state.last_status,
duration_ms=end_ms - start_ms,
error=job.state.last_error,
))
job.state.run_history = job.state.run_history[-self._MAX_RUN_HISTORY:]
# Handle one-shot jobs
if job.schedule.kind == "at":
if job.delete_after_run:
store = self._require_store()
store.jobs = [item for item in store.jobs if item.id != job.id]
else:
job.enabled = False
job.state.next_run_at_ms = None
else:
# Compute next run
job.state.next_run_at_ms = _compute_next_run(job.schedule, _now_ms())
def _append_action( def _append_action(
self, self,
@@ -715,7 +828,7 @@ class CronService:
store = self._require_store() store = self._require_store()
store.jobs.append(job) store.jobs.append(job)
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
else: else:
self._append_action("add", asdict(job)) self._append_action("add", asdict(job))
@@ -732,7 +845,7 @@ class CronService:
store.jobs = [j for j in store.jobs if j.id != job.id] store.jobs = [j for j in store.jobs if j.id != job.id]
store.jobs.append(job) store.jobs.append(job)
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
logger.info("Cron: registered system job '{}' ({})", job.name, job.id) logger.info("Cron: registered system job '{}' ({})", job.name, job.id)
return job return job
@@ -744,7 +857,7 @@ class CronService:
removed = len(store.jobs) < before removed = len(store.jobs) < before
if removed: if removed:
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
logger.info("Cron: removed system job {}", job_id) logger.info("Cron: removed system job {}", job_id)
return removed return removed
@@ -765,7 +878,7 @@ class CronService:
if removed: if removed:
if self._should_persist_store(): if self._should_persist_store():
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
else: else:
self._append_action("del", {"job_id": job_id}) self._append_action("del", {"job_id": job_id})
logger.info("Cron: removed job {}", job_id) logger.info("Cron: removed job {}", job_id)
@@ -787,7 +900,7 @@ class CronService:
job.state.next_run_at_ms = None job.state.next_run_at_ms = None
if self._should_persist_store(): if self._should_persist_store():
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
else: else:
self._append_action("update", asdict(job)) self._append_action("update", asdict(job))
return job return job
@@ -843,7 +956,7 @@ class CronService:
if self._should_persist_store(): if self._should_persist_store():
self._save_store() self._save_store()
self._arm_timer() self._request_timer_rearm()
else: else:
self._append_action("update", asdict(job)) self._append_action("update", asdict(job))
@@ -858,19 +971,31 @@ class CronService:
# A manual run is another side-effecting entrypoint. Do not start # A manual run is another side-effecting entrypoint. Do not start
# it while the result of a previous timer execution is still only # it while the result of a previous timer execution is still only
# in memory. # in memory.
if self._store_dirty: async with self._operation_lock:
self._save_store() if self._store_dirty:
store = self._require_store(reload_during_execution=reload_store) await shield_and_drain(asyncio.to_thread(self._save_store))
store = await shield_and_drain(
asyncio.to_thread(
self._require_store,
reload_during_execution=reload_store,
)
)
for job in store.jobs: for job in store.jobs:
if job.id == job_id: if job.id == job_id:
if self._is_unbound_agent_job(job): if self._is_unbound_agent_job(job):
self._enforce_agent_binding(job) async with self._operation_lock:
self._save_store() self._enforce_agent_binding(job)
await shield_and_drain(
asyncio.to_thread(self._save_store)
)
return False return False
if not force and not job.enabled: if not force and not job.enabled:
return False return False
await self._execute_job(job) executed = await self._execute_job(job)
self._save_store() if not executed:
return False
async with self._operation_lock:
await shield_and_drain(asyncio.to_thread(self._save_store))
return True return True
return False return False
finally: finally:
+1 -52
View File
@@ -6,7 +6,6 @@ from __future__ import annotations
import asyncio import asyncio
import hashlib import hashlib
import http.client
import json import json
import os import os
import subprocess import subprocess
@@ -39,33 +38,6 @@ GatewayLaunchMode = Literal["foreground", "background", "unknown"]
GatewayLifetime = Literal["explicit", "on_demand"] GatewayLifetime = Literal["explicit", "on_demand"]
def _gateway_health_ready(host: str, port: int, *, timeout_s: float = 0.4) -> bool:
"""Read readiness from the management listener without using proxy settings."""
connect_host = "127.0.0.1" if host in {"", "0.0.0.0"} else "::1" if host == "::" else host
connection = http.client.HTTPConnection(connect_host, port, timeout=timeout_s)
try:
connection.request("GET", "/health")
response = connection.getresponse()
body = response.read(1024)
except (OSError, http.client.HTTPException, TimeoutError):
return False
finally:
connection.close()
if response.status != 200:
return False
try:
raw_payload = cast(object, json.loads(body.decode("utf-8")))
except (UnicodeDecodeError, json.JSONDecodeError):
return False
if not isinstance(raw_payload, dict):
return False
payload = cast(dict[str, object], raw_payload)
return (
payload.get("status") == "ok"
and payload.get("ready") is not False
)
def _default_config_path() -> Path: def _default_config_path() -> Path:
return (Path.home() / ".nanobot" / "config.json").resolve(strict=False) return (Path.home() / ".nanobot" / "config.json").resolve(strict=False)
@@ -77,7 +49,6 @@ class GatewayStatus(ProcessStatus):
launch_mode: GatewayLaunchMode = "unknown" launch_mode: GatewayLaunchMode = "unknown"
lifetime: GatewayLifetime = "explicit" lifetime: GatewayLifetime = "explicit"
clients: int = 0 clients: int = 0
ready: bool | None = None
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -288,18 +259,6 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
raw_mode if raw_mode in {"foreground", "background"} else "unknown" raw_mode if raw_mode in {"foreground", "background"} else "unknown"
) )
lease = GatewayClientLease(self, kind="gateway-status").snapshot() lease = GatewayClientLease(self, kind="gateway-status").snapshot()
ready: bool | None = None
health_host = state.get("health_host") if state else None
if (
process.running
and process.pid != os.getpid()
and isinstance(health_host, str)
and process.port is not None
):
ready = _gateway_health_ready(health_host, process.port)
status_reason = process.reason
if ready is False and reason is None and status_reason == "running":
status_reason = "websocket_unavailable"
return GatewayStatus( return GatewayStatus(
running=process.running, running=process.running,
pid=process.pid, pid=process.pid,
@@ -308,22 +267,12 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
started_at=process.started_at, started_at=process.started_at,
port=process.port, port=process.port,
command=process.command, command=process.command,
reason=status_reason, reason=process.reason,
launch_mode=launch_mode, launch_mode=launch_mode,
lifetime="on_demand" if lease.auto_stop else "explicit", lifetime="on_demand" if lease.auto_stop else "explicit",
clients=lease.clients, clients=lease.clients,
ready=ready,
) )
def publish_health_host(self, host: str) -> None:
"""Record the management bind host for out-of-process readiness diagnostics."""
with self._lifecycle_lock():
state = self._read_state()
if not state or not self._record_matches_process(state, os.getpid()):
return
state["health_host"] = host
self._write_state(state)
@contextmanager @contextmanager
def foreground_instance(self, options: ProcessStartOptions) -> Generator[None]: def foreground_instance(self, options: ProcessStartOptions) -> Generator[None]:
"""Publish this foreground gateway while it is available to local clients.""" """Publish this foreground gateway while it is available to local clients."""
+3 -73
View File
@@ -252,13 +252,10 @@ class ProviderCallContext:
The regular ``chat`` contract stays provider-agnostic. Responses-capable The regular ``chat`` contract stays provider-agnostic. Responses-capable
providers consume this context through the opt-in ``chat_with_context`` providers consume this context through the opt-in ``chat_with_context``
hooks, while every other provider inherits the context-free delegation. hooks, while every other provider inherits the context-free delegation.
``session_id`` gives providers a stable conversation-scoped routing key
without exposing that identity in the public message transcript.
""" """
conversation_state: ProviderConversationState | None = field(default=None, repr=False) conversation_state: ProviderConversationState | None = field(default=None, repr=False)
context_window_tokens: int | None = None context_window_tokens: int | None = None
session_id: str | None = field(default=None, repr=False)
@dataclass(frozen=True, slots=True) @dataclass(frozen=True, slots=True)
@@ -603,6 +600,8 @@ _SYNTHETIC_USER_CONTENT = "(conversation continued)"
class LLMProvider(ABC): class LLMProvider(ABC):
"""Base class for LLM providers.""" """Base class for LLM providers."""
supports_progress_deltas = False
_CHAT_RETRY_DELAYS = (1, 2, 4) _CHAT_RETRY_DELAYS = (1, 2, 4)
_PERSISTENT_MAX_DELAY = 60 _PERSISTENT_MAX_DELAY = 60
_PERSISTENT_IDENTICAL_ERROR_LIMIT = 10 _PERSISTENT_IDENTICAL_ERROR_LIMIT = 10
@@ -1029,20 +1028,6 @@ class LLMProvider(ABC):
# Unknown 429 defaults to WAIT+retry. # Unknown 429 defaults to WAIT+retry.
return True return True
@staticmethod
def _content_as_blocks(content: Any) -> list[dict[str, Any]]:
"""Convert message content to blocks so mixed user content can be merged."""
if isinstance(content, list):
return [
dict(cast(dict[str, Any], item))
if isinstance(item, dict)
else {"type": "text", "text": str(item)}
for item in cast(list[object], content)
]
if content is None:
return []
return [{"type": "text", "text": str(content)}]
@staticmethod @staticmethod
def _enforce_role_alternation(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: def _enforce_role_alternation(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Merge consecutive same-role messages and drop trailing assistant messages. """Merge consecutive same-role messages and drop trailing assistant messages.
@@ -1077,13 +1062,6 @@ class LLMProvider(ABC):
curr_content = msg.get("content") or "" curr_content = msg.get("content") or ""
if isinstance(prev_content, str) and isinstance(curr_content, str): if isinstance(prev_content, str) and isinstance(curr_content, str):
prev["content"] = (prev_content + "\n\n" + curr_content).strip() prev["content"] = (prev_content + "\n\n" + curr_content).strip()
elif role == "user":
combined = dict(msg)
combined["content"] = [
*LLMProvider._content_as_blocks(prev_content),
*LLMProvider._content_as_blocks(curr_content),
]
merged[-1] = combined
else: else:
merged[-1] = dict(msg) merged[-1] = dict(msg)
else: else:
@@ -1281,53 +1259,6 @@ class LLMProvider(ABC):
"""Call chat_stream() and convert unexpected exceptions to error responses.""" """Call chat_stream() and convert unexpected exceptions to error responses."""
started_at_ms = time.time_ns() // 1_000_000 started_at_ms = time.time_ns() // 1_000_000
started_at_ns = time.monotonic_ns() started_at_ns = time.monotonic_ns()
first_output_at_ns: int | None = None
def _mark_output(delta: str) -> None:
nonlocal first_output_at_ns
if delta and first_output_at_ns is None:
first_output_at_ns = time.monotonic_ns()
if self._llm_call_observer is not None:
content_callback = kwargs.get("on_content_delta")
if callable(content_callback):
typed_content_callback = cast(
Callable[[str], Awaitable[None]],
content_callback,
)
async def _timed_content_delta(delta: str) -> None:
_mark_output(delta)
await typed_content_callback(delta)
kwargs["on_content_delta"] = _timed_content_delta
thinking_callback = kwargs.get("on_thinking_delta")
if callable(thinking_callback):
typed_thinking_callback = cast(
Callable[[str], Awaitable[None]],
thinking_callback,
)
async def _timed_thinking_delta(delta: str) -> None:
_mark_output(delta)
await typed_thinking_callback(delta)
kwargs["on_thinking_delta"] = _timed_thinking_delta
def _attach_stream_timing(response: LLMResponse) -> LLMResponse:
if first_output_at_ns is None:
return response
finished_at_ns = time.monotonic_ns()
if response.ttft_ms is None:
response.ttft_ms = max(0, round((first_output_at_ns - started_at_ns) / 1_000_000))
if response.generation_ms is None:
response.generation_ms = max(
1,
round((finished_at_ns - first_output_at_ns) / 1_000_000),
)
return response
try: try:
provider_context = kwargs.pop("provider_context", None) provider_context = kwargs.pop("provider_context", None)
if isinstance(provider_context, ProviderCallContext): if isinstance(provider_context, ProviderCallContext):
@@ -1353,7 +1284,7 @@ class LLMProvider(ABC):
except Exception as exc: except Exception as exc:
response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error") response = LLMResponse(content=f"Error calling LLM: {exc}", finish_reason="error")
return self._observe_llm_call( return self._observe_llm_call(
_attach_stream_timing(response), response,
kwargs, kwargs,
started_at_ms=started_at_ms, started_at_ms=started_at_ms,
started_at_ns=started_at_ns, started_at_ns=started_at_ns,
@@ -1662,7 +1593,6 @@ class LLMProvider(ABC):
context_window_tokens=( context_window_tokens=(
provider_context.context_window_tokens provider_context.context_window_tokens
), ),
session_id=provider_context.session_id,
) )
if stripped is not None or stripped_context is not None: if stripped is not None or stripped_context is not None:
logger.warning( logger.warning(
+3 -50
View File
@@ -11,7 +11,6 @@ from nanobot.providers.base import (
ProviderCallContext, ProviderCallContext,
ProviderConversationState, ProviderConversationState,
) )
from nanobot.utils.helpers import estimate_prompt_tokens_chain
_PROVIDER_STATE_OUTPUT_META = "provider_state_output" _PROVIDER_STATE_OUTPUT_META = "provider_state_output"
_PROVIDER_STATE_BOUNDARY_META = "provider_state_boundary" _PROVIDER_STATE_BOUNDARY_META = "provider_state_boundary"
@@ -43,11 +42,9 @@ class ProviderConversationStateController:
model: str | None, model: str | None,
messages: list[dict[str, Any]], messages: list[dict[str, Any]],
state: ProviderConversationState | None = None, state: ProviderConversationState | None = None,
session_id: str | None = None,
) -> None: ) -> None:
self._provider = provider self._provider = provider
self._model = model self._model = model
self._session_id = session_id
self._state = ( self._state = (
state state
if state is not None if state is not None
@@ -63,43 +60,9 @@ class ProviderConversationStateController:
context_window_tokens: int | None, context_window_tokens: int | None,
) -> ProviderCallContext | None: ) -> ProviderCallContext | None:
"""Return typed provider context for a request that does not resume state.""" """Return typed provider context for a request that does not resume state."""
if context_window_tokens is None and self._session_id is None: if context_window_tokens is None:
return None return None
return ProviderCallContext( return ProviderCallContext(context_window_tokens=context_window_tokens)
context_window_tokens=context_window_tokens,
session_id=self._session_id,
)
def estimate_request_context_tokens(
self,
messages: list[dict[str, Any]],
*,
model_messages: list[dict[str, Any]] | None = None,
supplemental_messages: list[dict[str, Any]] | None = None,
tool_definitions: list[dict[str, Any]] | None = None,
) -> int | None:
"""Estimate resumed state plus the pending delta for the next request."""
state = self.checkpoint(messages, model_messages=model_messages)
if state is None:
return None
context_tokens = state.payload.get("context_tokens")
if (
isinstance(context_tokens, bool)
or not isinstance(context_tokens, int)
or context_tokens < 0
):
return None
pending_messages = [
*state.pending_messages,
*(supplemental_messages or []),
]
delta_tokens, _ = estimate_prompt_tokens_chain(
self._provider,
self._model,
pending_messages,
tool_definitions,
)
return context_tokens + max(0, delta_tokens)
def prepare_request( def prepare_request(
self, self,
@@ -108,20 +71,11 @@ class ProviderConversationStateController:
context_window_tokens: int | None, context_window_tokens: int | None,
model_messages: list[dict[str, Any]] | None = None, model_messages: list[dict[str, Any]] | None = None,
supplemental_messages: list[dict[str, Any]] | None = None, supplemental_messages: list[dict[str, Any]] | None = None,
resume_state: bool = True,
) -> ProviderCallContext | None: ) -> ProviderCallContext | None:
"""Build context for the next request and remember its durable delta. """Build typed context for the next request and remember its durable delta."""
``resume_state=False`` abandons opaque history when local request
fitting has produced a new independent model-facing context.
"""
independent_context = self.independent_request_context( independent_context = self.independent_request_context(
context_window_tokens=context_window_tokens, context_window_tokens=context_window_tokens,
) )
if not resume_state:
self._state = None
self._request_messages = []
return independent_context
if self._state is None: if self._state is None:
self._request_messages = [] self._request_messages = []
return independent_context return independent_context
@@ -158,7 +112,6 @@ class ProviderConversationStateController:
if independent_context is not None if independent_context is not None
else None else None
), ),
session_id=self._session_id,
) )
def observe_response( def observe_response(
+4 -2
View File
@@ -157,6 +157,10 @@ class FallbackProvider(LLMProvider):
super().set_llm_call_observer(observer) super().set_llm_call_observer(observer)
self._primary.set_llm_call_observer(observer) self._primary.set_llm_call_observer(observer)
@property
def supports_progress_deltas(self) -> bool:
return bool(getattr(self._primary, "supports_progress_deltas", False))
def can_resume_conversation_state( def can_resume_conversation_state(
self, self,
state: ProviderConversationState, state: ProviderConversationState,
@@ -182,7 +186,6 @@ class FallbackProvider(LLMProvider):
return ProviderCallContext( return ProviderCallContext(
conversation_state=provider_context.conversation_state, conversation_state=provider_context.conversation_state,
context_window_tokens=context_window_tokens, context_window_tokens=context_window_tokens,
session_id=provider_context.session_id,
) )
def _primary_available(self) -> bool: def _primary_available(self) -> bool:
@@ -538,7 +541,6 @@ class FallbackProvider(LLMProvider):
fallback_kwargs["provider_context"] = ProviderCallContext( fallback_kwargs["provider_context"] = ProviderCallContext(
conversation_state=state, conversation_state=state,
context_window_tokens=context_window_tokens, context_window_tokens=context_window_tokens,
session_id=provider_context.session_id,
) )
if fallback.reasoning_effort is None: if fallback.reasoning_effort is None:
fallback_kwargs.pop("reasoning_effort", None) fallback_kwargs.pop("reasoning_effort", None)
+4 -183
View File
@@ -5,7 +5,6 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
import hashlib
import os import os
import time import time
import webbrowser import webbrowser
@@ -18,12 +17,7 @@ from oauth_cli_kit.models import OAuthToken
from oauth_cli_kit.storage import FileTokenStorage from oauth_cli_kit.storage import FileTokenStorage
from nanobot.providers.base import LLMResponse, ProviderCallContext from nanobot.providers.base import LLMResponse, ProviderCallContext
from nanobot.providers.oauth_model_catalog import (
OAuthModelCatalog,
OAuthModelCatalogSnapshot,
)
from nanobot.providers.openai_compat_provider import OpenAICompatProvider from nanobot.providers.openai_compat_provider import OpenAICompatProvider
from nanobot.providers.registry import ProviderModelSpec, find_by_name
DEFAULT_GITHUB_DEVICE_CODE_URL = "https://github.com/login/device/code" DEFAULT_GITHUB_DEVICE_CODE_URL = "https://github.com/login/device/code"
DEFAULT_GITHUB_ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token" DEFAULT_GITHUB_ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token"
@@ -102,9 +96,7 @@ def login_github_copilot(
device_code = str(payload["device_code"]) device_code = str(payload["device_code"])
user_code = str(payload["user_code"]) user_code = str(payload["user_code"])
verify_url = str( verify_url = str(payload.get("verification_uri") or payload.get("verification_uri_complete") or "")
payload.get("verification_uri") or payload.get("verification_uri_complete") or ""
)
verify_complete = str(payload.get("verification_uri_complete") or verify_url) verify_complete = str(payload.get("verification_uri_complete") or verify_url)
interval = max(1, int(payload.get("interval") or 5)) interval = max(1, int(payload.get("interval") or 5))
expires_in = int(payload.get("expires_in") or 900) expires_in = int(payload.get("expires_in") or 900)
@@ -188,6 +180,8 @@ class GitHubCopilotProvider(OpenAICompatProvider):
*, *,
provider_name: str = "github_copilot", provider_name: str = "github_copilot",
): ):
from nanobot.providers.registry import find_by_name
self._copilot_access_token: str | None = None self._copilot_access_token: str | None = None
self._copilot_expires_at: float = 0.0 self._copilot_expires_at: float = 0.0
self._copilot_token_lock: asyncio.Lock = asyncio.Lock() self._copilot_token_lock: asyncio.Lock = asyncio.Lock()
@@ -223,9 +217,7 @@ class GitHubCopilotProvider(OpenAICompatProvider):
) )
timeout = httpx.Timeout(20.0, connect=20.0) timeout = httpx.Timeout(20.0, connect=20.0)
async with httpx.AsyncClient( async with httpx.AsyncClient(timeout=timeout, follow_redirects=True, trust_env=True) as client:
timeout=timeout, follow_redirects=True, trust_env=True
) as client:
response = await client.get( response = await client.get(
_resolve("NANOBOT_COPILOT_TOKEN_URL", DEFAULT_COPILOT_TOKEN_URL), _resolve("NANOBOT_COPILOT_TOKEN_URL", DEFAULT_COPILOT_TOKEN_URL),
headers=_copilot_headers(github_token.access), headers=_copilot_headers(github_token.access),
@@ -304,174 +296,3 @@ class GitHubCopilotProvider(OpenAICompatProvider):
on_tool_call_delta=on_tool_call_delta, on_tool_call_delta=on_tool_call_delta,
provider_context=provider_context, provider_context=provider_context,
) )
def get_github_copilot_model_catalog(
proxy: str | None = None,
) -> OAuthModelCatalogSnapshot:
storage = get_storage()
token = storage.load()
account_key = _catalog_account_key(getattr(token, "account_id", None))
cache_key = (
f"{storage.get_token_path()}\0{account_key}\0"
f"{_resolve('NANOBOT_COPILOT_BASE_URL', DEFAULT_COPILOT_BASE_URL)}\0{proxy or ''}"
)
return _GITHUB_COPILOT_MODEL_CATALOG.get(cache_key=cache_key, proxy=proxy)
def invalidate_github_copilot_model_catalog() -> None:
_GITHUB_COPILOT_MODEL_CATALOG.invalidate()
def _fetch_github_copilot_models(proxy: str | None) -> tuple[ProviderModelSpec, ...]:
github_token = get_storage().load()
if not github_token or not github_token.access:
raise RuntimeError("GitHub Copilot is not logged in")
common_headers = {
"Accept": "application/json",
"User-Agent": USER_AGENT,
"Editor-Version": EDITOR_VERSION,
"Editor-Plugin-Version": EDITOR_PLUGIN_VERSION,
}
client_kwargs: dict[str, Any] = {"timeout": 20.0, "follow_redirects": True}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
exchange = client.get(
_resolve("NANOBOT_COPILOT_TOKEN_URL", DEFAULT_COPILOT_TOKEN_URL),
headers={**common_headers, "Authorization": f"token {github_token.access}"},
)
exchange.raise_for_status()
exchange_mapping = _catalog_mapping(exchange.json())
copilot_token = exchange_mapping.get("token")
if not isinstance(copilot_token, str) or not copilot_token:
raise RuntimeError("GitHub Copilot token exchange returned no token")
endpoint_base = _catalog_first_text(
_catalog_mapping(exchange_mapping.get("endpoints")),
"api",
)
base_url = endpoint_base or _resolve(
"NANOBOT_COPILOT_BASE_URL",
DEFAULT_COPILOT_BASE_URL,
)
models_url = (
base_url
if base_url.rstrip("/").endswith("/models")
else f"{base_url.rstrip('/')}/models"
)
response = client.get(
models_url,
headers={**common_headers, "Authorization": f"Bearer {copilot_token}"},
)
response.raise_for_status()
return _parse_github_copilot_models(response.json())
def _parse_github_copilot_models(payload: Any) -> tuple[ProviderModelSpec, ...]:
rows = cast(dict[str, Any], payload).get("data") if isinstance(payload, dict) else None
if not isinstance(rows, list):
return ()
fallback_models = _oauth_fallback_models("github_copilot")
fallback_by_id = {model.id.split("/", 1)[-1]: model for model in fallback_models}
models: list[ProviderModelSpec] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
wire_id = _catalog_first_text(row, "id")
policy = _catalog_mapping(row.get("policy"))
endpoints = row.get("supported_endpoints")
if (
not wire_id
or wire_id in seen
or row.get("model_picker_enabled") is not True
or policy.get("state") == "disabled"
or not _copilot_transport_supported(wire_id, endpoints)
):
continue
seen.add(wire_id)
capabilities = _catalog_mapping(row.get("capabilities"))
supports = _catalog_mapping(capabilities.get("supports"))
limits = _catalog_mapping(capabilities.get("limits"))
fallback = fallback_by_id.get(wire_id)
models.append(
ProviderModelSpec(
id=f"github-copilot/{wire_id}",
label=(
_catalog_first_text(row, "name")
or (fallback.label if fallback is not None else wire_id)
),
description=(fallback.description if fallback is not None else ""),
owned_by="GitHub Copilot",
context_window=(
_catalog_positive_int(limits, "max_context_window_tokens")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=_catalog_reasoning_efforts(supports.get("reasoning_effort")),
)
)
return tuple(models)
def _copilot_transport_supported(wire_id: str, endpoints: object) -> bool:
if not isinstance(endpoints, list):
return True
supported = cast(list[object], endpoints)
if "/chat/completions" in supported:
return True
model = wire_id.lower()
return "/responses" in supported and any(
token in model for token in ("gpt-5", "o1", "o3", "o4")
)
def _oauth_fallback_models(provider_name: str) -> tuple[ProviderModelSpec, ...]:
spec = find_by_name(provider_name)
assert spec is not None
return spec.builtin_models
def _catalog_account_key(account_id: object) -> str:
value = account_id if isinstance(account_id, str) else ""
return hashlib.sha256(value.encode()).hexdigest()[:16] if value else "anonymous"
def _catalog_mapping(value: Any) -> dict[str, Any]:
return cast(dict[str, Any], value) if isinstance(value, dict) else {}
def _catalog_first_text(row: dict[str, Any], *keys: str) -> str:
for key in keys:
value = row.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _catalog_positive_int(row: dict[str, Any], *keys: str) -> int | None:
for key in keys:
value = row.get(key)
if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0:
return int(value)
return None
def _catalog_reasoning_efforts(value: Any) -> tuple[str, ...]:
if not isinstance(value, list):
return ()
return tuple(
dict.fromkeys(
item.strip()
for item in cast(list[object], value)
if isinstance(item, str) and item.strip()
)
)
_GITHUB_COPILOT_MODEL_CATALOG = OAuthModelCatalog(
fallback_models=_oauth_fallback_models("github_copilot"),
fetch=_fetch_github_copilot_models,
)
+2 -2
View File
@@ -20,7 +20,7 @@ from nanobot.providers.registry import find_by_name
from nanobot.security.network import ( from nanobot.security.network import (
PinnedDNSAsyncTransport, PinnedDNSAsyncTransport,
UnsafeURLRequestError, UnsafeURLRequestError,
resolve_url_target, async_resolve_url_target,
) )
from nanobot.utils.helpers import detect_image_mime from nanobot.utils.helpers import detect_image_mime
@@ -174,7 +174,7 @@ async def _download_image_data_url(
current_url = url current_url = url
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1): for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
if proxy: if proxy:
ok, error, _ = resolve_url_target( ok, error, _ = await async_resolve_url_target(
current_url, current_url,
trust_remote_dns=True, trust_remote_dns=True,
) )
-224
View File
@@ -1,224 +0,0 @@
"""Shared cache seam for OAuth provider model discovery."""
from __future__ import annotations
import threading
import time
from collections.abc import Callable, Sequence
from dataclasses import dataclass, replace
from typing import Literal
from loguru import logger
from nanobot.providers.registry import ProviderModelSpec
CatalogSource = Literal["remote", "cache", "stale", "fallback"]
@dataclass(frozen=True, slots=True)
class OAuthModelCatalogSnapshot:
"""One usable catalog view, including where it came from."""
models: tuple[ProviderModelSpec, ...]
source: CatalogSource
fetched_at: float
message: str | None = None
def find(self, model: str) -> ProviderModelSpec | None:
wire_id = model.split("/", 1)[-1]
return next(
(item for item in self.models if item.id.split("/", 1)[-1] == wire_id),
None,
)
@dataclass(frozen=True, slots=True)
class _CacheEntry:
snapshot: OAuthModelCatalogSnapshot
stored_at: float
class OAuthModelCatalog:
"""Cache one provider's discovery behind a small failure-tolerant interface."""
def __init__(
self,
*,
fallback_models: Sequence[ProviderModelSpec],
fetch: Callable[[str | None], Sequence[ProviderModelSpec]],
fresh_ttl_s: float = 5 * 60,
stale_ttl_s: float = 24 * 60 * 60,
failure_ttl_s: float = 30,
max_entries: int = 8,
monotonic: Callable[[], float] = time.monotonic,
wall_clock: Callable[[], float] = time.time,
) -> None:
if fresh_ttl_s < 0 or stale_ttl_s < fresh_ttl_s or failure_ttl_s < 0:
raise ValueError("catalog cache TTLs are invalid")
if max_entries < 1:
raise ValueError("catalog cache must allow at least one entry")
self._fallback_models = tuple(fallback_models)
self._fetch = fetch
self._fresh_ttl_s = fresh_ttl_s
self._stale_ttl_s = stale_ttl_s
self._failure_ttl_s = failure_ttl_s
self._max_entries = max_entries
self._monotonic = monotonic
self._wall_clock = wall_clock
self._condition = threading.Condition()
self._entries: dict[str, _CacheEntry] = {}
self._failures: dict[str, float] = {}
self._inflight: set[str] = set()
self._generation = 0
def get(self, *, cache_key: str, proxy: str | None = None) -> OAuthModelCatalogSnapshot:
"""Return a fresh catalog, sharing concurrent work and retaining a fallback."""
with self._condition:
generation = self._generation
cached = self._cached_result(cache_key)
if cached is not None:
return cached
while cache_key in self._inflight:
self._condition.wait()
if generation != self._generation:
return self._stale_or_fallback(None, self._monotonic())
cached = self._cached_result(cache_key)
if cached is not None:
return cached
self._inflight.add(cache_key)
try:
models = tuple(self._fetch(proxy))
if not models:
raise ValueError("provider returned an empty model catalog")
except Exception as exc:
logger.warning("OAuth model catalog refresh failed: type={}", type(exc).__name__)
with self._condition:
result = (
self._stale_or_fallback(None, self._monotonic())
if generation != self._generation
else self._failure_result(cache_key)
)
else:
now = self._monotonic()
result = OAuthModelCatalogSnapshot(
models=models,
source="remote",
fetched_at=self._wall_clock(),
)
with self._condition:
if generation != self._generation:
result = self._stale_or_fallback(None, now)
else:
self._store(cache_key, _CacheEntry(snapshot=result, stored_at=now))
self._failures.pop(cache_key, None)
finally:
with self._condition:
self._inflight.discard(cache_key)
self._condition.notify_all()
return result
def invalidate(self) -> None:
"""Drop cached work and prevent an older identity refresh from being stored."""
with self._condition:
self._generation += 1
self._entries.clear()
self._failures.clear()
self._condition.notify_all()
def _cached_result(self, cache_key: str) -> OAuthModelCatalogSnapshot | None:
now = self._monotonic()
entry = self._entries.get(cache_key)
if entry is not None and now - entry.stored_at < self._fresh_ttl_s:
return replace(entry.snapshot, source="cache")
failure_until = self._failures.get(cache_key)
if failure_until is not None and failure_until <= now:
self._failures.pop(cache_key, None)
elif failure_until is not None:
return self._stale_or_fallback(entry, now)
return None
def _failure_result(self, cache_key: str) -> OAuthModelCatalogSnapshot:
now = self._monotonic()
self._reserve(cache_key)
self._failures[cache_key] = now + self._failure_ttl_s
return self._stale_or_fallback(self._entries.get(cache_key), now)
def _stale_or_fallback(
self,
entry: _CacheEntry | None,
now: float,
) -> OAuthModelCatalogSnapshot:
if entry is not None and now - entry.stored_at < self._stale_ttl_s:
return replace(
entry.snapshot,
source="stale",
message="Could not refresh the online model list; showing cached models.",
)
return OAuthModelCatalogSnapshot(
models=self._fallback_models,
source="fallback",
fetched_at=self._wall_clock(),
message="Could not load the online model list; showing built-in fallback models.",
)
def _store(self, cache_key: str, entry: _CacheEntry) -> None:
self._reserve(cache_key)
self._entries[cache_key] = entry
def _reserve(self, cache_key: str) -> None:
known = set(self._entries) | set(self._failures)
if cache_key in known or len(known) < self._max_entries:
return
oldest = min(
known,
key=lambda key: (
self._entries[key].stored_at
if key in self._entries
else self._failures[key] - self._failure_ttl_s
),
)
self._entries.pop(oldest, None)
self._failures.pop(oldest, None)
def get_oauth_model_catalog(
provider_name: str,
*,
proxy: str | None = None,
) -> OAuthModelCatalogSnapshot:
"""Discover models through the owning provider module."""
if provider_name == "openai_codex":
from nanobot.providers.openai_codex_provider import get_openai_codex_model_catalog
return get_openai_codex_model_catalog(proxy)
if provider_name == "xai_grok":
from nanobot.providers.xai_grok_provider import get_xai_grok_model_catalog
return get_xai_grok_model_catalog(proxy)
if provider_name == "github_copilot":
from nanobot.providers.github_copilot_provider import get_github_copilot_model_catalog
return get_github_copilot_model_catalog(proxy)
raise ValueError(f"OAuth model discovery is not available for {provider_name}")
def invalidate_oauth_model_catalog(provider_name: str) -> None:
"""Invalidate provider discovery after its OAuth identity changes."""
if provider_name == "openai_codex":
from nanobot.providers.openai_codex_provider import (
invalidate_openai_codex_model_catalog,
)
invalidate_openai_codex_model_catalog()
elif provider_name == "xai_grok":
from nanobot.providers.xai_grok_provider import invalidate_xai_grok_model_catalog
invalidate_xai_grok_model_catalog()
elif provider_name == "github_copilot":
from nanobot.providers.github_copilot_provider import (
invalidate_github_copilot_model_catalog,
)
invalidate_github_copilot_model_catalog()
+34 -178
View File
@@ -14,10 +14,7 @@ from typing import Any, cast
import httpx import httpx
from loguru import logger from loguru import logger
from oauth_cli_kit import get_token as get_codex_token from oauth_cli_kit import get_token as get_codex_token
from oauth_cli_kit.providers import OPENAI_CODEX_PROVIDER
from oauth_cli_kit.storage import FileTokenStorage
from nanobot import __version__
from nanobot.providers.base import ( from nanobot.providers.base import (
LLMProvider, LLMProvider,
LLMResponse, LLMResponse,
@@ -25,10 +22,6 @@ from nanobot.providers.base import (
ProviderConversationState, ProviderConversationState,
resolve_stream_idle_timeout_s, resolve_stream_idle_timeout_s,
) )
from nanobot.providers.oauth_model_catalog import (
OAuthModelCatalog,
OAuthModelCatalogSnapshot,
)
from nanobot.providers.openai_responses import ( from nanobot.providers.openai_responses import (
ResponsesStreamCapture, ResponsesStreamCapture,
build_responses_state, build_responses_state,
@@ -42,11 +35,8 @@ from nanobot.providers.openai_responses import (
responses_state_items, responses_state_items,
responses_state_matches, responses_state_matches,
) )
from nanobot.providers.registry import ProviderModelSpec, find_by_name
DEFAULT_CODEX_URL = "https://chatgpt.com/backend-api/codex/responses" DEFAULT_CODEX_URL = "https://chatgpt.com/backend-api/codex/responses"
DEFAULT_OPENAI_CODEX_MODELS_URL = "https://chatgpt.com/backend-api/codex/models"
OPENAI_CODEX_CATALOG_CLIENT_VERSION = "0.144.0"
DEFAULT_ORIGINATOR = "nanobot" DEFAULT_ORIGINATOR = "nanobot"
_COMPACTION_RETAINED_CHAR_BUDGET = 256_000 _COMPACTION_RETAINED_CHAR_BUDGET = 256_000
@@ -54,6 +44,8 @@ _COMPACTION_RETAINED_CHAR_BUDGET = 256_000
class OpenAICodexProvider(LLMProvider): class OpenAICodexProvider(LLMProvider):
"""Use Codex OAuth to call the Responses API.""" """Use Codex OAuth to call the Responses API."""
supports_progress_deltas = True
def __init__( def __init__(
self, self,
default_model: str = "openai-codex/gpt-5.6-sol", default_model: str = "openai-codex/gpt-5.6-sol",
@@ -97,7 +89,9 @@ class OpenAICodexProvider(LLMProvider):
model = model or self.default_model model = model or self.default_model
sanitized_messages = self._sanitize_empty_content(messages) sanitized_messages = self._sanitize_empty_content(messages)
sanitized_state = ( sanitized_state = (
provider_context.conversation_state if provider_context is not None else None provider_context.conversation_state
if provider_context is not None
else None
) )
if sanitized_state is not None: if sanitized_state is not None:
sanitized_state = sanitized_state.with_pending_messages( sanitized_state = sanitized_state.with_pending_messages(
@@ -109,7 +103,6 @@ class OpenAICodexProvider(LLMProvider):
provider=self._responses_state_provider(), provider=self._responses_state_provider(),
model=_strip_model_prefix(model), model=_strip_model_prefix(model),
) )
session_id = provider_context.session_id if provider_context is not None else None
body: dict[str, Any] = { body: dict[str, Any] = {
"model": _strip_model_prefix(model), "model": _strip_model_prefix(model),
@@ -118,11 +111,10 @@ class OpenAICodexProvider(LLMProvider):
"instructions": system_prompt, "instructions": system_prompt,
"input": input_items, "input": input_items,
"text": {"verbosity": "medium"}, "text": {"verbosity": "medium"},
"prompt_cache_key": _prompt_cache_key(messages[:2]),
"tool_choice": tool_choice or "auto", "tool_choice": tool_choice or "auto",
"parallel_tool_calls": True, "parallel_tool_calls": True,
} }
if session_id:
body["prompt_cache_key"] = _prompt_cache_key(session_id)
body["include"] = ["reasoning.encrypted_content"] body["include"] = ["reasoning.encrypted_content"]
reasoning_options = _build_reasoning_options(reasoning_effort) reasoning_options = _build_reasoning_options(reasoning_effort)
if replayed and "gpt-5.6" in _strip_model_prefix(model).lower(): if replayed and "gpt-5.6" in _strip_model_prefix(model).lower():
@@ -176,7 +168,11 @@ class OpenAICodexProvider(LLMProvider):
) )
compact_threshold = resolve_compact_threshold( compact_threshold = resolve_compact_threshold(
(provider_context.context_window_tokens if provider_context is not None else None), (
provider_context.context_window_tokens
if provider_context is not None
else None
),
max_tokens, max_tokens,
) )
if ( if (
@@ -240,12 +236,8 @@ class OpenAICodexProvider(LLMProvider):
return response return response
async def chat( async def chat(
self, self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
messages: list[dict[str, Any]], model: str | None = None, max_tokens: int = 4096, temperature: float = 0.7,
tools: list[dict[str, Any]] | None = None,
model: str | None = None,
max_tokens: int = 4096,
temperature: float = 0.7,
reasoning_effort: str | None = None, reasoning_effort: str | None = None,
tool_choice: str | dict[str, Any] | None = None, tool_choice: str | dict[str, Any] | None = None,
provider_context: ProviderCallContext | None = None, provider_context: ProviderCallContext | None = None,
@@ -272,12 +264,8 @@ class OpenAICodexProvider(LLMProvider):
) )
async def chat_stream( async def chat_stream(
self, self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
messages: list[dict[str, Any]], model: str | None = None, max_tokens: int = 4096, temperature: float = 0.7,
tools: list[dict[str, Any]] | None = None,
model: str | None = None,
max_tokens: int = 4096,
temperature: float = 0.7,
reasoning_effort: str | None = None, reasoning_effort: str | None = None,
tool_choice: str | dict[str, Any] | None = None, tool_choice: str | dict[str, Any] | None = None,
on_content_delta: Callable[[str], Awaitable[None]] | None = None, on_content_delta: Callable[[str], Awaitable[None]] | None = None,
@@ -356,7 +344,11 @@ def _without_response_item_ids(
sanitized_input.append(raw_item) sanitized_input.append(raw_item)
continue continue
item = cast(dict[str, Any], raw_item) item = cast(dict[str, Any], raw_item)
sanitized_input.append({key: value for key, value in item.items() if key != "id"}) sanitized_input.append({
key: value
for key, value in item.items()
if key != "id"
})
body = dict(request_body) body = dict(request_body)
body["input"] = sanitized_input body["input"] = sanitized_input
@@ -452,12 +444,15 @@ async def _request_codex(
raw = text.decode("utf-8", "ignore") raw = text.decode("utf-8", "ignore")
retry_after = LLMProvider._extract_retry_after_from_headers(response.headers) retry_after = LLMProvider._extract_retry_after_from_headers(response.headers)
error_type, error_code = LLMProvider._extract_error_type_code(raw) error_type, error_code = LLMProvider._extract_error_type_code(raw)
compaction_unsupported = response.status_code in {400, 404, 422} and any( compaction_unsupported = (
marker in raw.lower() response.status_code in {400, 404, 422}
for marker in ( and any(
"context_management", marker in raw.lower()
"compact_threshold", for marker in (
"compaction_trigger", "context_management",
"compact_threshold",
"compaction_trigger",
)
) )
) )
raise _CodexHTTPError( raise _CodexHTTPError(
@@ -466,9 +461,7 @@ async def _request_codex(
retry_after=retry_after, retry_after=retry_after,
error_type=error_type, error_type=error_type,
error_code=error_code, error_code=error_code,
should_retry=_should_retry_status( should_retry=_should_retry_status(response.status_code, error_type, error_code, raw),
response.status_code, error_type, error_code, raw
),
compaction_unsupported=compaction_unsupported, compaction_unsupported=compaction_unsupported,
) )
capture = ResponsesStreamCapture() capture = ResponsesStreamCapture()
@@ -503,8 +496,9 @@ async def _request_codex(
return result return result
def _prompt_cache_key(session_id: str) -> str: def _prompt_cache_key(messages: list[dict[str, Any]]) -> str:
return hashlib.sha256(session_id.encode("utf-8")).hexdigest() raw = json.dumps(messages, ensure_ascii=True, sort_keys=True)
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
def _friendly_error(status_code: int, raw: str) -> str: def _friendly_error(status_code: int, raw: str) -> str:
@@ -541,9 +535,7 @@ def _codex_error_response(exc: Exception) -> LLMResponse:
default_detail = "HTTP request failed" default_detail = "HTTP request failed"
if status_code is not None and should_retry is None: if status_code is not None and should_retry is None:
retry_content = ( retry_content = None if int(status_code) == 429 and isinstance(exc, _CodexHTTPError) else detail
None if int(status_code) == 429 and isinstance(exc, _CodexHTTPError) else detail
)
should_retry = _should_retry_status( should_retry = _should_retry_status(
int(status_code), int(status_code),
getattr(exc, "error_type", None), getattr(exc, "error_type", None),
@@ -601,139 +593,3 @@ def _should_retry_status(
) )
) )
return status_code in LLMProvider._RETRYABLE_STATUS_CODES or status_code >= 500 return status_code in LLMProvider._RETRYABLE_STATUS_CODES or status_code >= 500
def get_openai_codex_model_catalog(
proxy: str | None = None,
) -> OAuthModelCatalogSnapshot:
storage = FileTokenStorage(token_filename=OPENAI_CODEX_PROVIDER.token_filename)
token = storage.load()
account_id = getattr(token, "account_id", None)
account_key = _catalog_account_key(account_id)
cache_key = f"{storage.get_token_path()}\0{account_key}\0{proxy or ''}"
return _OPENAI_CODEX_MODEL_CATALOG.get(cache_key=cache_key, proxy=proxy)
def invalidate_openai_codex_model_catalog() -> None:
_OPENAI_CODEX_MODEL_CATALOG.invalidate()
def _fetch_openai_codex_models(proxy: str | None) -> tuple[ProviderModelSpec, ...]:
token = get_codex_token(proxy=proxy)
account_id = getattr(token, "account_id", None)
if not isinstance(account_id, str) or not account_id:
raise RuntimeError("OpenAI Codex OAuth token has no account ID")
client_kwargs: dict[str, Any] = {"timeout": 10.0, "follow_redirects": False}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
response = client.get(
DEFAULT_OPENAI_CODEX_MODELS_URL,
params={"client_version": OPENAI_CODEX_CATALOG_CLIENT_VERSION},
headers={
"Authorization": f"Bearer {token.access}",
"chatgpt-account-id": account_id,
"originator": DEFAULT_ORIGINATOR,
"User-Agent": f"nanobot/{__version__} (python)",
"accept": "application/json",
},
)
response.raise_for_status()
return _parse_openai_codex_models(response.json())
def _parse_openai_codex_models(payload: Any) -> tuple[ProviderModelSpec, ...]:
rows = cast(dict[str, Any], payload).get("models") if isinstance(payload, dict) else None
if not isinstance(rows, list):
return ()
fallback_models = _oauth_fallback_models("openai_codex")
fallback_by_id = {model.id.split("/", 1)[-1]: model for model in fallback_models}
parsed: list[tuple[int, ProviderModelSpec]] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
wire_id = _catalog_first_text(row, "slug", "id")
if not wire_id or wire_id in seen or row.get("visibility") in {"hide", "none"}:
continue
seen.add(wire_id)
fallback = fallback_by_id.get(wire_id)
priority = row.get("priority")
parsed.append(
(
priority if isinstance(priority, int) and not isinstance(priority, bool) else 2**31,
ProviderModelSpec(
id=f"openai-codex/{wire_id}",
label=(
_catalog_first_text(row, "display_name", "name")
or (fallback.label if fallback is not None else wire_id)
),
description=(
_catalog_first_text(row, "description")
or (fallback.description if fallback is not None else "")
),
owned_by="OpenAI Codex",
context_window=(
_catalog_positive_int(row, "context_window")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=(
_catalog_reasoning_efforts(row.get("supported_reasoning_levels"))
or (fallback.reasoning_efforts if fallback is not None else ())
),
),
)
)
parsed.sort(key=lambda item: item[0])
return tuple(model for _, model in parsed)
def _oauth_fallback_models(provider_name: str) -> tuple[ProviderModelSpec, ...]:
spec = find_by_name(provider_name)
assert spec is not None
return spec.builtin_models
def _catalog_account_key(account_id: object) -> str:
value = account_id if isinstance(account_id, str) else ""
return hashlib.sha256(value.encode()).hexdigest()[:16] if value else "anonymous"
def _catalog_first_text(row: dict[str, Any], *keys: str) -> str:
for key in keys:
value = row.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _catalog_positive_int(row: dict[str, Any], *keys: str) -> int | None:
for key in keys:
value = row.get(key)
if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0:
return int(value)
return None
def _catalog_reasoning_efforts(value: Any) -> tuple[str, ...]:
if not isinstance(value, list):
return ()
efforts: list[str] = []
for item in cast(list[object], value):
if isinstance(item, str):
effort = item.strip()
elif isinstance(item, dict):
effort = _catalog_first_text(cast(dict[str, Any], item), "effort", "value", "id")
else:
effort = ""
if effort and effort not in efforts:
efforts.append(effort)
return tuple(efforts)
_OPENAI_CODEX_MODEL_CATALOG = OAuthModelCatalog(
fallback_models=_oauth_fallback_models("openai_codex"),
fetch=_fetch_openai_codex_models,
)
+8 -41
View File
@@ -20,15 +20,12 @@ from pydantic.alias_generators import to_snake
@dataclass(frozen=True) @dataclass(frozen=True)
class ProviderModelSpec: class ProviderModelSpec:
"""Curated model metadata used for fixed catalogs or online fallback.""" """A curated model exposed by providers without a model-list endpoint."""
id: str id: str
label: str = "" label: str = ""
description: str = "" description: str = ""
owned_by: str = ""
context_window: int | None = None context_window: int | None = None
reasoning_efforts: tuple[str, ...] = ()
supports_backend_search: bool = False
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -45,7 +42,7 @@ class ProviderSpec:
keywords: tuple[str, ...] # model-name keywords for matching (lowercase) keywords: tuple[str, ...] # model-name keywords for matching (lowercase)
env_key: str # env var for API key, e.g. "DASHSCOPE_API_KEY" env_key: str # env var for API key, e.g. "DASHSCOPE_API_KEY"
display_name: str = "" # shown in `nanobot status` display_name: str = "" # shown in `nanobot status`
model_catalog: str = "auto" # WebUI model-list source, including builtin/hybrid model_catalog: str = "auto" # WebUI model-list source
builtin_models: tuple[ProviderModelSpec, ...] = () builtin_models: tuple[ProviderModelSpec, ...] = ()
settings_alias_for: str = "" # compatibility alias grouped under this provider in Settings settings_alias_for: str = "" # compatibility alias grouped under this provider in Settings
@@ -410,56 +407,45 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
keywords=("openai-codex",), keywords=("openai-codex",),
env_key="", env_key="",
display_name="OpenAI Codex", display_name="OpenAI Codex",
model_catalog="hybrid", model_catalog="builtin",
builtin_models=( builtin_models=(
ProviderModelSpec( ProviderModelSpec(
id="openai-codex/gpt-5.6-sol", id="openai-codex/gpt-5.6-sol",
label="GPT-5.6-Sol", label="GPT-5.6-Sol",
description="Latest frontier agentic coding model.", description="Latest frontier agentic coding model.",
context_window=272_000, context_window=372000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max", "ultra"),
), ),
ProviderModelSpec( ProviderModelSpec(
id="openai-codex/gpt-5.6-terra", id="openai-codex/gpt-5.6-terra",
label="GPT-5.6-Terra", label="GPT-5.6-Terra",
description="Balanced agentic coding model for everyday work.", description="Balanced agentic coding model for everyday work.",
context_window=272_000, context_window=372000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max", "ultra"),
), ),
ProviderModelSpec( ProviderModelSpec(
id="openai-codex/gpt-5.6-luna", id="openai-codex/gpt-5.6-luna",
label="GPT-5.6-Luna", label="GPT-5.6-Luna",
description="Fast and affordable agentic coding model.", description="Fast and affordable agentic coding model.",
context_window=272_000, context_window=372000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max"),
), ),
ProviderModelSpec( ProviderModelSpec(
id="openai-codex/gpt-5.5", id="openai-codex/gpt-5.5",
label="GPT-5.5", label="GPT-5.5",
description="Frontier model for complex coding, research, and real-world work.", description="Frontier model for complex coding, research, and real-world work.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
), ),
ProviderModelSpec( ProviderModelSpec(
id="openai-codex/gpt-5.4", id="openai-codex/gpt-5.4",
label="GPT-5.4", label="GPT-5.4",
description="Strong model for everyday coding.", description="Strong model for everyday coding.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
), ),
ProviderModelSpec( ProviderModelSpec(
id="openai-codex/gpt-5.4-mini", id="openai-codex/gpt-5.4-mini",
label="GPT-5.4-Mini", label="GPT-5.4-Mini",
description="Small, fast, and cost-efficient model for simpler coding tasks.", description="Small, fast, and cost-efficient model for simpler coding tasks.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
), ),
ProviderModelSpec( ProviderModelSpec(
id="openai-codex/gpt-5.3-codex-spark", id="openai-codex/gpt-5.3-codex-spark",
label="GPT-5.3-Codex-Spark", label="GPT-5.3-Codex-Spark",
description="Ultra-fast coding model.", description="Ultra-fast coding model.",
context_window=128_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
), ),
), ),
backend="openai_codex", backend="openai_codex",
@@ -473,19 +459,13 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
keywords=("xai-grok", "xai_grok"), keywords=("xai-grok", "xai_grok"),
env_key="", env_key="",
display_name="xAI Grok", display_name="xAI Grok",
model_catalog="hybrid", model_catalog="builtin",
builtin_models=( builtin_models=(
ProviderModelSpec(
id="xai-grok/grok-4.6",
label="Grok 4.6",
description="Grok via xAI subscription; X Search is enabled when supported.",
context_window=500_000,
),
ProviderModelSpec( ProviderModelSpec(
id="xai-grok/grok-4.5", id="xai-grok/grok-4.5",
label="Grok 4.5", label="Grok 4.5",
description="Grok via xAI subscription; X Search is enabled when supported.", description="Grok via xAI subscription; X Search is enabled when supported.",
context_window=500_000, context_window=500000,
), ),
), ),
backend="xai_grok", backend="xai_grok",
@@ -498,19 +478,6 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
keywords=("github_copilot", "copilot"), keywords=("github_copilot", "copilot"),
env_key="", env_key="",
display_name="Github Copilot", display_name="Github Copilot",
model_catalog="hybrid",
builtin_models=(
ProviderModelSpec(
id="github-copilot/gpt-5.4-mini",
label="GPT-5.4 Mini",
description="GitHub Copilot Responses model.",
),
ProviderModelSpec(
id="github-copilot/gpt-4.1",
label="GPT-4.1",
description="GitHub Copilot chat model.",
),
),
backend="github_copilot", backend="github_copilot",
default_api_base="https://api.githubcopilot.com", default_api_base="https://api.githubcopilot.com",
strip_model_prefix=True, strip_model_prefix=True,
+188 -392
View File
@@ -4,9 +4,9 @@ from __future__ import annotations
import asyncio import asyncio
import base64 import base64
import hashlib
import json import json
import re import re
import time
import uuid import uuid
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
from typing import Any, cast from typing import Any, cast
@@ -22,24 +22,21 @@ from nanobot.providers.base import (
ToolCallRequest, ToolCallRequest,
resolve_stream_idle_timeout_s, resolve_stream_idle_timeout_s,
) )
from nanobot.providers.oauth_model_catalog import OAuthModelCatalog, OAuthModelCatalogSnapshot
from nanobot.providers.openai_responses import ( from nanobot.providers.openai_responses import (
consume_sse_with_reasoning, consume_sse_with_reasoning,
convert_messages, convert_messages,
convert_tools, convert_tools,
) )
from nanobot.providers.registry import ProviderModelSpec, find_by_name
from nanobot.providers.xai_oauth import ( from nanobot.providers.xai_oauth import (
XAI_CLIENT_VERSION, XAI_CLIENT_VERSION,
get_xai_oauth_login_status, XAIToken,
get_xai_oauth_storage_path,
get_xai_oauth_token, get_xai_oauth_token,
) )
DEFAULT_XAI_GROK_MODEL = "xai-grok/grok-4.6"
DEFAULT_XAI_GROK_URL = "https://cli-chat-proxy.grok.com/v1/responses" DEFAULT_XAI_GROK_URL = "https://cli-chat-proxy.grok.com/v1/responses"
DEFAULT_XAI_GROK_MODELS_URL = "https://cli-chat-proxy.grok.com/v1/models" DEFAULT_XAI_GROK_MODELS_URL = "https://cli-chat-proxy.grok.com/v1/models"
_HOSTED_SEARCH_MAX_TURNS = 5 DEFAULT_XAI_GROK_MODEL = "xai-grok/grok-4.5"
_MODEL_CAPABILITIES_TTL_S = 5 * 60
_MAX_ERROR_BODY_CHARS = 1000 _MAX_ERROR_BODY_CHARS = 1000
_SENSITIVE_ERROR_KEYS = { _SENSITIVE_ERROR_KEYS = {
"accesstoken", "accesstoken",
@@ -66,9 +63,7 @@ def _is_named_x_search_tool(value: object) -> bool:
class XAIGrokProvider(LLMProvider): class XAIGrokProvider(LLMProvider):
"""Call xAI's subscription proxy and expose supported hosted tools.""" """Call xAI's subscription proxy and expose supported hosted tools."""
# An incomplete hosted-tool stream can already have emitted answer text. Let the supports_progress_deltas = True
# provider close that stream segment before its one bounded recovery attempt.
supports_stream_recover_callback = True
def __init__( def __init__(
self, self,
@@ -82,19 +77,37 @@ class XAIGrokProvider(LLMProvider):
self.default_model = default_model self.default_model = default_model
self.proxy = proxy or None self.proxy = proxy or None
self._extra_body = dict(extra_body or {}) self._extra_body = dict(extra_body or {})
self._model_capabilities: dict[str, bool] | None = None
self._model_capabilities_fetched_at = 0.0
async def _supports_backend_search(self, model: str) -> bool: async def _supports_backend_search(self, token: XAIToken, model: str) -> bool:
catalog = await asyncio.to_thread( now = time.monotonic()
get_xai_grok_model_catalog, capabilities = self._model_capabilities
self.proxy, if (
) capabilities is None
if catalog.message: or now - self._model_capabilities_fetched_at >= _MODEL_CAPABILITIES_TTL_S
logger.warning( ):
"xAI model catalog unavailable; hosted X Search disabled unless cached: {}", try:
catalog.message, capabilities = await _fetch_xai_model_capabilities(
) DEFAULT_XAI_GROK_MODELS_URL,
info = catalog.find(model) _build_model_headers(token),
return bool(info and info.supports_backend_search) proxy=self.proxy,
)
except Exception as exc:
logger.warning(
"xAI model capability lookup failed; hosted X Search disabled for model {}: "
"type={} error={}",
model,
type(exc).__name__,
str(exc).strip() or "unexpected error",
)
capabilities = {}
self._model_capabilities = capabilities
self._model_capabilities_fetched_at = now
else:
self._model_capabilities = capabilities
self._model_capabilities_fetched_at = now
return capabilities.get(model, False)
async def _call_xai( async def _call_xai(
self, self,
@@ -108,7 +121,6 @@ class XAIGrokProvider(LLMProvider):
on_content_delta: Callable[[str], Awaitable[None]] | None = None, on_content_delta: Callable[[str], Awaitable[None]] | None = None,
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None, on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
on_stream_recover: Callable[[], Awaitable[None]] | None = None,
) -> LLMResponse: ) -> LLMResponse:
wire_model = _strip_model_prefix(model or self.default_model) wire_model = _strip_model_prefix(model or self.default_model)
system_prompt, input_items = convert_messages(messages) system_prompt, input_items = convert_messages(messages)
@@ -118,13 +130,17 @@ class XAIGrokProvider(LLMProvider):
token = await asyncio.to_thread(get_xai_oauth_token, proxy=self.proxy) token = await asyncio.to_thread(get_xai_oauth_token, proxy=self.proxy)
configured_tools = self._extra_body.get("tools") configured_tools = self._extra_body.get("tools")
tools_are_explicit = "tools" in self._extra_body tools_are_explicit = "tools" in self._extra_body
configured_hosted_search = isinstance(configured_tools, list) and any( configured_hosted_search = (
_is_hosted_x_search_tool(tool) for tool in cast(list[object], configured_tools) isinstance(configured_tools, list)
and any(
_is_hosted_x_search_tool(tool)
for tool in cast(list[object], configured_tools)
)
) )
supports_backend_search = False supports_backend_search = False
if not tools_are_explicit: if not tools_are_explicit:
stage = "model_capabilities" stage = "model_capabilities"
supports_backend_search = await self._supports_backend_search(wire_model) supports_backend_search = await self._supports_backend_search(token, wire_model)
converted_tools = convert_tools(tools or []) converted_tools = convert_tools(tools or [])
if isinstance(configured_tools, list): if isinstance(configured_tools, list):
converted_tools.extend(cast(list[dict[str, Any]], configured_tools)) converted_tools.extend(cast(list[dict[str, Any]], configured_tools))
@@ -135,8 +151,6 @@ class XAIGrokProvider(LLMProvider):
if supports_backend_search: if supports_backend_search:
converted_tools.append({"type": "x_search"}) converted_tools.append({"type": "x_search"})
hosted_search_enabled = supports_backend_search or configured_hosted_search
body: dict[str, Any] = { body: dict[str, Any] = {
"model": wire_model, "model": wire_model,
"store": False, "store": False,
@@ -152,65 +166,51 @@ class XAIGrokProvider(LLMProvider):
"temperature": temperature, "temperature": temperature,
"reasoning": _build_reasoning_options(reasoning_effort), "reasoning": _build_reasoning_options(reasoning_effort),
} }
if hosted_search_enabled:
# xAI's global default is intentionally unspecified. Five turns is
# their documented balanced setting and prevents a search from
# stopping after a single unsuccessful lookup.
body["max_turns"] = _HOSTED_SEARCH_MAX_TURNS
if self._extra_body: if self._extra_body:
body.update( body.update({
{key: value for key, value in self._extra_body.items() if key != "tools"} key: value
) for key, value in self._extra_body.items()
if key != "tools"
})
if tools_are_explicit and not isinstance(configured_tools, list): if tools_are_explicit and not isinstance(configured_tools, list):
body["tools"] = configured_tools body["tools"] = configured_tools
headers = _build_headers(token.access, wire_model) headers = _build_headers(token.access, wire_model)
stage = "xai_request" stage = "xai_request"
auth_retried = False try:
hosted_tool_retried = False result = await _request_xai(
retry_usage: LLMUsage | None = None DEFAULT_XAI_GROK_URL,
while True: headers,
try: body,
result = await _request_xai( proxy=self.proxy,
DEFAULT_XAI_GROK_URL, on_content_delta=on_content_delta,
headers, on_thinking_delta=on_thinking_delta,
body, on_tool_call_delta=on_tool_call_delta,
proxy=self.proxy, )
on_content_delta=on_content_delta, except _XAIHTTPError as exc:
on_thinking_delta=on_thinking_delta, if exc.status_code != 401:
on_tool_call_delta=on_tool_call_delta, raise
) stage = "oauth_refresh"
break token = await asyncio.to_thread(
except _XAIHTTPError as exc: get_xai_oauth_token,
if exc.status_code != 401 or auth_retried: proxy=self.proxy,
raise force_refresh=True,
auth_retried = True )
stage = "oauth_refresh" self._model_capabilities = None
token = await asyncio.to_thread( self._model_capabilities_fetched_at = 0.0
get_xai_oauth_token, headers = _build_headers(token.access, wire_model)
proxy=self.proxy, stage = "xai_request_retry"
force_refresh=True, result = await _request_xai(
) DEFAULT_XAI_GROK_URL,
headers = _build_headers(token.access, wire_model) headers,
stage = "xai_request_after_oauth_refresh" body,
except _XAIIncompleteHostedToolError as exc: proxy=self.proxy,
retry_usage = _combine_usage(retry_usage, exc.usage) on_content_delta=on_content_delta,
cannot_recover_stream = exc.stream_output_emitted and on_stream_recover is None on_thinking_delta=on_thinking_delta,
if hosted_tool_retried or cannot_recover_stream: on_tool_call_delta=on_tool_call_delta,
exc.usage = retry_usage )
raise
hosted_tool_retried = True
stage = "hosted_tool_recovery"
logger.warning(
"xAI response ended with unfinished hosted tool(s): {}; retrying once",
", ".join(exc.tool_names),
)
if on_stream_recover is not None:
await on_stream_recover()
headers = _build_headers(token.access, wire_model)
content, tool_calls, finish_reason, usage, reasoning_content = result content, tool_calls, finish_reason, usage, reasoning_content = result
usage = _combine_usage(retry_usage, usage)
return LLMResponse( return LLMResponse(
content=content, content=content,
tool_calls=tool_calls, tool_calls=tool_calls,
@@ -259,7 +259,6 @@ class XAIGrokProvider(LLMProvider):
on_content_delta: Callable[[str], Awaitable[None]] | None = None, on_content_delta: Callable[[str], Awaitable[None]] | None = None,
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None, on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
on_stream_recover: Callable[[], Awaitable[None]] | None = None,
) -> LLMResponse: ) -> LLMResponse:
return await self._call_xai( return await self._call_xai(
messages, messages,
@@ -272,7 +271,6 @@ class XAIGrokProvider(LLMProvider):
on_content_delta, on_content_delta,
on_thinking_delta, on_thinking_delta,
on_tool_call_delta, on_tool_call_delta,
on_stream_recover,
) )
def get_default_model(self) -> str: def get_default_model(self) -> str:
@@ -292,14 +290,6 @@ def _build_reasoning_options(reasoning_effort: str | None) -> dict[str, str]:
return options return options
def _combine_usage(left: LLMUsage | None, right: LLMUsage | None) -> LLMUsage | None:
if left is None:
return right
if right is None:
return left
return left + right
def _build_headers(token: str, model: str) -> dict[str, str]: def _build_headers(token: str, model: str) -> dict[str, str]:
conversation_id = str(uuid.uuid4()) conversation_id = str(uuid.uuid4())
return { return {
@@ -320,6 +310,44 @@ def _build_headers(token: str, model: str) -> dict[str, str]:
} }
def _build_model_headers(token: XAIToken) -> dict[str, str]:
headers = {
"Authorization": f"Bearer {token.access}",
"X-XAI-Token-Auth": "xai-grok-cli",
"x-grok-client-version": XAI_CLIENT_VERSION,
"x-grok-client-identifier": "nanobot",
"x-grok-client-mode": "headless",
"User-Agent": f"nanobot/{__version__} (python)",
"accept": "application/json",
}
claims = _decode_access_token_claims(token.access)
user_id = claims.get("sub")
if claims.get("principal_type") == "Team":
user_id = claims.get("principal_id") or user_id
if isinstance(user_id, str) and user_id:
headers["x-userid"] = user_id
email = claims.get("email")
if not isinstance(email, str) or "@" not in email:
email = token.account_id if token.account_id and "@" in token.account_id else None
if email:
headers["x-email"] = email
return headers
def _decode_access_token_claims(token: str) -> dict[str, Any]:
"""Read identity hints from the signed token; the server still authenticates it."""
parts = token.split(".")
if len(parts) < 2 or not parts[1]:
return {}
payload = parts[1]
try:
decoded = base64.urlsafe_b64decode(payload + "=" * (-len(payload) % 4))
claims = json.loads(decoded)
except (ValueError, TypeError):
return {}
return cast(dict[str, Any], claims) if isinstance(claims, dict) else {}
class _XAIHTTPError(RuntimeError): class _XAIHTTPError(RuntimeError):
def __init__( def __init__(
self, self,
@@ -341,25 +369,65 @@ class _XAIHTTPError(RuntimeError):
self.response_body = response_body self.response_body = response_body
class _XAIIncompleteHostedToolError(RuntimeError): async def _fetch_xai_model_capabilities(
"""A nominally successful xAI stream ended before a hosted tool did.""" url: str,
headers: dict[str, str],
*,
proxy: str | None = None,
) -> dict[str, bool]:
client_kwargs: dict[str, Any] = {"timeout": 10.0, "follow_redirects": False}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
async with httpx.AsyncClient(**client_kwargs) as client:
response = await client.get(url, headers=headers)
if response.status_code != 200:
raw = response.content.decode("utf-8", "ignore")
raise _build_xai_http_error(response.status_code, response.headers, raw)
try:
payload = response.json()
except ValueError as exc:
raise RuntimeError("xAI model catalog returned invalid JSON.") from exc
return _parse_xai_model_capabilities(payload)
should_retry = False # _call_xai already performs the one safe recovery attempt.
def __init__( def _parse_xai_model_capabilities(payload: Any) -> dict[str, bool]:
self, if isinstance(payload, dict):
active_tools: list[dict[str, Any]], payload = cast(dict[str, Any], payload)
*, rows: object = payload.get("data")
usage: LLMUsage | None, if not isinstance(rows, list):
stream_output_emitted: bool = False, rows = payload.get("models")
) -> None: else:
names = [str(event.get("name") or "hosted_tool") for event in active_tools] rows = payload
super().__init__( if not isinstance(rows, list):
"xAI ended the response before its hosted tool completed: " + ", ".join(names) return {}
capabilities: dict[str, bool] = {}
for row_value in cast(list[object], rows):
if not isinstance(row_value, dict):
continue
row = cast(dict[str, Any], row_value)
meta_value = row.get("_meta")
meta = cast(dict[str, Any], meta_value) if isinstance(meta_value, dict) else {}
support_value = row.get("supportsBackendSearch")
if not isinstance(support_value, bool):
support_value = row.get("supports_backend_search")
if not isinstance(support_value, bool):
support_value = meta.get("supportsBackendSearch")
if not isinstance(support_value, bool):
support_value = meta.get("supports_backend_search")
supports_backend_search = support_value if isinstance(support_value, bool) else False
identifiers = (
row.get("model"),
row.get("modelId"),
row.get("id"),
meta.get("model"),
meta.get("modelId"),
) )
self.tool_names = tuple(names) for identifier in identifiers:
self.usage = usage if isinstance(identifier, str) and identifier.strip():
self.stream_output_emitted = stream_output_emitted capabilities[_strip_model_prefix(identifier.strip())] = supports_backend_search
return capabilities
async def _request_xai( async def _request_xai(
@@ -372,39 +440,10 @@ async def _request_xai(
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None, on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None, on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]: ) -> tuple[str, list[ToolCallRequest], str, LLMUsage | None, str | None]:
active_hosted_tools: dict[str, dict[str, Any]] = {}
stream_output_emitted = False
async def _forward_content_delta(delta: str) -> None:
nonlocal stream_output_emitted
if delta:
stream_output_emitted = True
if on_content_delta is not None:
await on_content_delta(delta)
async def _forward_thinking_delta(delta: str) -> None:
nonlocal stream_output_emitted
if delta:
stream_output_emitted = True
if on_thinking_delta is not None:
await on_thinking_delta(delta)
async def _track_and_forward_tool_event(event: dict[str, Any]) -> None:
if event.get("kind") == "hosted_tool":
call_id = event.get("call_id")
if call_id:
call_id = str(call_id)
if event.get("phase") == "start":
active_hosted_tools[call_id] = dict(event)
elif event.get("phase") in {"end", "error"}:
active_hosted_tools.pop(call_id, None)
if on_tool_call_delta is not None:
await on_tool_call_delta(event)
async def _on_response_event(event: dict[str, Any]) -> None: async def _on_response_event(event: dict[str, Any]) -> None:
hosted_event = _xai_hosted_tool_event(event) hosted_event = _xai_hosted_tool_event(event)
if hosted_event is not None: if hosted_event is not None and on_tool_call_delta is not None:
await _track_and_forward_tool_event(hosted_event) await on_tool_call_delta(hosted_event)
client_kwargs: dict[str, Any] = {"timeout": resolve_stream_idle_timeout_s()} client_kwargs: dict[str, Any] = {"timeout": resolve_stream_idle_timeout_s()}
if proxy: if proxy:
@@ -415,34 +454,13 @@ async def _request_xai(
content = await response.aread() content = await response.aread()
raw = content.decode("utf-8", "ignore") raw = content.decode("utf-8", "ignore")
raise _build_xai_http_error(response.status_code, response.headers, raw) raise _build_xai_http_error(response.status_code, response.headers, raw)
result = await consume_sse_with_reasoning( return await consume_sse_with_reasoning(
response, response,
on_content_delta=(_forward_content_delta if on_content_delta is not None else None), on_content_delta=on_content_delta,
# Always observe tool events so protocol validation also works for on_tool_call_delta=on_tool_call_delta,
# non-streaming callers that did not request UI progress callbacks. on_reasoning_delta=on_thinking_delta,
on_tool_call_delta=_track_and_forward_tool_event, on_response_event=_on_response_event if on_tool_call_delta else None,
on_reasoning_delta=(
_forward_thinking_delta if on_thinking_delta is not None else None
),
on_response_event=_on_response_event,
) )
if result[2] != "error" and active_hosted_tools:
active = list(active_hosted_tools.values())
for event in active:
await _track_and_forward_tool_event(
{
**event,
"phase": "error",
"result": None,
"error": "xAI ended the response before this hosted tool completed.",
}
)
raise _XAIIncompleteHostedToolError(
active,
usage=result[3],
stream_output_emitted=stream_output_emitted,
)
return result
def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None: def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
@@ -456,33 +474,19 @@ def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
"phase": "start", "phase": "start",
"call_id": str(call_id), "call_id": str(call_id),
"name": "x_search", "name": "x_search",
"arguments": _xai_hosted_tool_arguments(event.get("input", event.get("arguments"))), "arguments": _xai_hosted_tool_arguments(
event.get("input", event.get("arguments"))
),
"result": None, "result": None,
} }
if event_type not in {"response.output_item.added", "response.output_item.done"}: if event_type != "response.output_item.done":
return None return None
item = event.get("item") item = event.get("item")
if not isinstance(item, dict): if not isinstance(item, dict):
return None return None
item = cast(dict[str, Any], item) item = cast(dict[str, Any], item)
item_type = item.get("type") if item.get("type") != "custom_tool_call":
if item_type == "x_search_call":
call_id = item.get("id") or item.get("call_id") or event.get("item_id")
if not call_id:
return None
phase = "start" if event_type == "response.output_item.added" else "end"
return {
"kind": "hosted_tool",
"phase": phase,
"call_id": str(call_id),
"name": "x_search",
"arguments": _xai_hosted_tool_arguments(item.get("action")),
"result": (
{"status": str(item.get("status") or "completed")} if phase == "end" else None
),
}
if event_type != "response.output_item.done" or item_type != "custom_tool_call":
return None return None
tool_name = item.get("name") tool_name = item.get("name")
if not isinstance(tool_name, str) or not tool_name.startswith("x_"): if not isinstance(tool_name, str) or not tool_name.startswith("x_"):
@@ -495,7 +499,9 @@ def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
"phase": "end", "phase": "end",
"call_id": str(call_id), "call_id": str(call_id),
"name": "x_search", "name": "x_search",
"arguments": _xai_hosted_tool_arguments(item.get("input", item.get("arguments"))), "arguments": _xai_hosted_tool_arguments(
item.get("input", item.get("arguments"))
),
# Keep the useful search subtype, but do not persist large hosted results # Keep the useful search subtype, but do not persist large hosted results
# in WebUI activity messages. The model answer already carries citations. # in WebUI activity messages. The model answer already carries citations.
"result": {"name": tool_name}, "result": {"name": tool_name},
@@ -604,8 +610,6 @@ def _xai_error_response(exc: Exception) -> LLMResponse:
should_retry = True if should_retry is None else should_retry should_retry = True if should_retry is None else should_retry
elif isinstance(exc, _XAIHTTPError): elif isinstance(exc, _XAIHTTPError):
error_kind = "http" error_kind = "http"
elif isinstance(exc, _XAIIncompleteHostedToolError):
error_kind = "provider"
if status_code is not None and should_retry is None: if status_code is not None and should_retry is None:
should_retry = _should_retry_status( should_retry = _should_retry_status(
int(status_code), int(status_code),
@@ -615,11 +619,9 @@ def _xai_error_response(exc: Exception) -> LLMResponse:
) )
message = str(exc).strip() or "unexpected error" message = str(exc).strip() or "unexpected error"
retry_after = getattr(exc, "retry_after", None) retry_after = getattr(exc, "retry_after", None)
usage = getattr(exc, "usage", None)
return LLMResponse( return LLMResponse(
content=f"Error calling xAI ({type(exc).__name__}): {message}", content=f"Error calling xAI ({type(exc).__name__}): {message}",
finish_reason="error", finish_reason="error",
usage=usage if isinstance(usage, LLMUsage) else None,
retry_after=retry_after, retry_after=retry_after,
error_status_code=int(status_code) if status_code is not None else None, error_status_code=int(status_code) if status_code is not None else None,
error_kind=error_kind, error_kind=error_kind,
@@ -647,209 +649,3 @@ def _should_retry_status(
) )
) )
return status_code in LLMProvider._RETRYABLE_STATUS_CODES or status_code >= 500 # pyright: ignore[reportPrivateUsage] return status_code in LLMProvider._RETRYABLE_STATUS_CODES or status_code >= 500 # pyright: ignore[reportPrivateUsage]
def get_xai_grok_model_catalog(proxy: str | None = None) -> OAuthModelCatalogSnapshot:
token = get_xai_oauth_login_status()
account_key = _catalog_account_key(getattr(token, "account_id", None))
cache_key = f"{get_xai_oauth_storage_path()}\0{account_key}\0{proxy or ''}"
return _XAI_GROK_MODEL_CATALOG.get(cache_key=cache_key, proxy=proxy)
def invalidate_xai_grok_model_catalog() -> None:
_XAI_GROK_MODEL_CATALOG.invalidate()
def _fetch_xai_grok_models(proxy: str | None) -> tuple[ProviderModelSpec, ...]:
token = get_xai_oauth_token(proxy=proxy)
client_kwargs: dict[str, Any] = {"timeout": 10.0, "follow_redirects": False}
if proxy:
client_kwargs.update(proxy=proxy, trust_env=False)
with httpx.Client(**client_kwargs) as client:
response = client.get(
DEFAULT_XAI_GROK_MODELS_URL,
headers=_build_xai_model_headers(token.access, token.account_id),
)
response.raise_for_status()
return _parse_xai_grok_models(response.json())
def _parse_xai_grok_models(payload: Any) -> tuple[ProviderModelSpec, ...]:
if isinstance(payload, dict):
payload_mapping = cast(dict[str, Any], payload)
rows: object = payload_mapping.get("data")
if not isinstance(rows, list):
rows = payload_mapping.get("models")
else:
rows = payload
if not isinstance(rows, list):
return ()
fallback_models = _oauth_fallback_models("xai_grok")
fallback_by_id = {model.id.split("/", 1)[-1]: model for model in fallback_models}
models: list[ProviderModelSpec] = []
seen: set[str] = set()
for value in cast(list[object], rows):
if not isinstance(value, dict):
continue
row = cast(dict[str, Any], value)
meta = _catalog_mapping(row.get("_meta"))
raw_id = next(
(
candidate.strip()
for candidate in (
row.get("id"),
row.get("model"),
row.get("modelId"),
row.get("name"),
meta.get("id"),
meta.get("model"),
meta.get("modelId"),
)
if isinstance(candidate, str) and candidate.strip()
),
None,
)
if raw_id is None:
continue
wire_id = raw_id.split("/", 1)[-1]
if wire_id in seen:
continue
seen.add(wire_id)
fallback = fallback_by_id.get(wire_id)
label = _catalog_first_text(row, "display_name", "label", "name") or _catalog_first_text(
meta,
"display_name",
"label",
"name",
)
if not label or label == raw_id:
label = fallback.label if fallback is not None else wire_id
models.append(
ProviderModelSpec(
id=f"xai-grok/{wire_id}",
label=label,
description=(
_catalog_first_text(row, "description")
or _catalog_first_text(meta, "description")
or (fallback.description if fallback is not None else "")
),
owned_by=(
_catalog_first_text(row, "owned_by", "owner", "organization")
or _catalog_first_text(meta, "owned_by", "owner", "organization")
or (fallback.owned_by if fallback is not None else "xAI")
),
context_window=(
_catalog_positive_int(row, "context_window", "context_length")
or _catalog_positive_int(meta, "context_window", "context_length")
or (fallback.context_window if fallback is not None else None)
),
reasoning_efforts=_catalog_reasoning_efforts(
row.get("reasoning_efforts", meta.get("reasoning_efforts"))
),
supports_backend_search=_catalog_bool_field(
row,
"supports_backend_search",
"supportsBackendSearch",
),
)
)
return tuple(models)
def _build_xai_model_headers(access_token: str, account_id: str | None) -> dict[str, str]:
headers = {
"Authorization": f"Bearer {access_token}",
"X-XAI-Token-Auth": "xai-grok-cli",
"x-grok-client-version": XAI_CLIENT_VERSION,
"x-grok-client-identifier": "nanobot",
"x-grok-client-mode": "headless",
"User-Agent": f"nanobot/{__version__} (python)",
"accept": "application/json",
}
claims = _decode_access_token_claims(access_token)
user_id = claims.get("sub")
if claims.get("principal_type") == "Team":
user_id = claims.get("principal_id") or user_id
if isinstance(user_id, str) and user_id:
headers["x-userid"] = user_id
email = claims.get("email")
if not isinstance(email, str) or "@" not in email:
email = account_id if account_id and "@" in account_id else None
if email:
headers["x-email"] = email
return headers
def _decode_access_token_claims(token: str) -> dict[str, Any]:
parts = token.split(".")
if len(parts) < 2 or not parts[1]:
return {}
try:
decoded = base64.urlsafe_b64decode(parts[1] + "=" * (-len(parts[1]) % 4))
claims = json.loads(decoded)
except (ValueError, TypeError):
return {}
return cast(dict[str, Any], claims) if isinstance(claims, dict) else {}
def _oauth_fallback_models(provider_name: str) -> tuple[ProviderModelSpec, ...]:
spec = find_by_name(provider_name)
assert spec is not None
return spec.builtin_models
def _catalog_account_key(account_id: object) -> str:
value = account_id if isinstance(account_id, str) else ""
return hashlib.sha256(value.encode()).hexdigest()[:16] if value else "anonymous"
def _catalog_mapping(value: Any) -> dict[str, Any]:
return cast(dict[str, Any], value) if isinstance(value, dict) else {}
def _catalog_first_text(row: dict[str, Any], *keys: str) -> str:
for key in keys:
value = row.get(key)
if isinstance(value, str) and value.strip():
return value.strip()
return ""
def _catalog_positive_int(row: dict[str, Any], *keys: str) -> int | None:
for key in keys:
value = row.get(key)
if isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0:
return int(value)
return None
def _catalog_bool_field(row: dict[str, Any], *keys: str) -> bool:
for key in keys:
value = row.get(key)
if isinstance(value, bool):
return value
meta = row.get("_meta")
return _catalog_bool_field(_catalog_mapping(meta), *keys) if isinstance(meta, dict) else False
def _catalog_reasoning_efforts(value: Any) -> tuple[str, ...]:
if not isinstance(value, list):
return ()
efforts: list[str] = []
for item in cast(list[object], value):
if isinstance(item, str):
effort = item.strip()
elif isinstance(item, dict):
effort = _catalog_first_text(cast(dict[str, Any], item), "effort", "value", "id")
else:
effort = ""
if effort and effort not in efforts:
efforts.append(effort)
return tuple(efforts)
_XAI_GROK_MODEL_CATALOG = OAuthModelCatalog(
fallback_models=_oauth_fallback_models("xai_grok"),
fetch=_fetch_xai_grok_models,
)
+7 -5
View File
@@ -210,18 +210,20 @@ class RuntimeClient:
async def compact_session(self, session_key: str) -> SessionSnapshot: async def compact_session(self, session_key: str) -> SessionSnapshot:
"""Run token consolidation for one session.""" """Run token consolidation for one session."""
session = self._loop.sessions.get_or_create(session_key) session = await self._loop.sessions.get_or_create_async(session_key)
runtime = self._loop.runtime_for_session(session) runtime = await self._loop.runtime_for_session_async(session)
await self._loop.consolidator.maybe_consolidate_by_tokens( await self._loop.consolidator.maybe_consolidate_by_tokens(
session, session,
runtime=runtime, runtime=runtime,
) )
return snapshot_from_session(self._loop.sessions.get_or_create(session_key)) return snapshot_from_session(
await self._loop.sessions.get_or_create_async(session_key)
)
async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None: async def compact_idle_session(self, session_key: str, *, max_suffix: int = 8) -> str | None:
"""Run idle-session compaction for one session and return the summary.""" """Run idle-session compaction for one session and return the summary."""
session = self._loop.sessions.get_or_create(session_key) session = await self._loop.sessions.get_or_create_async(session_key)
runtime = self._loop.runtime_for_session(session) runtime = await self._loop.runtime_for_session_async(session)
return await self._loop.consolidator.compact_idle_session( return await self._loop.consolidator.compact_idle_session(
session_key, session_key,
runtime=runtime, runtime=runtime,
+144 -42
View File
@@ -29,6 +29,7 @@ _BLOCKED_NETWORKS = [
_URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE) _URL_RE = re.compile(r"https?://[^\s\"'`;|<>]+", re.IGNORECASE)
_allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = [] _allowed_networks: list[ipaddress.IPv4Network | ipaddress.IPv6Network] = []
_DNS_RESOLUTION_TIMEOUT_SECONDS = 5.0
def is_loopback_host(host: str) -> bool: def is_loopback_host(host: str) -> bool:
@@ -75,6 +76,63 @@ def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
return any(normalized in net for net in _BLOCKED_NETWORKS) return any(normalized in net for net in _BLOCKED_NETWORKS)
def _parse_url_hostname(url: str) -> tuple[str | None, str | None]:
try:
parsed = urlparse(url)
except Exception as exc:
return None, str(exc)
if parsed.scheme not in ("http", "https"):
return None, f"Only http/https allowed, got '{parsed.scheme or 'none'}'"
if not parsed.netloc:
return None, "Missing domain"
if not parsed.hostname:
return None, "Missing hostname"
return parsed.hostname, None
def _unresolved_target_result(
hostname: str,
*,
trust_remote_dns: bool,
) -> tuple[bool, str, tuple[str, ...]]:
if not trust_remote_dns:
return False, f"Cannot resolve hostname: {hostname}", ()
normalized_hostname = hostname.rstrip(".").lower()
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
return False, f"Blocked local/internal hostname: {hostname}", ()
try:
literal_addr = ipaddress.ip_address(normalized_hostname)
except ValueError:
return True, "", ()
if _is_private(literal_addr):
return False, f"Blocked private/internal address: {literal_addr}", ()
return True, "", (str(_normalize_addr(literal_addr)),)
def _resolved_target_result(
hostname: str,
infos: list[Any],
*,
allow_loopback: bool,
) -> tuple[bool, str, tuple[str, ...]]:
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
for info in infos:
try:
addr = ipaddress.ip_address(info[4][0])
except (IndexError, TypeError, ValueError):
continue
addrs.append(addr)
if allow_loopback and _is_allowed_loopback_target(hostname, addrs):
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
for addr in addrs:
if _is_private(addr):
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", ()
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs))
def resolve_url_target( def resolve_url_target(
url: str, url: str,
*, *,
@@ -97,52 +155,43 @@ def resolve_url_target(
resolved_ips contains the public IPs that were validated for this URL, or resolved_ips contains the public IPs that were validated for this URL, or
is empty when an unresolved hostname is delegated to a trusted proxy. is empty when an unresolved hostname is delegated to a trusted proxy.
""" """
try: hostname, error = _parse_url_hostname(url)
p = urlparse(url) if hostname is None:
except Exception as e: return False, error or "Missing hostname", ()
return False, str(e), ()
if p.scheme not in ("http", "https"):
return False, f"Only http/https allowed, got '{p.scheme or 'none'}'", ()
if not p.netloc:
return False, "Missing domain", ()
hostname = p.hostname
if not hostname:
return False, "Missing hostname", ()
try: try:
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM) infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
except socket.gaierror: except socket.gaierror:
if not trust_remote_dns: return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
return False, f"Cannot resolve hostname: {hostname}", () return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
normalized_hostname = hostname.rstrip(".").lower()
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
return False, f"Blocked local/internal hostname: {hostname}", ()
try: async def async_resolve_url_target(
literal_addr = ipaddress.ip_address(normalized_hostname) url: str,
except ValueError: *,
return True, "", () allow_loopback: bool = False,
if _is_private(literal_addr): trust_remote_dns: bool = False,
return False, f"Blocked private/internal address: {literal_addr}", () timeout_s: float = _DNS_RESOLUTION_TIMEOUT_SECONDS,
return True, "", (str(_normalize_addr(literal_addr)),) ) -> tuple[bool, str, tuple[str, ...]]:
"""Resolve and validate an HTTP target without blocking the event loop."""
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = [] hostname, error = _parse_url_hostname(url)
for info in infos: if hostname is None:
try: return False, error or "Missing hostname", ()
addr = ipaddress.ip_address(info[4][0]) loop = asyncio.get_running_loop()
except ValueError: try:
continue infos = await asyncio.wait_for(
addrs.append(addr) loop.getaddrinfo(
if allow_loopback and _is_allowed_loopback_target(hostname, addrs): hostname,
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs)) None,
for addr in addrs: family=socket.AF_UNSPEC,
if _is_private(addr): type=socket.SOCK_STREAM,
return False, f"Blocked: {hostname} resolves to private/internal address {addr}", () ),
timeout=timeout_s,
return True, "", tuple(dict.fromkeys(str(_normalize_addr(addr)) for addr in addrs)) )
except asyncio.TimeoutError:
return False, f"Timed out resolving hostname: {hostname}", ()
except socket.gaierror:
return _unresolved_target_result(hostname, trust_remote_dns=trust_remote_dns)
return _resolved_target_result(hostname, infos, allow_loopback=allow_loopback)
def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]: def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str]:
@@ -151,6 +200,16 @@ def validate_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool
return ok, error return ok, error
async def async_validate_url_target(
url: str,
*,
allow_loopback: bool = False,
) -> tuple[bool, str]:
"""Validate a URL using the event loop's asynchronous resolver."""
ok, error, _ = await async_resolve_url_target(url, allow_loopback=allow_loopback)
return ok, error
def env_proxy_applies_to_url(url: str) -> bool: def env_proxy_applies_to_url(url: str) -> bool:
"""Return True when process proxy settings would proxy this URL.""" """Return True when process proxy settings would proxy this URL."""
try: try:
@@ -277,7 +336,10 @@ class PinnedDNSAsyncTransport(httpx.AsyncBaseTransport):
async def handle_async_request(self, request: httpx.Request) -> httpx.Response: async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
url = str(request.url) url = str(request.url)
ok, error, resolved_ips = resolve_url_target(url, allow_loopback=self._allow_loopback) ok, error, resolved_ips = await async_resolve_url_target(
url,
allow_loopback=self._allow_loopback,
)
if not ok: if not ok:
raise UnsafeURLRequestError(error, request=request) raise UnsafeURLRequestError(error, request=request)
async with self._resolver_lock: async with self._resolver_lock:
@@ -320,6 +382,46 @@ def validate_resolved_url(url: str) -> tuple[bool, str]:
return True, "" return True, ""
async def async_validate_resolved_url(url: str) -> tuple[bool, str]:
"""Validate a redirect target without blocking on domain resolution."""
try:
parsed = urlparse(url)
except Exception:
return True, ""
hostname = parsed.hostname
if not hostname:
return True, ""
try:
addr = ipaddress.ip_address(hostname)
except ValueError:
loop = asyncio.get_running_loop()
try:
infos = await asyncio.wait_for(
loop.getaddrinfo(
hostname,
None,
family=socket.AF_UNSPEC,
type=socket.SOCK_STREAM,
),
timeout=_DNS_RESOLUTION_TIMEOUT_SECONDS,
)
except asyncio.TimeoutError:
return False, f"Timed out resolving redirect hostname: {hostname}"
except socket.gaierror:
return True, ""
for info in infos:
try:
addr = ipaddress.ip_address(info[4][0])
except (IndexError, TypeError, ValueError):
continue
if _is_private(addr):
return False, f"Redirect target {hostname} resolves to private address {addr}"
return True, ""
if _is_private(addr):
return False, f"Redirect target is a private address: {addr}"
return True, ""
def contains_internal_url(command: str, *, allow_loopback: bool = False) -> bool: def contains_internal_url(command: str, *, allow_loopback: bool = False) -> bool:
"""Return True if the command string contains a URL targeting an internal/private address.""" """Return True if the command string contains a URL targeting an internal/private address."""
for m in _URL_RE.finditer(command): for m in _URL_RE.finditer(command):
+29
View File
@@ -0,0 +1,29 @@
"""Compatibility bridge for asynchronous SessionManager operations."""
import asyncio
import inspect
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar, cast
from nanobot.utils.cancellation import shield_and_drain
_SessionResult = TypeVar("_SessionResult")
async def call_session_manager(
manager: object,
async_method_name: str,
sync_method: Callable[..., _SessionResult],
/,
*args: Any,
**kwargs: Any,
) -> _SessionResult:
"""Prefer a class-declared coroutine, or offload the established sync contract."""
class_async_method = inspect.getattr_static(type(manager), async_method_name, None)
if inspect.iscoroutinefunction(class_async_method):
async_method = cast(
Callable[..., Awaitable[_SessionResult]],
getattr(manager, async_method_name),
)
return await async_method(*args, **kwargs)
return await shield_and_drain(asyncio.to_thread(sync_method, *args, **kwargs))
+138 -150
View File
@@ -1,5 +1,6 @@
"""Session management for conversation history.""" """Session management for conversation history."""
import asyncio
import base64 import base64
import errno import errno
import hashlib import hashlib
@@ -28,6 +29,7 @@ from nanobot.runtime_context import (
public_history_message, public_history_message,
) )
from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY from nanobot.session.model_selection import SESSION_MODEL_PRESET_METADATA_KEY
from nanobot.utils.cancellation import shield_and_drain
from nanobot.utils.helpers import ( from nanobot.utils.helpers import (
content_with_media_breadcrumbs, content_with_media_breadcrumbs,
ensure_dir, ensure_dir,
@@ -71,6 +73,7 @@ _WORKSPACE_STATE_DIR = ".nanobot"
_WORKSPACE_ID_FILE = "workspace-id" _WORKSPACE_ID_FILE = "workspace-id"
_WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$") _WORKSPACE_ID_RE = re.compile(r"^[0-9a-f]{32}$")
_SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30 _SESSION_MIGRATION_LOCK_TIMEOUT_SECONDS = 30
_SESSION_FILES_LOCK_TIMEOUT_SECONDS = 5
_SESSION_FILES_LOCK_FILENAME = ".session-files.lock" _SESSION_FILES_LOCK_FILENAME = ".session-files.lock"
_COPY_CHUNK_SIZE = 1024 * 1024 _COPY_CHUNK_SIZE = 1024 * 1024
@@ -82,120 +85,6 @@ def _json_object(value: object) -> dict[str, Any]:
return cast(dict[str, Any], value) return cast(dict[str, Any], value)
def _archive_offset(data: dict[str, Any]) -> int:
"""Read the Memory archive watermark across the field-name migration."""
for key in ("last_archived", "last_consolidated"):
offset = cast(object, data.get(key))
if isinstance(offset, int) and not isinstance(offset, bool):
return offset
return 0
# TODO(0.3.2): Remove the write_stdin replay migration after 0.3.1.
def _migrate_legacy_exec_arguments(container: dict[str, Any]) -> bool:
raw_arguments = cast(object, container.get("arguments"))
encoded = isinstance(raw_arguments, str)
if encoded:
try:
decoded: object = json.loads(raw_arguments)
except json.JSONDecodeError:
return False
else:
decoded = raw_arguments
if not isinstance(decoded, dict):
return False
arguments = cast(dict[str, Any], decoded)
changed = False
if "chars" in arguments:
if "input" not in arguments:
arguments["input"] = arguments["chars"]
arguments.pop("chars")
changed = True
wait_key = (
"wait_timeout_ms"
if arguments.get("wait_for") or arguments.get("until_exit")
else "yield_time_ms"
)
if "timeout_ms" not in arguments and wait_key in arguments:
arguments["timeout_ms"] = arguments[wait_key]
for key in ("yield_time_ms", "wait_timeout_ms", "max_output_chars", "max_output_tokens"):
if key in arguments:
arguments.pop(key)
changed = True
if changed:
container["arguments"] = (
json.dumps(arguments, ensure_ascii=False, separators=(",", ":"))
if encoded
else arguments
)
return changed
def _migrate_legacy_exec_tool_call(value: object) -> bool:
if not isinstance(value, dict):
return False
tool_call = cast(dict[str, Any], value)
function_value = cast(object, tool_call.get("function"))
function = (
cast(dict[str, Any], function_value)
if isinstance(function_value, dict)
else tool_call
)
name = function.get("name")
if name not in {"write_stdin", "exec_session"}:
return False
changed = name == "write_stdin"
if changed:
function["name"] = "exec_session"
return _migrate_legacy_exec_arguments(function) or changed
def _migrate_legacy_exec_message(message: dict[str, Any]) -> bool:
changed = False
if message.get("name") == "write_stdin":
message["name"] = "exec_session"
changed = True
tool_calls = cast(object, message.get("tool_calls"))
if isinstance(tool_calls, list):
for tool_call in cast(list[object], tool_calls):
changed = _migrate_legacy_exec_tool_call(tool_call) or changed
return changed
def _migrate_legacy_exec_session_records(
messages: list[dict[str, Any]],
metadata: dict[str, Any],
) -> bool:
changed = False
for message in messages:
changed = _migrate_legacy_exec_message(message) or changed
checkpoint_value = cast(object, metadata.get(_RUNTIME_CHECKPOINT_KEY))
if not isinstance(checkpoint_value, dict):
return changed
checkpoint = cast(dict[str, Any], checkpoint_value)
assistant = cast(object, checkpoint.get("assistant_message"))
if isinstance(assistant, dict):
changed = _migrate_legacy_exec_message(cast(dict[str, Any], assistant)) or changed
pending = cast(object, checkpoint.get("pending_tool_calls"))
if isinstance(pending, list):
for tool_call in cast(list[object], pending):
changed = _migrate_legacy_exec_tool_call(tool_call) or changed
completed = cast(object, checkpoint.get("completed_tool_results"))
if isinstance(completed, list):
for result in cast(list[object], completed):
if isinstance(result, dict):
result_data = cast(dict[str, Any], result)
if result_data.get("name") == "write_stdin":
result_data["name"] = "exec_session"
changed = True
return changed
def _is_provider_state_record_line(line: str) -> bool: def _is_provider_state_record_line(line: str) -> bool:
"""Recognize the canonical private record without decoding its opaque payload.""" """Recognize the canonical private record without decoding its opaque payload."""
return _PROVIDER_STATE_RECORD_PREFIX_RE.match(line) is not None return _PROVIDER_STATE_RECORD_PREFIX_RE.match(line) is not None
@@ -286,10 +175,7 @@ class Session:
created_at: datetime = field(default_factory=datetime.now) created_at: datetime = field(default_factory=datetime.now)
updated_at: datetime = field(default_factory=datetime.now) updated_at: datetime = field(default_factory=datetime.now)
metadata: dict[str, Any] = field(default_factory=dict) metadata: dict[str, Any] = field(default_factory=dict)
# Legacy storage name for the Memory ingestion watermark. New code should last_consolidated: int = 0 # Number of messages already consolidated to files
# use ``last_archived`` so this progress is not confused with model-context
# compaction. Keep the field while persisted sessions and SDK callers migrate.
last_consolidated: int = 0
provider_state: ProviderConversationState | None = field(default=None, repr=False) provider_state: ProviderConversationState | None = field(default=None, repr=False)
policy: SessionPolicy = field(default_factory=SessionPolicy, repr=False, compare=False) policy: SessionPolicy = field(default_factory=SessionPolicy, repr=False, compare=False)
@@ -307,15 +193,6 @@ class Session:
): ):
self.last_consolidated = 0 self.last_consolidated = 0
@property
def last_archived(self) -> int:
"""Number of transcript messages already written to the Memory journal."""
return self.last_consolidated
@last_archived.setter
def last_archived(self, value: int) -> None:
self.last_consolidated = value
def add_message(self, role: str, content: str, **kwargs: Any) -> None: def add_message(self, role: str, content: str, **kwargs: Any) -> None:
"""Add a message to the session.""" """Add a message to the session."""
msg = { msg = {
@@ -340,9 +217,9 @@ class Session:
A positive ``max_messages`` applies an explicit caller-owned count A positive ``max_messages`` applies an explicit caller-owned count
limit. The normal model path relies on ``max_tokens`` instead. limit. The normal model path relies on ``max_tokens`` instead.
""" """
replay_start = self.last_archived replay_start = self.last_consolidated
if replay_start: if replay_start:
# ``last_archived`` is archive progress, not a replay boundary. # ``last_consolidated`` is archive progress, not a replay boundary.
# Keep a small raw suffix for continuity, extending back to the user # Keep a small raw suffix for continuity, extending back to the user
# that started an assistant/tool sequence when necessary. # that started an assistant/tool sequence when necessary.
recent_start = recent_message_start_index( recent_start = recent_message_start_index(
@@ -356,8 +233,8 @@ class Session:
if max_messages <= 0: if max_messages <= 0:
start_idx = 0 start_idx = 0
else: else:
unarchived_count = len(self.messages) - self.last_archived unarchived_count = len(self.messages) - self.last_consolidated
if replay_start < self.last_archived and unarchived_count < max_messages: if replay_start < self.last_consolidated and unarchived_count < max_messages:
# The archived replay suffix can exceed the nominal count when one # The archived replay suffix can exceed the nominal count when one
# tool-heavy turn spans the boundary. Preserve that complete turn. # tool-heavy turn spans the boundary. Preserve that complete turn.
start_idx = 0 start_idx = 0
@@ -480,7 +357,7 @@ class Session:
def clear(self) -> None: def clear(self) -> None:
"""Clear all messages and reset session to initial state.""" """Clear all messages and reset session to initial state."""
self.messages = [] self.messages = []
self.last_archived = 0 self.last_consolidated = 0
self.provider_state = None self.provider_state = None
self.updated_at = datetime.now() self.updated_at = datetime.now()
self.metadata.pop("_last_summary", None) self.metadata.pop("_last_summary", None)
@@ -495,11 +372,11 @@ class Session:
Returns a RetentionResult with dropped messages and how many of those Returns a RetentionResult with dropped messages and how many of those
were in the already-consolidated prefix. This method mutates were in the already-consolidated prefix. This method mutates
self.messages and self.last_archived in place. self.messages and self.last_consolidated in place.
""" """
if max_messages <= 0: if max_messages <= 0:
dropped = list(self.messages) dropped = list(self.messages)
lc = self.last_archived lc = self.last_consolidated
self.clear() self.clear()
return RetentionResult( return RetentionResult(
dropped=dropped, dropped=dropped,
@@ -512,7 +389,7 @@ class Session:
) )
original = list(self.messages) original = list(self.messages)
before_lc = self.last_archived before_lc = self.last_consolidated
start_idx = max(0, len(self.messages) - max_messages) start_idx = max(0, len(self.messages) - max_messages)
if extend_to_user: if extend_to_user:
@@ -572,7 +449,7 @@ class Session:
if i < before_lc and id(m) not in retained_ids if i < before_lc and id(m) not in retained_ids
) )
# New last_archived = count of retained messages that were inside # New last_consolidated = count of retained messages that were inside
# the old consolidated prefix. # the old consolidated prefix.
new_lc = sum( new_lc = sum(
1 for i, m in enumerate(original) 1 for i, m in enumerate(original)
@@ -580,7 +457,7 @@ class Session:
) )
self.messages = retained self.messages = retained
self.last_archived = new_lc self.last_consolidated = new_lc
if dropped: if dropped:
self.provider_state = None self.provider_state = None
self.updated_at = datetime.now() self.updated_at = datetime.now()
@@ -686,7 +563,8 @@ class JsonlSessionStore:
self.sessions_dir = ensure_dir(root / workspace_id) self.sessions_dir = ensure_dir(root / workspace_id)
self.legacy_sessions_dir = get_legacy_sessions_dir() self.legacy_sessions_dir = get_legacy_sessions_dir()
self._session_files_lock = FileLock( self._session_files_lock = FileLock(
str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME) str(self.sessions_dir / _SESSION_FILES_LOCK_FILENAME),
timeout=_SESSION_FILES_LOCK_TIMEOUT_SECONDS,
) )
with self._session_files_lock: with self._session_files_lock:
self._migrate_from_workspace(canonical_workspace) self._migrate_from_workspace(canonical_workspace)
@@ -1188,7 +1066,12 @@ class JsonlSessionStore:
if isinstance(updated_at_value, str) and updated_at_value if isinstance(updated_at_value, str) and updated_at_value
else None else None
) )
last_consolidated = _archive_offset(data) offset = cast(object, data.get("last_consolidated", 0))
last_consolidated = (
offset
if isinstance(offset, int) and not isinstance(offset, bool)
else 0
)
elif record_type == _PROVIDER_STATE_RECORD_TYPE: elif record_type == _PROVIDER_STATE_RECORD_TYPE:
provider_state = ProviderConversationState.from_private_record( provider_state = ProviderConversationState.from_private_record(
data.get("state") data.get("state")
@@ -1206,8 +1089,6 @@ class JsonlSessionStore:
provider_state=provider_state, provider_state=provider_state,
) )
self._overlay_runtime_checkpoint_unlocked(session, path) self._overlay_runtime_checkpoint_unlocked(session, path)
if _migrate_legacy_exec_session_records(session.messages, session.metadata):
session.provider_state = None
return session return session
except _SESSION_DATA_ERRORS as e: except _SESSION_DATA_ERRORS as e:
logger.warning("Failed to load session {}: {}", key, e) logger.warning("Failed to load session {}: {}", key, e)
@@ -1270,7 +1151,12 @@ class JsonlSessionStore:
if isinstance(updated_at_value, str) and updated_at_value: if isinstance(updated_at_value, str) and updated_at_value:
with suppress(ValueError): with suppress(ValueError):
updated_at = datetime.fromisoformat(updated_at_value) updated_at = datetime.fromisoformat(updated_at_value)
last_consolidated = _archive_offset(data) offset = cast(object, data.get("last_consolidated", 0))
last_consolidated = (
offset
if isinstance(offset, int) and not isinstance(offset, bool)
else 0
)
elif record_type == _PROVIDER_STATE_RECORD_TYPE: elif record_type == _PROVIDER_STATE_RECORD_TYPE:
candidate = ProviderConversationState.from_private_record( candidate = ProviderConversationState.from_private_record(
data.get("state") data.get("state")
@@ -1298,8 +1184,6 @@ class JsonlSessionStore:
provider_state=provider_state, provider_state=provider_state,
) )
self._overlay_runtime_checkpoint_unlocked(session, path) self._overlay_runtime_checkpoint_unlocked(session, path)
if _migrate_legacy_exec_session_records(session.messages, session.metadata):
session.provider_state = None
return session return session
except _SESSION_DATA_ERRORS as e: except _SESSION_DATA_ERRORS as e:
logger.warning("Repair failed for session {}: {}", key, e) logger.warning("Repair failed for session {}: {}", key, e)
@@ -1430,9 +1314,6 @@ class JsonlSessionStore:
"created_at": session.created_at.isoformat(), "created_at": session.created_at.isoformat(),
"updated_at": session.updated_at.isoformat(), "updated_at": session.updated_at.isoformat(),
"metadata": session.metadata, "metadata": session.metadata,
"last_archived": session.last_archived,
# Keep old nanobot releases able to read sessions written
# during the field-name migration.
"last_consolidated": session.last_consolidated, "last_consolidated": session.last_consolidated,
} }
f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n") f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n")
@@ -1577,7 +1458,6 @@ class JsonlSessionStore:
continue continue
else: else:
messages.append(data) messages.append(data)
_migrate_legacy_exec_session_records(messages, metadata)
return { return {
"key": stored_key or key, "key": stored_key or key,
"created_at": created_at, "created_at": created_at,
@@ -1766,6 +1646,7 @@ class SessionManager:
self._cache: OrderedDict[str, Session] = OrderedDict() self._cache: OrderedDict[str, Session] = OrderedDict()
# Preserve identity for sessions held by active callers without retaining idle ones. # Preserve identity for sessions held by active callers without retaining idle ones.
self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary() self._overflow_cache: WeakValueDictionary[str, Session] = WeakValueDictionary()
self._async_session_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary()
self._max_cached_sessions = SESSION_CACHE_MAX_SIZE self._max_cached_sessions = SESSION_CACHE_MAX_SIZE
self._delete_observer: Callable[[str], None] | None = None self._delete_observer: Callable[[str], None] | None = None
@@ -1865,6 +1746,28 @@ class SessionManager:
self._remember(session) self._remember(session)
return session return session
def _async_session_lock(self, key: str) -> asyncio.Lock:
lock = self._async_session_locks.get(key)
if lock is None:
lock = asyncio.Lock()
self._async_session_locks[key] = lock
return lock
async def get_or_create_async(self, key: str) -> Session:
"""Load a session without running file I/O or lock waits on the event loop."""
cached = self.get_cached(key)
if cached is not None:
return cached
async with self._async_session_lock(key):
cached = self.get_cached(key)
if cached is not None:
return cached
session = await asyncio.to_thread(self._load, key)
if session is None:
session = Session(key=key)
self._remember(session)
return session
def get_or_create_transient( def get_or_create_transient(
self, self,
key: str, key: str,
@@ -1898,6 +1801,17 @@ class SessionManager:
self._store.save(session, fsync=fsync) self._store.save(session, fsync=fsync)
self._remember(session) self._remember(session)
async def save_async(self, session: Session, *, fsync: bool = False) -> None:
"""Persist a session without blocking the caller's event loop."""
if not session.policy.persist:
return
async def save_and_remember() -> None:
await asyncio.to_thread(self._store.save, session, fsync=fsync)
self._remember(session)
await shield_and_drain(save_and_remember())
def save_runtime_checkpoint(self, session: Session) -> None: def save_runtime_checkpoint(self, session: Session) -> None:
"""Persist volatile recovery state without rewriting long history.""" """Persist volatile recovery state without rewriting long history."""
if not session.policy.persist: if not session.policy.persist:
@@ -1910,6 +1824,23 @@ class SessionManager:
# they opt into a dedicated checkpoint primitive. # they opt into a dedicated checkpoint primitive.
self.save(session) self.save(session)
async def save_runtime_checkpoint_async(self, session: Session) -> None:
"""Persist an in-flight checkpoint without blocking the event loop."""
if not session.policy.persist:
return
async def save_and_remember() -> None:
if self._store is self._jsonl_store:
await asyncio.to_thread(
self._jsonl_store.save_runtime_checkpoint,
session,
)
else:
await asyncio.to_thread(self._store.save, session)
self._remember(session)
await shield_and_drain(save_and_remember())
def rename_model_preset(self, old_name: str, new_name: str) -> int: def rename_model_preset(self, old_name: str, new_name: str) -> int:
"""Rename a session-scoped model preset across durable and live sessions.""" """Rename a session-scoped model preset across durable and live sessions."""
if old_name == new_name: if old_name == new_name:
@@ -1951,6 +1882,21 @@ class SessionManager:
raise raise
return len(changed) return len(changed)
async def flush_all_async(self) -> int:
"""Re-save every cached session without blocking the event loop."""
cached = dict(self._overflow_cache.items())
cached.update(self._cache)
flushed = 0
for key, session in cached.items():
try:
await shield_and_drain(
asyncio.to_thread(self._store.save, session, fsync=True)
)
flushed += 1
except Exception:
logger.warning("Failed to flush session {}", key, exc_info=True)
return flushed
def flush_all(self) -> int: def flush_all(self) -> int:
"""Re-save every cached session with fsync for durable shutdown. """Re-save every cached session with fsync for durable shutdown.
@@ -1982,6 +1928,18 @@ class SessionManager:
self._delete_observer(key) self._delete_observer(key)
return deleted return deleted
async def delete_session_async(self, key: str) -> bool:
"""Delete a session without blocking the event loop."""
async def delete_and_notify() -> bool:
self.invalidate(key)
deleted = await asyncio.to_thread(self._store.delete, key)
if self._delete_observer is not None:
self._delete_observer(key)
return deleted
return await shield_and_drain(delete_and_notify())
def restore_sessions_to_workspace(self) -> SessionRestoreResult: def restore_sessions_to_workspace(self) -> SessionRestoreResult:
"""Restore session files to the pre-relocation path for an explicit rollback.""" """Restore session files to the pre-relocation path for an explicit rollback."""
return self._jsonl_store.restore_to_workspace() return self._jsonl_store.restore_to_workspace()
@@ -2025,8 +1983,8 @@ class SessionManager:
for key in _FORK_VOLATILE_METADATA_KEYS: for key in _FORK_VOLATILE_METADATA_KEYS:
metadata.pop(key, None) metadata.pop(key, None)
last_consolidated = min(source.last_archived, len(copied)) last_consolidated = min(source.last_consolidated, len(copied))
if source.last_archived > len(copied): if source.last_consolidated > len(copied):
metadata.pop("_last_summary", None) metadata.pop("_last_summary", None)
last_consolidated = 0 last_consolidated = 0
@@ -2054,6 +2012,10 @@ class SessionManager:
"""Read session metadata without loading the transcript.""" """Read session metadata without loading the transcript."""
return cast(dict[str, Any] | None, self._store.read_metadata(key)) return cast(dict[str, Any] | None, self._store.read_metadata(key))
async def read_session_metadata_async(self, key: str) -> dict[str, Any] | None:
"""Read session metadata without blocking the event loop."""
return await asyncio.to_thread(self.read_session_metadata, key)
def update_session_metadata( def update_session_metadata(
self, self,
key: str, key: str,
@@ -2067,5 +2029,31 @@ class SessionManager:
session.metadata.update(deepcopy(updates)) session.metadata.update(deepcopy(updates))
return updated return updated
async def update_session_metadata_async(
self,
key: str,
updates: dict[str, Any],
*,
fsync: bool = False,
) -> bool:
"""Update metadata without blocking the event loop."""
async def update_and_refresh_cache() -> bool:
updated = await asyncio.to_thread(
self._store.update_metadata,
key,
updates,
fsync=fsync,
)
if updated and (session := self.get_cached(key)) is not None:
session.metadata.update(deepcopy(updates))
return updated
return await shield_and_drain(update_and_refresh_cache())
def list_sessions(self) -> list[dict[str, Any]]: def list_sessions(self) -> list[dict[str, Any]]:
return cast(list[dict[str, Any]], self._store.list_sessions()) return cast(list[dict[str, Any]], self._store.list_sessions())
async def list_sessions_async(self) -> list[dict[str, Any]]:
"""List persisted sessions without blocking the event loop."""
return await asyncio.to_thread(self.list_sessions)
+62 -31
View File
@@ -1,8 +1,8 @@
"""Durable, side-effect-safe recovery for interrupted WebUI turns. """Durable, side-effect-safe recovery for interrupted WebUI turns.
The coordinator owns restart policy. Checkpoint materialization is a session The coordinator owns restart policy. AgentLoop only exposes checkpoint
operation shared with AgentLoop lifecycle boundaries, so transport code never materialization and an admission hook, so transport code never has to guess
has to guess whether an interrupted tool call is safe to replay. whether an interrupted tool call is safe to replay.
""" """
from __future__ import annotations from __future__ import annotations
@@ -25,10 +25,10 @@ from nanobot.bus.outbound_events import (
) )
from nanobot.bus.queue import MessageBus from nanobot.bus.queue import MessageBus
from nanobot.session import turn_continuation from nanobot.session import turn_continuation
from nanobot.session.async_compat import call_session_manager
from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata from nanobot.session.keys import UNIFIED_SESSION_KEY, last_channel_from_metadata
from nanobot.session.manager import Session, SessionManager from nanobot.session.manager import Session, SessionManager
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
from nanobot.webui.session_identity import webui_chat_id, webui_session_key
RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint" RUNTIME_CHECKPOINT_KEY = "runtime_checkpoint"
PENDING_USER_TURN_KEY = "pending_user_turn" PENDING_USER_TURN_KEY = "pending_user_turn"
@@ -461,6 +461,37 @@ class RecoveryCoordinator:
repr=False, repr=False,
) )
async def _get_or_create_session(self, key: str) -> Session:
return await call_session_manager(
self.sessions,
"get_or_create_async",
self.sessions.get_or_create,
key,
)
async def _save_session(self, session: Session) -> None:
await call_session_manager(
self.sessions,
"save_async",
self.sessions.save,
session,
)
async def _read_session_metadata(self, key: str) -> dict[str, Any] | None:
return await call_session_manager(
self.sessions,
"read_session_metadata_async",
self.sessions.read_session_metadata,
key,
)
async def _list_sessions(self) -> list[dict[str, Any]]:
return await call_session_manager(
self.sessions,
"list_sessions_async",
self.sessions.list_sessions,
)
def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None: def register_recovery_task(self, session_key: str, task: asyncio.Task[Any]) -> None:
"""Track the task that owns an explicit recovery continuation.""" """Track the task that owns an explicit recovery continuation."""
self._active_recovery_tasks[session_key] = task self._active_recovery_tasks[session_key] = task
@@ -483,8 +514,8 @@ class RecoveryCoordinator:
async def scan(self) -> None: async def scan(self) -> None:
"""Recover every interrupted WebUI session once at gateway startup.""" """Recover every interrupted WebUI session once at gateway startup."""
for key in self._recovery_candidates(): for key in await self._recovery_candidates():
metadata_payload = self.sessions.read_session_metadata(key) metadata_payload = await self._read_session_metadata(key)
raw_metadata = metadata_payload.get("metadata") if metadata_payload else None raw_metadata = metadata_payload.get("metadata") if metadata_payload else None
metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {} metadata = cast(dict[str, Any], raw_metadata) if isinstance(raw_metadata, dict) else {}
route = self._websocket_route_for(key, metadata) route = self._websocket_route_for(key, metadata)
@@ -493,7 +524,7 @@ class RecoveryCoordinator:
unfinished = self._has_unfinished_webui_transcript(key) unfinished = self._has_unfinished_webui_transcript(key)
if not self._needs_recovery(metadata) and not unfinished: if not self._needs_recovery(metadata) and not unfinished:
continue continue
session = self.sessions.get_or_create(key) session = await self._get_or_create_session(key)
try: try:
await self._recover_session(session, route[1]) await self._recover_session(session, route[1])
await self._requeue_pending_followups(session) await self._requeue_pending_followups(session)
@@ -508,14 +539,14 @@ class RecoveryCoordinator:
reason="recovery_failed", reason="recovery_failed",
can_continue=False, can_continue=False,
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(route[1], failed) await self._publish(route[1], failed)
def _recovery_candidates(self) -> list[str]: async def _recovery_candidates(self) -> list[str]:
"""Discover canonical and transcript-only WebUI sessions cheaply.""" """Discover canonical and transcript-only WebUI sessions cheaply."""
candidates = dict.fromkeys( candidates = dict.fromkeys(
key key
for item in self.sessions.list_sessions() for item in await self._list_sessions()
if isinstance((key := item.get("key")), str) if isinstance((key := item.get("key")), str)
) )
try: try:
@@ -524,7 +555,7 @@ class RecoveryCoordinator:
# duplicating its filename and migration rules here would drift. # duplicating its filename and migration rules here would drift.
from nanobot.webui.session_list_index import list_webui_sessions from nanobot.webui.session_list_index import list_webui_sessions
for item in list_webui_sessions(self.sessions): for item in await asyncio.to_thread(list_webui_sessions, self.sessions):
key = item.get("key") key = item.get("key")
if isinstance(key, str): if isinstance(key, str):
candidates.setdefault(key, None) candidates.setdefault(key, None)
@@ -550,7 +581,7 @@ class RecoveryCoordinator:
"""Reject stale queued recoveries and let new user input supersede them.""" """Reject stale queued recoveries and let new user input supersede them."""
recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY) recovery_id = message.metadata.get(RECOVERY_INBOUND_METADATA_KEY)
if isinstance(recovery_id, str): if isinstance(recovery_id, str):
session = self.sessions.get_or_create(message.session_key) session = await self._get_or_create_session(message.session_key)
state = recovery_state_from_metadata(session.metadata) state = recovery_state_from_metadata(session.metadata)
return bool( return bool(
state state
@@ -559,7 +590,7 @@ class RecoveryCoordinator:
) )
if message.channel != "websocket": if message.channel != "websocket":
return True return True
session = self.sessions.get_or_create(message.session_key) session = await self._get_or_create_session(message.session_key)
state = recovery_state_from_metadata(session.metadata) state = recovery_state_from_metadata(session.metadata)
if state and state["status"] in {"resuming", "awaiting_user", "failed"}: if state and state["status"] in {"resuming", "awaiting_user", "failed"}:
await self._cancel_active_recovery(message.session_key) await self._cancel_active_recovery(message.session_key)
@@ -573,13 +604,13 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)), attempts=cast(int, state.get("attempts", 0)),
reason="superseded", reason="superseded",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(message.chat_id, recovered) await self._publish(message.chat_id, recovered)
return True return True
async def turn_completed(self, session_key: str) -> None: async def turn_completed(self, session_key: str) -> None:
"""Resolve a resuming state after the recovered turn commits.""" """Resolve a resuming state after the recovered turn commits."""
session = self.sessions.get_or_create(session_key) session = await self._get_or_create_session(session_key)
state = recovery_state_from_metadata(session.metadata) state = recovery_state_from_metadata(session.metadata)
if not state or state["status"] != "resuming": if not state or state["status"] != "resuming":
return return
@@ -593,7 +624,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)), attempts=cast(int, state.get("attempts", 0)),
reason="continued", reason="continued",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(route[1], recovered) await self._publish(route[1], recovered)
async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]: async def handle_action(self, action: str, payload: dict[str, Any]) -> dict[str, Any]:
@@ -604,7 +635,7 @@ class RecoveryCoordinator:
raise RecoveryActionError("missing chat_id") raise RecoveryActionError("missing chat_id")
if not isinstance(recovery_id, str) or not recovery_id: if not isinstance(recovery_id, str) or not recovery_id:
raise RecoveryActionError("missing recovery_id") raise RecoveryActionError("missing recovery_id")
session = self.sessions.get_or_create(self._session_key(chat_id)) session = await self._get_or_create_session(self._session_key(chat_id))
state = recovery_state_from_metadata(session.metadata) state = recovery_state_from_metadata(session.metadata)
if not state or state["recovery_id"] != recovery_id: if not state or state["recovery_id"] != recovery_id:
raise RecoveryActionError("recovery state is stale", status=409) raise RecoveryActionError("recovery state is stale", status=409)
@@ -619,7 +650,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 0)), attempts=cast(int, state.get("attempts", 0)),
reason="dismissed", reason="dismissed",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, next_state) await self._publish(chat_id, next_state)
return next_state return next_state
if action != "continue": if action != "continue":
@@ -636,7 +667,7 @@ class RecoveryCoordinator:
reason="user_confirmed", reason="user_confirmed",
resume_message_count=len(session.messages), resume_message_count=len(session.messages),
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, next_state) await self._publish(chat_id, next_state)
await self._queue_continuation(session, chat_id, next_state) await self._queue_continuation(session, chat_id, next_state)
return next_state return next_state
@@ -669,7 +700,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 1)), attempts=cast(int, state.get("attempts", 1)),
reason="loop_guard", reason="loop_guard",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, next_state) await self._publish(chat_id, next_state)
elif self._has_unfinished_webui_transcript(session.key): elif self._has_unfinished_webui_transcript(session.key):
# A normal last-client shutdown can materialize the checkpoint # A normal last-client shutdown can materialize the checkpoint
@@ -691,7 +722,7 @@ class RecoveryCoordinator:
), ),
can_continue=can_continue, can_continue=can_continue,
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
if state and state["status"] in {"awaiting_user", "failed"}: if state and state["status"] in {"awaiting_user", "failed"}:
@@ -707,7 +738,7 @@ class RecoveryCoordinator:
attempts=cast(int, state.get("attempts", 1)), attempts=cast(int, state.get("attempts", 1)),
reason="loop_guard", reason="loop_guard",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
@@ -725,7 +756,7 @@ class RecoveryCoordinator:
reason="checkpoint_unknown", reason="checkpoint_unknown",
can_continue=False, can_continue=False,
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint): if checkpoint is not None and not _runtime_checkpoint_is_well_formed(checkpoint):
@@ -739,7 +770,7 @@ class RecoveryCoordinator:
reason="checkpoint_invalid", reason="checkpoint_invalid",
can_continue=False, can_continue=False,
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
if phase == "final_response": if phase == "final_response":
@@ -751,7 +782,7 @@ class RecoveryCoordinator:
attempts=0, attempts=0,
reason="answer_restored", reason="answer_restored",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, recovered) await self._publish(chat_id, recovered)
return return
if phase in _UNCERTAIN_TOOL_PHASES or pending_calls: if phase in _UNCERTAIN_TOOL_PHASES or pending_calls:
@@ -763,7 +794,7 @@ class RecoveryCoordinator:
attempts=0, attempts=0,
reason="tool_state_unknown", reason="tool_state_unknown",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
return return
# A gateway restart is a lifecycle boundary. Never enqueue model work # A gateway restart is a lifecycle boundary. Never enqueue model work
@@ -778,7 +809,7 @@ class RecoveryCoordinator:
attempts=0, attempts=0,
reason="restart_requires_confirmation", reason="restart_requires_confirmation",
) )
self.sessions.save(session) await self._save_session(session)
await self._publish(chat_id, waiting) await self._publish(chat_id, waiting)
async def _queue_continuation( async def _queue_continuation(
@@ -878,7 +909,7 @@ class RecoveryCoordinator:
return state return state
def _session_key(self, chat_id: str) -> str: def _session_key(self, chat_id: str) -> str:
return UNIFIED_SESSION_KEY if self.unified_session else webui_session_key(chat_id) return UNIFIED_SESSION_KEY if self.unified_session else f"websocket:{chat_id}"
@staticmethod @staticmethod
def _has_unfinished_webui_transcript(session_key: str) -> bool: def _has_unfinished_webui_transcript(session_key: str) -> bool:
@@ -930,9 +961,9 @@ class RecoveryCoordinator:
session_key: str, session_key: str,
metadata: Mapping[str, Any], metadata: Mapping[str, Any],
) -> tuple[str, str] | None: ) -> tuple[str, str] | None:
chat_id = webui_chat_id(session_key) if session_key.startswith("websocket:"):
if chat_id is not None: chat_id = session_key.split(":", 1)[1]
return ("websocket", chat_id) return ("websocket", chat_id) if chat_id else None
if session_key == UNIFIED_SESSION_KEY: if session_key == UNIFIED_SESSION_KEY:
route = last_channel_from_metadata(metadata) route = last_channel_from_metadata(metadata)
if route and route[0] == "websocket": if route and route[0] == "websocket":
+7 -3
View File
@@ -147,10 +147,10 @@ def prepare_save_boundary(ctx: TurnContext) -> None:
if ctx.session is not None: if ctx.session is not None:
clear_internal_continuation_state(ctx.session.metadata) clear_internal_continuation_state(ctx.session.metadata)
assert ctx.transcript_input is not None
ctx.save_skip = _save_skip_for_turn( ctx.save_skip = _save_skip_for_turn(
message_metadata=ctx.msg.metadata, message_metadata=ctx.msg.metadata,
initial_message_count=ctx.transcript_input.message_count, initial_message_count=len(ctx.initial_messages),
history_count=len(ctx.history),
input_persisted_early=ctx.input_persisted_early, input_persisted_early=ctx.input_persisted_early,
) )
@@ -185,6 +185,7 @@ def _save_skip_for_turn(
*, *,
message_metadata: Mapping[str, Any] | None, message_metadata: Mapping[str, Any] | None,
initial_message_count: int, initial_message_count: int,
history_count: int,
input_persisted_early: bool, input_persisted_early: bool,
) -> int: ) -> int:
"""Return the persisted-message append boundary for this turn.""" """Return the persisted-message append boundary for this turn."""
@@ -192,7 +193,10 @@ def _save_skip_for_turn(
return initial_message_count return initial_message_count
if internal_continuation_inbound(message_metadata): if internal_continuation_inbound(message_metadata):
return initial_message_count return initial_message_count
if not input_persisted_early: # build_messages may merge the current message into a same-role history tail.
# Runner-appended messages start at initial_message_count in either shape.
has_standalone_current = initial_message_count > 1 + history_count
if has_standalone_current and not input_persisted_early:
return initial_message_count - 1 return initial_message_count - 1
return initial_message_count return initial_message_count
+22 -72
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import re import re
import time import time
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
@@ -56,7 +57,6 @@ from nanobot.webui.metadata import (
WEBSOCKET_TURN_OWNER_METADATA_KEY, WEBSOCKET_TURN_OWNER_METADATA_KEY,
WEBUI_TURN_METADATA_KEY, WEBUI_TURN_METADATA_KEY,
) )
from nanobot.webui.session_identity import is_webui_session_key
from nanobot.webui.transcript import append_session_message_input from nanobot.webui.transcript import append_session_message_input
WEBUI_SESSION_METADATA_KEY = "webui" WEBUI_SESSION_METADATA_KEY = "webui"
@@ -169,76 +169,30 @@ def _title_inputs(session: Session) -> tuple[str, str]:
return user_text, assistant_text return user_text, assistant_text
def _latest_title_inputs(session: Session) -> tuple[str, str]:
"""Latest user/assistant texts, for turns executed on a shared session."""
user_text = ""
assistant_text = ""
for message in reversed(session.messages):
if message.get("_command") is True:
continue
if is_hidden_history_message(message):
continue
message = public_history_message(message)
role = message.get("role")
content = message.get("content")
if not isinstance(content, str) or not content.strip():
continue
content = strip_think(content)
if not content:
continue
if role == "user" and not user_text:
user_text = content.strip()
elif role == "assistant" and not assistant_text:
assistant_text = content.strip()
if user_text and assistant_text:
break
return user_text, assistant_text
async def maybe_generate_webui_title( async def maybe_generate_webui_title(
*, *,
sessions: SessionManager, sessions: SessionManager,
session_key: str, session_key: str,
provider: LLMProvider, provider: LLMProvider,
model: str, model: str,
target_session_key: str | None = None,
) -> bool: ) -> bool:
"""Generate and persist a short title for WebUI-owned sessions. """Generate and persist a short title for WebUI-owned sessions only."""
session = await sessions.get_or_create_async(session_key)
``session_key`` owns the conversation content. Under unified-session if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True:
routing this is the shared session while WebUI renders per-chat sessions,
so pass ``target_session_key`` to project the title onto that per-chat
session instead of storing it on the shared one.
"""
routed_session = sessions.get_or_create(session_key)
target_is_routed = target_session_key is None or target_session_key == session_key
if target_is_routed or target_session_key is None:
target_session = routed_session
else:
target_session = sessions.get_or_create(target_session_key)
if (
routed_session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True
and target_session.metadata.get(WEBUI_SESSION_METADATA_KEY) is not True
):
return False return False
if target_session.metadata.get(WEBUI_TITLE_USER_EDITED_METADATA_KEY) is True: if session.metadata.get(WEBUI_TITLE_USER_EDITED_METADATA_KEY) is True:
return False return False
current_title = target_session.metadata.get(WEBUI_TITLE_METADATA_KEY) current_title = session.metadata.get(WEBUI_TITLE_METADATA_KEY)
if isinstance(current_title, str) and current_title.strip(): if isinstance(current_title, str) and current_title.strip():
cleaned_current_title = clean_generated_title(current_title) cleaned_current_title = clean_generated_title(current_title)
if cleaned_current_title: if cleaned_current_title:
if cleaned_current_title != current_title: if cleaned_current_title != current_title:
target_session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title session.metadata[WEBUI_TITLE_METADATA_KEY] = cleaned_current_title
sessions.save(target_session) await sessions.save_async(session)
return False return False
target_session.metadata.pop(WEBUI_TITLE_METADATA_KEY, None) session.metadata.pop(WEBUI_TITLE_METADATA_KEY, None)
if target_is_routed: user_text, assistant_text = _title_inputs(session)
user_text, assistant_text = _title_inputs(routed_session)
else:
# Shared-session content mixes every channel; generation runs right
# after this turn, so its exchange is the latest pair.
user_text, assistant_text = _latest_title_inputs(routed_session)
if not user_text: if not user_text:
return False return False
@@ -287,15 +241,14 @@ async def maybe_generate_webui_title(
response.finish_reason, response.finish_reason,
) )
return False return False
target_session.metadata[WEBUI_TITLE_METADATA_KEY] = title session.metadata[WEBUI_TITLE_METADATA_KEY] = title
sessions.save(target_session) await sessions.save_async(session)
return True return True
async def maybe_generate_webui_title_after_turn( async def maybe_generate_webui_title_after_turn(
*, *,
channel: str, channel: str,
chat_id: str,
metadata: dict[str, Any], metadata: dict[str, Any],
sessions: SessionManager, sessions: SessionManager,
session_key: str, session_key: str,
@@ -304,15 +257,11 @@ async def maybe_generate_webui_title_after_turn(
) -> bool: ) -> bool:
if channel != "websocket" or metadata.get(WEBUI_SESSION_METADATA_KEY) is not True: if channel != "websocket" or metadata.get(WEBUI_SESSION_METADATA_KEY) is not True:
return False return False
origin_session_key = f"{channel}:{chat_id}"
return await maybe_generate_webui_title( return await maybe_generate_webui_title(
sessions=sessions, sessions=sessions,
session_key=session_key, session_key=session_key,
provider=provider, provider=provider,
model=model, model=model,
target_session_key=(
origin_session_key if origin_session_key != session_key else None
),
) )
@@ -489,8 +438,8 @@ class WebuiTurnRoutePolicy:
) )
and route.channel == "websocket" and route.channel == "websocket"
): ):
session = self.sessions.get_or_create(session_key) session = self.sessions.get_cached(session_key)
if session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True: if session is not None and session.metadata.get(WEBUI_SESSION_METADATA_KEY) is True:
metadata = dict(route.metadata) metadata = dict(route.metadata)
turn_prefix = "session-input" if internal_user_input else "subagent" turn_prefix = "session-input" if internal_user_input else "subagent"
metadata.update({ metadata.update({
@@ -629,10 +578,10 @@ class WebuiTurnCoordinator:
event.context.channel != "system" event.context.channel != "system"
or envelope is None or envelope is None
or envelope["target_session_key"] != session_key or envelope["target_session_key"] != session_key
or not is_webui_session_key(session_key) or not session_key.startswith("websocket:")
): ):
return return
persisted = self.sessions.read_session_metadata(session_key) persisted = await self.sessions.read_session_metadata_async(session_key)
metadata_value: object = persisted.get("metadata") if persisted is not None else None metadata_value: object = persisted.get("metadata") if persisted is not None else None
metadata = ( metadata = (
cast(dict[str, Any], metadata_value) cast(dict[str, Any], metadata_value)
@@ -643,7 +592,8 @@ class WebuiTurnCoordinator:
return return
public_metadata = _session_message_public_metadata(envelope) public_metadata = _session_message_public_metadata(envelope)
try: try:
append_session_message_input( await asyncio.to_thread(
append_session_message_input,
session_key, session_key,
content=event.content, content=event.content,
created_at_ms=envelope["created_at_ms"], created_at_ms=envelope["created_at_ms"],
@@ -668,8 +618,9 @@ class WebuiTurnCoordinator:
def _handle_session_turn_started(self, event: SessionTurnStarted) -> None: def _handle_session_turn_started(self, event: SessionTurnStarted) -> None:
if not self._is_websocket_event(event.context): if not self._is_websocket_event(event.context):
return return
session = self.sessions.get_or_create(event.context.session_key) session = self.sessions.get_cached(event.context.session_key)
mark_webui_session(session, event.context.metadata) if session is not None:
mark_webui_session(session, event.context.metadata)
async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None: async def _handle_run_status_changed(self, event: TurnRunStatusChanged) -> None:
if not self._is_websocket_event(event.context): if not self._is_websocket_event(event.context):
@@ -755,7 +706,7 @@ class WebuiTurnCoordinator:
if msg.channel != "websocket": if msg.channel != "websocket":
return return
session = self.sessions.get_or_create(session_key) session = await self.sessions.get_or_create_async(session_key)
await self.bus.publish_outbound( await self.bus.publish_outbound(
outbound_message_for_event( outbound_message_for_event(
channel=msg.channel, channel=msg.channel,
@@ -783,7 +734,6 @@ class WebuiTurnCoordinator:
) -> None: ) -> None:
generated = await maybe_generate_webui_title_after_turn( generated = await maybe_generate_webui_title_after_turn(
channel=event.context.channel, channel=event.context.channel,
chat_id=event.context.chat_id,
metadata=event.context.metadata, metadata=event.context.metadata,
sessions=self.sessions, sessions=self.sessions,
session_key=event.context.session_key, session_key=event.context.session_key,
+2 -2
View File
@@ -1,6 +1,6 @@
--- ---
name: my name: my
description: Inspect and optionally adjust the agent's runtime state. Use to check the current model or preset, context window and runtime limits, workspace and tool configuration, subagent status, and request routing metadata such as channel, chat ID, and sender ID; diagnose unavailable capabilities; change allowed runtime settings; or store temporary session scratchpad values. description: Inspect and optionally adjust the agent's runtime state. Use to check the current model or preset, context window, iteration progress and limits, token usage, workspace and tool configuration, subagent status, and request routing metadata such as channel, chat ID, and sender ID; diagnose unavailable capabilities; change allowed runtime settings; or store temporary session scratchpad values.
--- ---
# Self-Awareness # Self-Awareness
@@ -9,7 +9,7 @@ description: Inspect and optionally adjust the agent's runtime state. Use to che
1. **Identify the situation** from the categories below 1. **Identify the situation** from the categories below
2. **Call the my tool** with the appropriate action 2. **Call the my tool** with the appropriate action
3. **If set**, warn the user before changing impactful settings such as the model or runtime limits 3. **If set**, warn the user before changing impactful settings (model, iterations)
4. **For detailed examples**, read [references/examples.md](references/examples.md) 4. **For detailed examples**, read [references/examples.md](references/examples.md)
## When to check ## When to check
+11
View File
@@ -15,6 +15,8 @@ Concrete scenarios showing when and how to use the my tool effectively.
``` ```
→ my(action="check", key="max_iterations") → my(action="check", key="max_iterations")
→ 40 → 40
→ my(action="check", key="_last_usage")
→ {"input_tokens": 62000, "output_tokens": 3000}
→ "I hit the iteration limit (40). The task was complex. I can ask the user if they want to increase it." → "I hit the iteration limit (40). The task was complex. I can ask the user if they want to increase it."
``` ```
@@ -64,3 +66,12 @@ Concrete scenarios showing when and how to use the my tool effectively.
→ my(action="set", key="test_framework", value="pytest") → my(action="set", key="test_framework", value="pytest")
→ my(action="set", key="has_docker", value=true) → my(action="set", key="has_docker", value=true)
``` ```
## Budget Awareness
### Token-conscious behavior
```
→ my(action="check", key="_last_usage")
→ {"input_tokens": 58000, "output_tokens": 12000}
→ "I've consumed ~70k tokens. I'll keep my remaining responses focused."
```
+18 -35
View File
@@ -1,42 +1,25 @@
Create a compact replacement checkpoint for this session. Create a memory overview for only the final {{ archive_count }} conversation messages immediately before this instruction. Earlier messages are context for resolving references; do not summarize them again.
When `[Archived Context Summary]` appears in the system prompt, update that previous checkpoint to reflect the current conversation state. Use [skip] unless a fact meets all SNIP criteria:
- Signal: would the user need to repeat this if forgotten?
- Novel: not just a restatement of another fact in this same conversation chunk
- Important: prevents rework or captures preferences / rules
- Persistent: still relevant after 2 weeks
## Merge rules Format each fact as:
- [mark] fact content
- Use the latest correction or decision as the current version of a fact, and merge duplicates. Marks (choose the best match):
- Preserve exact names, identifiers, paths, commands, decisions, results, and unresolved blockers when they are needed to continue the session. - [permanent] Core preferences, personal traits, habits — never becomes stale
- Retain a fact already present in long-term memory when it is needed for session continuity. - [durable] Technical discoveries, project knowledge, config details — valid for months
- [ephemeral] Active task state, temporary decisions — may change in weeks
- [correction] Correction to a previous memory — state what changed
- [skip] Conversational filler, code/source facts derivable from the repo, or audit-only breadcrumbs
## What to retain Priority: user corrections and preferences > solutions > decisions > events > environment facts.
Always retain a compact working-state handoff: Do not output facts already present in the system prompt's Recent History.
- active objective
- current status
- completed results that constrain later work
- unresolved blockers
- next action
- exact identifiers needed for that action
Mark working-state facts `[ephemeral]`. Do not mark something [skip] merely because it might already exist in long-term memory.
For other facts, retain a candidate only when it meets all four SNIP criteria: Return only formatted fact lines, or `(nothing)` if nothing noteworthy happened.
- Signal: remembering it saves the user from repeating it
- Novel: it adds a distinct fact to this checkpoint
- Important: losing it would cause rework or discard a preference or rule
- Persistent: it is expected to remain useful for at least two weeks
Assign each retained fact its best current mark:
- `[permanent]` for core preferences, personal traits, and habits that remain relevant indefinitely
- `[durable]` for technical discoveries, project knowledge, and configuration that remains valid for months
- `[ephemeral]` for active task state and temporary decisions that may change within weeks
- `[correction]` for the current fact that supersedes conflicting earlier long-term memory
When space is limited, prioritize user corrections and preferences, then solutions, decisions, events, and environment facts.
## Output
Return one concise retained fact per line in this form:
- [mark] fact
Use `(nothing)` when no fact qualifies and there is no active working state.
+12 -6
View File
@@ -18,9 +18,11 @@
## Discovery and Reading ## Discovery and Reading
- Use `find_files` or `list_dir` for uncertain paths, `grep` for content, and `read_file` for a known path. - Use `find_files` or `list_dir` to locate workspace paths before `read_file` when a path is uncertain.
- `grep` returns matches with five context lines by default; use `files_with_matches` for paths or `count` for totals. - Use `grep` for content search inside the workspace; prefer it over shell grep for ordinary searches.
- `grep` defaults to `output_mode="files_with_matches"`; use `output_mode="content"` for matching lines with context.
- Use `fixed_strings=true` for literal keywords containing regex characters. - Use `fixed_strings=true` for literal keywords containing regex characters.
- Use `output_mode="count"` to size a broad search before reading full matches.
- Use `head_limit` and `offset` to page across large result sets. - Use `head_limit` and `offset` to page across large result sets.
- Search tools enforce binary and file-size limits and report skipped files in the result. - Search tools enforce binary and file-size limits and report skipped files in the result.
@@ -40,15 +42,19 @@
result with its original consumer or checker when one is available. result with its original consumer or checker when one is available.
- Use `apply_patch` as the default code editing tool, especially for multi-file changes, structural edits, generated code, moves, adds, or deletes. - Use `apply_patch` as the default code editing tool, especially for multi-file changes, structural edits, generated code, moves, adds, or deletes.
- Use `apply_patch dry_run=true` when the patch is uncertain and you want validation plus a change summary before writing. - Use `apply_patch dry_run=true` when the patch is uncertain and you want validation plus a change summary before writing.
- Use `edit_file` only for small exact replacements in one file, with `old_text` copied from `read_file`. - Use `edit_file` only for small exact replacements in one file, with `old_text` copied from `read_file`; when editing a specific numbered line, pass that exact line as `line_hint`; add `occurrence` or `expected_replacements` when ambiguity matters.
- Use `write_file` for new files or intentional full-file rewrites, not routine partial edits. - Use `write_file` for new files or intentional full-file rewrites, not routine partial edits.
- If `apply_patch` or `edit_file` fails, re-read with `force=true`, narrow the context, and try a smaller patch rather than switching to shell `sed` or `echo`. - If `apply_patch` or `edit_file` fails, re-read with `force=true`, narrow the context, and try a smaller patch rather than switching to shell `sed` or `echo`.
## Process Execution ## Process Execution
- Use `exec` for processes, not file inspection or editing. - Use `exec` for tests, builds, package commands, git commands, and other process execution.
- For interaction or early output, set `yield_time_ms` and continue with `exec_session` (`until_exit=true` when no further input is needed). - Prefer dedicated file/search tools over `cat`, shell `find`, shell `grep`, `sed`, or `echo` for ordinary workspace inspection and edits.
- Use `list_exec_sessions` to recover session IDs. - Use non-interactive flags such as `-y` or `--yes` when available.
- Commands have a configurable timeout (default 60s), dangerous commands are blocked, and output is truncated.
- For long-running or interactive commands, pass `yield_time_ms`; if the process keeps running, continue with `write_stdin`.
- Use `write_stdin` to poll, provide stdin, close stdin, wait for expected output with `wait_for`, or terminate an existing exec session.
- Use `list_exec_sessions` to recover active session IDs after context shifts.
## CLI App Attachments ## CLI App Attachments
+215 -59
View File
@@ -5,17 +5,24 @@ from __future__ import annotations
import asyncio import asyncio
import uuid import uuid
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
from typing import Any from contextlib import suppress
from typing import Any, TypeVar
from loguru import logger from loguru import logger
from nanobot.agent.automation_turns import AutomationTurnError from nanobot.agent.automation_turns import (
AutomationTurnAcceptedCancellation,
AutomationTurnError,
)
from nanobot.bus.events import InboundMessage, OutboundMessage from nanobot.bus.events import InboundMessage, OutboundMessage
from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META from nanobot.triggers.local_session_turns import LOCAL_TRIGGER_META
from nanobot.triggers.local_store import LocalTriggerStore from nanobot.triggers.local_store import LocalTriggerStore
from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery from nanobot.triggers.local_types import LocalTrigger, TriggerDelivery
from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY from nanobot.webui.metadata import WEBUI_MESSAGE_SOURCE_METADATA_KEY, WEBUI_TURN_METADATA_KEY
_T = TypeVar("_T")
async def run_local_trigger_queue( async def run_local_trigger_queue(
*, *,
@@ -29,14 +36,16 @@ async def run_local_trigger_queue(
if submit_turn is None: if submit_turn is None:
raise ValueError("run_local_trigger_queue requires submit_turn") raise ValueError("run_local_trigger_queue requires submit_turn")
logger.info("Local trigger queue started") logger.info("Local trigger queue started")
recovered = store.recover_processing_deliveries() recovered = await shield_and_drain(asyncio.to_thread(store.recover_processing_deliveries))
if recovered: if recovered:
logger.warning( logger.warning(
"Trigger: recovered {} interrupted delivery file(s) from processing", "Trigger: recovered {} interrupted delivery file(s) from processing",
recovered, recovered,
) )
while True: while True:
deliveries = store.claim_deliveries(limit=batch_size) deliveries = await shield_and_drain(
asyncio.to_thread(store.claim_deliveries, limit=batch_size)
)
if not deliveries: if not deliveries:
await asyncio.sleep(poll_interval_s) await asyncio.sleep(poll_interval_s)
continue continue
@@ -49,30 +58,24 @@ async def run_local_trigger_queue(
submit_turn=submit_turn, submit_turn=submit_turn,
is_channel_enabled=is_channel_enabled, is_channel_enabled=is_channel_enabled,
) )
store.complete_delivery(delivery) except _DeliverySettledOnCancellation:
raise
except asyncio.CancelledError as exc: except asyncio.CancelledError as exc:
store.retry_delivery(delivery, str(exc) or exc.__class__.__name__) error = str(exc) or exc.__class__.__name__
_write_delivery_run_record( await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
await _write_delivery_run_record(
store, store,
delivery, delivery,
status="interrupted", status="interrupted",
error=str(exc) or exc.__class__.__name__, error=error,
) )
raise raise
except _TerminalDeliveryError as exc: except _TerminalDeliveryError as exc:
store.record_delivery( await _await_delivery_settlement(
delivery.trigger_id, _settle_failed_delivery(store, delivery, error=str(exc)),
status="error", store=store,
error=str(exc), delivery=delivery,
run_at_ms=delivery.created_at_ms,
) )
_write_delivery_run_record(
store,
delivery,
status="error",
error=str(exc),
)
store.complete_delivery(delivery)
logger.warning( logger.warning(
"Trigger: dropped delivery {} for {}: {}", "Trigger: dropped delivery {} for {}: {}",
delivery.id, delivery.id,
@@ -81,19 +84,11 @@ async def run_local_trigger_queue(
) )
except AutomationTurnError as exc: except AutomationTurnError as exc:
error = str(exc) or exc.__class__.__name__ error = str(exc) or exc.__class__.__name__
store.record_delivery( await _await_delivery_settlement(
delivery.trigger_id, _settle_failed_delivery(store, delivery, error=error),
status="error", store=store,
error=error, delivery=delivery,
run_at_ms=delivery.created_at_ms,
) )
_write_delivery_run_record(
store,
delivery,
status="error",
error=error,
)
store.complete_delivery(delivery)
logger.warning( logger.warning(
"Trigger: delivery {} for {} reached the agent but failed: {}", "Trigger: delivery {} for {} reached the agent but failed: {}",
delivery.id, delivery.id,
@@ -102,18 +97,10 @@ async def run_local_trigger_queue(
) )
except Exception as exc: except Exception as exc:
error = str(exc) or exc.__class__.__name__ error = str(exc) or exc.__class__.__name__
retried = store.retry_delivery(delivery, error) retried = await _await_delivery_settlement(
_write_delivery_run_record( _settle_retryable_delivery(store, delivery, error=error),
store, store=store,
delivery, delivery=delivery,
status="retrying" if retried else "error",
error=error,
)
store.record_delivery(
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
) )
logger.exception( logger.exception(
"Trigger: failed delivery {} for {}{}", "Trigger: failed delivery {} for {}{}",
@@ -127,6 +114,10 @@ class _TerminalDeliveryError(RuntimeError):
pass pass
class _DeliverySettledOnCancellation(asyncio.CancelledError):
"""Cancellation reported only after an already-submitted delivery is settled."""
async def _deliver_delivery( async def _deliver_delivery(
store: LocalTriggerStore, store: LocalTriggerStore,
delivery: TriggerDelivery, delivery: TriggerDelivery,
@@ -134,7 +125,7 @@ async def _deliver_delivery(
submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]], submit_turn: Callable[[InboundMessage], Awaitable[OutboundMessage | None]],
is_channel_enabled: Callable[[str], bool], is_channel_enabled: Callable[[str], bool],
) -> None: ) -> None:
trigger = store.get(delivery.trigger_id) trigger = await asyncio.to_thread(store.get, delivery.trigger_id)
if trigger is None: if trigger is None:
raise _TerminalDeliveryError("trigger not found") raise _TerminalDeliveryError("trigger not found")
if not trigger.enabled: if not trigger.enabled:
@@ -142,7 +133,14 @@ async def _deliver_delivery(
if not is_channel_enabled(trigger.channel): if not is_channel_enabled(trigger.channel):
raise _TerminalDeliveryError(f"target channel is not enabled: {trigger.channel}") raise _TerminalDeliveryError(f"target channel is not enabled: {trigger.channel}")
store.write_delivery_run_record(delivery, trigger=trigger, status="processing") await shield_and_drain(
asyncio.to_thread(
store.write_delivery_run_record,
delivery,
trigger=trigger,
status="processing",
)
)
msg = InboundMessage( msg = InboundMessage(
channel=trigger.channel, channel=trigger.channel,
sender_id=trigger.sender_id, sender_id=trigger.sender_id,
@@ -151,22 +149,177 @@ async def _deliver_delivery(
metadata=_delivery_metadata(trigger, delivery), metadata=_delivery_metadata(trigger, delivery),
session_key_override=trigger.session_key, session_key_override=trigger.session_key,
) )
response = await submit_turn(msg) try:
store.record_delivery( response = await submit_turn(msg)
trigger.id, except AutomationTurnAcceptedCancellation:
status="ok", try:
run_at_ms=delivery.created_at_ms, await _await_delivery_settlement(
_settle_accepted_delivery(store, delivery, trigger=trigger),
store=store,
delivery=delivery,
)
except _DeliverySettledOnCancellation:
raise
except Exception:
logger.exception(
"Trigger: failed to persist accepted delivery {}; dropping retry",
delivery.id,
)
with suppress(Exception):
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
raise _DeliverySettledOnCancellation from None
try:
await _await_delivery_settlement(
_settle_submitted_delivery(store, delivery, trigger=trigger, response=response),
store=store,
delivery=delivery,
)
except Exception:
logger.exception(
"Trigger: failed to persist status for submitted delivery {}; dropping retry",
delivery.id,
)
with suppress(Exception):
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
async def _await_delivery_settlement(
operation: Awaitable[_T],
*,
store: LocalTriggerStore,
delivery: TriggerDelivery,
) -> _T:
settlement = asyncio.ensure_future(operation)
try:
return await asyncio.shield(settlement)
except asyncio.CancelledError:
while not settlement.done():
try:
await asyncio.shield(settlement)
except asyncio.CancelledError:
continue
try:
settlement.result()
except Exception:
logger.exception(
"Trigger: failed to settle delivery {} during cancellation",
delivery.id,
)
completion = asyncio.create_task(
shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
)
while not completion.done():
try:
await asyncio.shield(completion)
except asyncio.CancelledError:
continue
with suppress(Exception):
completion.result()
raise _DeliverySettledOnCancellation from None
async def _settle_failed_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
error: str,
) -> None:
await _write_delivery_run_record(
store,
delivery,
status="error",
error=error,
) )
_write_delivery_run_record( await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
# Publish the terminal status only after the durable delivery state is settled.
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
)
)
async def _settle_retryable_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
error: str,
) -> bool:
retried = await shield_and_drain(asyncio.to_thread(store.retry_delivery, delivery, error))
await _write_delivery_run_record(
store,
delivery,
status="retrying" if retried else "error",
error=error,
)
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
delivery.trigger_id,
status="error",
error=error,
run_at_ms=delivery.created_at_ms,
)
)
return retried
async def _settle_accepted_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
trigger: LocalTrigger,
) -> None:
"""Commit an accepted delivery without claiming the agent turn completed."""
await _write_delivery_run_record(
store,
delivery,
trigger=trigger,
status="accepted",
)
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
trigger.id,
status="ok",
run_at_ms=delivery.created_at_ms,
)
)
async def _settle_submitted_delivery(
store: LocalTriggerStore,
delivery: TriggerDelivery,
*,
trigger: LocalTrigger,
response: OutboundMessage | None,
) -> None:
await _write_delivery_run_record(
store, store,
delivery, delivery,
trigger=trigger, trigger=trigger,
status="ok", status="ok",
response=response.content if response else "", response=response.content if response else "",
) )
await shield_and_drain(asyncio.to_thread(store.complete_delivery, delivery))
# last_status is the externally visible commit marker for a settled delivery.
await shield_and_drain(
asyncio.to_thread(
store.record_delivery,
trigger.id,
status="ok",
run_at_ms=delivery.created_at_ms,
)
)
def _write_delivery_run_record( async def _write_delivery_run_record(
store: LocalTriggerStore, store: LocalTriggerStore,
delivery: TriggerDelivery, delivery: TriggerDelivery,
*, *,
@@ -176,12 +329,15 @@ def _write_delivery_run_record(
response: str | None = None, response: str | None = None,
) -> None: ) -> None:
try: try:
store.write_delivery_run_record( await shield_and_drain(
delivery, asyncio.to_thread(
trigger=trigger, store.write_delivery_run_record,
status=status, delivery,
error=error, trigger=trigger,
response=response, status=status,
error=error,
response=response,
)
) )
except Exception: except Exception:
logger.exception( logger.exception(
+5 -1
View File
@@ -24,6 +24,7 @@ _MAX_RUN_HISTORY = 20
_MAX_DELIVERY_ATTEMPTS = 10 _MAX_DELIVERY_ATTEMPTS = 10
_RUN_RECORD_TEXT_MAX_CHARS = 4000 _RUN_RECORD_TEXT_MAX_CHARS = 4000
_PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing" _PROCESSING_RECOVERY_ERROR = "delivery was recovered from interrupted processing"
_FILE_LOCK_TIMEOUT_SECONDS = 5
class TriggerStoreError(RuntimeError): class TriggerStoreError(RuntimeError):
@@ -49,7 +50,10 @@ class LocalTriggerStore:
self.processing_dir = self.root / "processing" self.processing_dir = self.root / "processing"
self.failed_dir = self.root / "failed" self.failed_dir = self.root / "failed"
self.runs_dir = self.root / "runs" self.runs_dir = self.root / "runs"
self._lock = FileLock(str(self.root / ".lock")) self._lock = FileLock(
str(self.root / ".lock"),
timeout=_FILE_LOCK_TIMEOUT_SECONDS,
)
def create( def create(
self, self,
+40
View File
@@ -3,8 +3,48 @@
from __future__ import annotations from __future__ import annotations
import asyncio import asyncio
from collections.abc import Awaitable
from typing import TypeVar
_T = TypeVar("_T")
def task_is_cancelling() -> bool: def task_is_cancelling() -> bool:
task = asyncio.current_task() task = asyncio.current_task()
return task is not None and task.cancelling() > 0 return task is not None and task.cancelling() > 0
async def shield_and_drain(awaitable: Awaitable[_T]) -> _T:
"""Delay caller cancellation until an accepted operation has fully settled.
``asyncio.to_thread`` cannot stop a worker that has already started. Shielding
keeps cancellation from detaching that worker, and draining also lets any
post-write in-memory settlement in ``awaitable`` finish. Cancellation is still
re-raised as soon as the accepted operation is done.
"""
settlement = asyncio.ensure_future(awaitable)
cancellation: asyncio.CancelledError | None = None
while not settlement.done():
try:
result = await asyncio.shield(settlement)
except asyncio.CancelledError as exc:
if cancellation is None:
cancellation = exc
except BaseException:
if cancellation is None:
raise
break
else:
if cancellation is not None:
raise cancellation
return result
if cancellation is not None:
try:
settlement.result()
except BaseException:
# The caller's cancellation wins once settlement has been observed.
pass
raise cancellation
return settlement.result()
+117 -360
View File
@@ -66,10 +66,6 @@ class DocxSafetyError(Exception):
"""Raised when a DOCX table exceeds a parser safety boundary.""" """Raised when a DOCX table exceeds a parser safety boundary."""
class DocumentExtractionError(Exception):
"""Raised when a document cannot be opened for incremental extraction."""
@dataclass(frozen=True, slots=True) @dataclass(frozen=True, slots=True)
class PdfExtraction: class PdfExtraction:
text: str text: str
@@ -78,24 +74,6 @@ class PdfExtraction:
end_page: int end_page: int
@dataclass(frozen=True, slots=True)
class LocatedDocumentLine:
"""One searchable document line with a stable, human-readable locator."""
text: str
extracted_line: int
locator: str
searchable: bool = True
@dataclass(frozen=True, slots=True)
class DocumentLineSource:
"""Incremental document lines plus an optional next PDF page range."""
lines: Iterator[LocatedDocumentLine]
continuation: str | None = None
def extract_text(path: str | Path) -> str | None: def extract_text(path: str | Path) -> str | None:
"""Extract text from a file. """Extract text from a file.
@@ -107,8 +85,13 @@ def extract_text(path: str | Path) -> str | None:
or error string for failures. or error string for failures.
""" """
path = Path(path) path = Path(path)
if error := _extraction_path_error(path): if not path.exists():
return error return f"[error: file not found: {path}]"
try:
if path.stat().st_size > _MAX_EXTRACT_FILE_SIZE:
return f"[error: file exceeds {_MAX_EXTRACT_FILE_SIZE // (1024 * 1024)} MB limit]"
except OSError as e:
return f"[error: failed to inspect file: {e!s}]"
ext = path.suffix.lower() ext = path.suffix.lower()
@@ -132,303 +115,6 @@ def extract_text(path: str | Path) -> str | None:
return None return None
def open_document_line_source(
path: str | Path,
*,
pages: str | None = None,
) -> DocumentLineSource | None:
"""Open a document as an incremental stream of extracted lines.
Unlike :func:`extract_text`, this interface does not apply the attachment
text preview limit. Parser/file safety limits still apply. Lines that are
useful only for the rendered document view (for example sheet headers and
blank separators) have ``searchable=False`` so range reads can retain them
without making grep match synthetic text.
"""
path = Path(path)
ext = path.suffix.lower()
if ext not in {".pdf", ".docx", ".xlsx", ".pptx"}:
return None
if error := _extraction_path_error(path):
raise DocumentExtractionError(_clean_extraction_error(error))
if ext == ".pdf":
return _open_pdf_line_source(path, pages)
if ext == ".docx":
return _open_docx_line_source(path)
if ext == ".xlsx":
return _open_xlsx_line_source(path)
return _open_pptx_line_source(path)
def _clean_extraction_error(error: str) -> str:
if error.startswith("[error:") and error.endswith("]"):
return error[len("[error:") : -1].strip()
return error
def _check_office_archive(path: Path) -> None:
if error := _office_archive_error(path):
raise DocumentExtractionError(_clean_extraction_error(error))
def _open_pdf_line_source(path: Path, pages: str | None) -> DocumentLineSource:
try:
from pypdf import PdfReader
reader = PdfReader(path, strict=False)
total_pages = len(reader.pages)
if total_pages == 0:
return DocumentLineSource(iter(()))
start, requested_end = _parse_pdf_page_range(pages, total_pages)
except PdfPageRangeError:
raise
except Exception as e:
raise DocumentExtractionError(f"failed to open PDF: {e!s}") from e
end = min(requested_end, start + _MAX_PDF_ATTACHMENT_PAGES - 1)
continuation = None
if end < total_pages - 1:
next_start = end + 2
next_end = min(end + 1 + _MAX_PDF_ATTACHMENT_PAGES, total_pages)
continuation = f"pages='{next_start}-{next_end}'"
def iter_lines() -> Iterator[LocatedDocumentLine]:
extracted_line = 0
wrote_page = False
for index in range(start, end + 1):
page = reader.pages[index]
contents = page.get_contents()
if contents is not None:
stream_size = len(contents.get_data())
if stream_size > _MAX_PDF_CONTENT_STREAM_SIZE:
raise PdfSafetyError(
f"page {index + 1} content stream exceeds "
f"{_MAX_PDF_CONTENT_STREAM_SIZE // (1024 * 1024)} MB limit"
)
text = (page.extract_text() or "").strip()
if not text:
continue
if wrote_page:
extracted_line += 1
yield LocatedDocumentLine("", extracted_line, "", searchable=False)
extracted_line += 1
yield LocatedDocumentLine(
f"--- Page {index + 1} ---",
extracted_line,
"",
searchable=False,
)
page_line = 0
for text_line in text.splitlines():
extracted_line += 1
if not text_line:
yield LocatedDocumentLine("", extracted_line, "", searchable=False)
continue
page_line += 1
yield LocatedDocumentLine(
text_line,
extracted_line,
f"page={index + 1},line={page_line}",
)
wrote_page = True
return DocumentLineSource(iter_lines(), continuation=continuation)
def _open_xlsx_line_source(path: Path) -> DocumentLineSource:
_check_office_archive(path)
try:
from openpyxl import load_workbook
except ImportError as e:
raise DocumentExtractionError("openpyxl not installed") from e
try:
workbook = load_workbook(path, read_only=True, data_only=True)
except Exception as e:
raise DocumentExtractionError(f"failed to open XLSX: {e!s}") from e
def iter_lines() -> Iterator[LocatedDocumentLine]:
extracted_line = 0
wrote_document_content = False
try:
for sheet_name in workbook.sheetnames:
worksheet = workbook[sheet_name]
wrote_header = False
for row_index, row in enumerate(worksheet.iter_rows(values_only=True), 1):
row_text = "\t".join(
str(cell) if cell is not None else "" for cell in row
)
if not row_text.strip():
continue
if not wrote_header:
if wrote_document_content:
extracted_line += 1
yield LocatedDocumentLine(
"", extracted_line, "", searchable=False
)
extracted_line += 1
yield LocatedDocumentLine(
f"--- Sheet: {sheet_name} ---",
extracted_line,
"",
searchable=False,
)
wrote_header = True
wrote_document_content = True
extracted_line += 1
yield LocatedDocumentLine(
row_text,
extracted_line,
f"sheet={sheet_name!r},row={row_index}",
)
finally:
workbook.close()
return DocumentLineSource(iter_lines())
def _open_pptx_line_source(path: Path) -> DocumentLineSource:
_check_office_archive(path)
try:
from pptx import Presentation as PptxPresentation
except ImportError as e:
raise DocumentExtractionError("python-pptx not installed") from e
try:
presentation = PptxPresentation(str(path))
except Exception as e:
raise DocumentExtractionError(f"failed to open PPTX: {e!s}") from e
def iter_lines() -> Iterator[LocatedDocumentLine]:
extracted_line = 0
wrote_slide = False
for slide_number, slide in enumerate(presentation.slides, 1):
slide_text: list[str] = []
for shape in slide.shapes:
_collect_pptx_shape_text(shape, slide_text)
rendered_lines = [line for text in slide_text for line in text.splitlines()]
if not rendered_lines:
continue
if wrote_slide:
extracted_line += 1
yield LocatedDocumentLine("", extracted_line, "", searchable=False)
extracted_line += 1
yield LocatedDocumentLine(
f"--- Slide {slide_number} ---",
extracted_line,
"",
searchable=False,
)
slide_line = 0
for text_line in rendered_lines:
extracted_line += 1
if not text_line:
yield LocatedDocumentLine("", extracted_line, "", searchable=False)
continue
slide_line += 1
yield LocatedDocumentLine(
text_line,
extracted_line,
f"slide={slide_number},line={slide_line}",
)
wrote_slide = True
return DocumentLineSource(iter_lines())
def _open_docx_line_source(path: Path) -> DocumentLineSource:
_check_office_archive(path)
try:
from docx import Document as DocxDocument
from docx.table import Table, _Cell # pyright: ignore[reportPrivateUsage]
from docx.text.paragraph import Paragraph
except ImportError as e:
raise DocumentExtractionError("python-docx not installed") from e
try:
document = DocxDocument(str(path))
except Exception as e:
raise DocumentExtractionError(f"failed to open DOCX: {e!s}") from e
def iter_lines() -> Iterator[LocatedDocumentLine]:
table_cell_count = 0
def cell_text(cell: _Cell, depth: int) -> str:
parts: list[str] = []
for block in cell.iter_inner_content():
if isinstance(block, Paragraph):
text = " ".join(block.text.split())
if text:
parts.append(text)
elif isinstance(block, Table): # pyright: ignore[reportUnnecessaryIsInstance]
parts.extend(
row.replace("\t", " | ") for row in table_rows(block, depth + 1)
)
return " ".join(parts)
def table_rows(table: Table, depth: int) -> Iterator[str]:
nonlocal table_cell_count
if depth > _MAX_DOCX_TABLE_DEPTH:
raise DocxSafetyError(
f"table nesting exceeds {_MAX_DOCX_TABLE_DEPTH} levels"
)
for row in table.rows:
cells: list[str] = []
for tc in row._tr.tc_lst: # pyright: ignore[reportPrivateUsage]
table_cell_count += 1
if table_cell_count > _MAX_DOCX_TABLE_CELLS:
raise DocxSafetyError(
f"document contains more than {_MAX_DOCX_TABLE_CELLS} table cells"
)
cells.append(cell_text(_Cell(tc, table), depth))
if any(cells):
yield "\t".join(cells)
def blocks() -> Iterator[tuple[str, bool]]:
for block in document.iter_inner_content():
if isinstance(block, Paragraph):
text = block.text.strip()
if text:
yield text, True
continue
if not isinstance(block, Table): # pyright: ignore[reportUnnecessaryIsInstance]
continue
first_row = True
for row_text in table_rows(block, 1):
yield row_text, first_row
first_row = False
extracted_line = 0
paragraph = 0
wrote_content = False
for text, separate in blocks():
if wrote_content and separate:
extracted_line += 1
yield LocatedDocumentLine("", extracted_line, "", searchable=False)
for text_line in text.splitlines():
extracted_line += 1
if not text_line:
yield LocatedDocumentLine("", extracted_line, "", searchable=False)
continue
paragraph += 1
yield LocatedDocumentLine(
text_line,
extracted_line,
f"paragraph={paragraph}",
)
wrote_content = True
return DocumentLineSource(iter_lines())
def _extraction_path_error(path: Path) -> str | None:
if not path.exists():
return f"[error: file not found: {path}]"
try:
if path.stat().st_size > _MAX_EXTRACT_FILE_SIZE:
return f"[error: file exceeds {_MAX_EXTRACT_FILE_SIZE // (1024 * 1024)} MB limit]"
except OSError as e:
return f"[error: failed to inspect file: {e!s}]"
return None
def _extract_pdf(path: Path) -> str: def _extract_pdf(path: Path) -> str:
"""Extract text from PDF using pypdf.""" """Extract text from PDF using pypdf."""
try: try:
@@ -484,73 +170,144 @@ def extract_pdf_pages(
def _parse_pdf_page_range(pages: str | None, total_pages: int) -> tuple[int, int]: def _parse_pdf_page_range(pages: str | None, total_pages: int) -> tuple[int, int]:
if not pages: if not pages:
return 0, total_pages - 1 return 0, total_pages - 1
page_word = "page" if total_pages == 1 else "pages"
guidance = (
f"document has {total_pages} {page_word}; "
f"use a page number or range within 1-{total_pages}"
)
values = pages.strip().split("-") values = pages.strip().split("-")
if len(values) not in {1, 2}: if len(values) not in {1, 2}:
raise PdfPageRangeError(guidance) raise PdfPageRangeError(f"invalid page range: {pages}")
try: try:
start = int(values[0]) start = int(values[0])
end = int(values[-1]) end = int(values[-1])
except ValueError as e: except ValueError as e:
raise PdfPageRangeError(guidance) from e raise PdfPageRangeError(f"invalid page range: {pages}") from e
if start < 1 or end < start or start > total_pages: if start < 1 or end < start or start > total_pages:
raise PdfPageRangeError(guidance) raise PdfPageRangeError(f"invalid page range: {pages}")
return start - 1, min(end, total_pages) - 1 return start - 1, min(end, total_pages) - 1
def _render_document_preview(source: DocumentLineSource) -> str:
"""Render a bounded attachment preview from the canonical line stream."""
collector = _TextCollector(_MAX_TEXT_LENGTH)
iterator = source.lines
first_line = True
try:
for line in iterator:
if not first_line and not collector.add("\n"):
break
first_line = False
if line.text and not collector.add(line.text):
break
return collector.render()
finally:
close = getattr(iterator, "close", None)
if close is not None:
close()
def _extract_docx(path: Path) -> str: def _extract_docx(path: Path) -> str:
"""Extract a bounded DOCX attachment preview.""" """Extract text from DOCX using python-docx."""
try: try:
return _render_document_preview(_open_docx_line_source(path)) from docx import Document as DocxDocument
from docx.table import Table, _Cell # pyright: ignore[reportPrivateUsage]
from docx.text.paragraph import Paragraph
except ImportError:
return "[error: python-docx not installed]"
try:
if error := _office_archive_error(path):
return error
doc = DocxDocument(str(path))
collector = _TextCollector(_MAX_TEXT_LENGTH)
table_cell_count = 0
def cell_text(cell: _Cell, depth: int) -> str:
parts: list[str] = []
for block in cell.iter_inner_content():
if isinstance(block, Paragraph):
text = " ".join(block.text.split())
if text:
parts.append(text)
elif isinstance(block, Table): # pyright: ignore[reportUnnecessaryIsInstance]
parts.extend(row.replace("\t", " | ") for row in table_rows(block, depth + 1))
return " ".join(parts)
def table_rows(table: Table, depth: int) -> Iterator[str]:
nonlocal table_cell_count
if depth > _MAX_DOCX_TABLE_DEPTH:
raise DocxSafetyError(
f"table nesting exceeds {_MAX_DOCX_TABLE_DEPTH} levels"
)
for row in table.rows:
cells: list[str] = []
# row.cells expands w:gridSpan before callers can apply a bound.
# Physical w:tc elements keep malformed documents proportional to XML size.
for tc in row._tr.tc_lst: # pyright: ignore[reportPrivateUsage]
table_cell_count += 1
if table_cell_count > _MAX_DOCX_TABLE_CELLS:
raise DocxSafetyError(
f"document contains more than {_MAX_DOCX_TABLE_CELLS} table cells"
)
cells.append(cell_text(_Cell(tc, table), depth))
if any(cells):
yield "\t".join(cells)
for block in doc.iter_inner_content():
if isinstance(block, Paragraph):
text = block.text.strip()
if text and not collector.add(text, separator="\n\n"):
break
continue
if not isinstance(block, Table): # pyright: ignore[reportUnnecessaryIsInstance]
continue
first_row = True
for row_text in table_rows(block, 1):
separator = "\n\n" if first_row else "\n"
first_row = False
if not collector.add(row_text, separator=separator):
return collector.render()
return collector.render()
except DocxSafetyError as e: except DocxSafetyError as e:
return f"[error: unsafe DOCX: {e!s}]" return f"[error: unsafe DOCX: {e!s}]"
except DocumentExtractionError as e:
return f"[error: {e!s}]"
except Exception as e: except Exception as e:
logger.exception("Failed to extract DOCX {}", path) logger.exception("Failed to extract DOCX {}", path)
return f"[error: failed to extract DOCX: {e!s}]" return f"[error: failed to extract DOCX: {e!s}]"
def _extract_xlsx(path: Path) -> str: def _extract_xlsx(path: Path) -> str:
"""Extract a bounded XLSX attachment preview.""" """Extract text from XLSX using openpyxl."""
try: try:
return _render_document_preview(_open_xlsx_line_source(path)) from openpyxl import load_workbook
except DocumentExtractionError as e: except ImportError:
return f"[error: {e!s}]" return "[error: openpyxl not installed]"
try:
if error := _office_archive_error(path):
return error
wb = load_workbook(path, read_only=True, data_only=True)
try:
collector = _TextCollector(_MAX_TEXT_LENGTH)
for sheet_name in wb.sheetnames:
ws = wb[sheet_name]
wrote_header = False
for row in ws.iter_rows(values_only=True):
row_text = "\t".join(str(cell) if cell is not None else "" for cell in row)
if row_text.strip():
if not wrote_header:
if not collector.add(
f"--- Sheet: {sheet_name} ---",
separator="\n\n",
):
return collector.render()
wrote_header = True
if not collector.add(row_text, separator="\n"):
return collector.render()
return collector.render()
finally:
wb.close()
except Exception as e: except Exception as e:
logger.exception("Failed to extract XLSX {}", path) logger.exception("Failed to extract XLSX {}", path)
return f"[error: failed to extract XLSX: {e!s}]" return f"[error: failed to extract XLSX: {e!s}]"
def _extract_pptx(path: Path) -> str: def _extract_pptx(path: Path) -> str:
"""Extract a bounded PPTX attachment preview.""" """Extract text from PPTX using python-pptx."""
try: try:
return _render_document_preview(_open_pptx_line_source(path)) from pptx import Presentation as PptxPresentation
except DocumentExtractionError as e: except ImportError:
return f"[error: {e!s}]" return "[error: python-pptx not installed]"
try:
if error := _office_archive_error(path):
return error
prs = PptxPresentation(str(path))
collector = _TextCollector(_MAX_TEXT_LENGTH)
for i, slide in enumerate(prs.slides, 1):
slide_text: list[str] = []
for shape in slide.shapes:
_collect_pptx_shape_text(shape, slide_text)
if slide_text:
if not collector.add(
f"--- Slide {i} ---\n" + "\n".join(slide_text),
separator="\n\n",
):
break
return collector.render()
except Exception as e: except Exception as e:
logger.exception("Failed to extract PPTX {}", path) logger.exception("Failed to extract PPTX {}", path)
return f"[error: failed to extract PPTX: {e!s}]" return f"[error: failed to extract PPTX: {e!s}]"
+5 -5
View File
@@ -133,13 +133,12 @@ class GitStore:
try: try:
from dulwich import porcelain from dulwich import porcelain
# Stage first so Dulwich refreshes the content hashes. A status # .gitignore excludes everything except tracked files,
# check can miss rapid same-size rewrites when the filesystem also # so any staged/unstaged change must be in our files.
# preserves the file's mtime.
porcelain.add(str(self._workspace), paths=self._staging_paths(*self._tracked_files))
st = porcelain.status(str(self._workspace)) st = porcelain.status(str(self._workspace))
unstaged = cast(list[object], st.unstaged)
staged = cast(dict[object, list[object]], st.staged) staged = cast(dict[object, list[object]], st.staged)
if not any(staged.values()): if not unstaged and not any(staged.values()):
return None return None
message_value = cast(object, message) message_value = cast(object, message)
@@ -148,6 +147,7 @@ class GitStore:
if isinstance(message_value, str) if isinstance(message_value, str)
else cast(bytes, message_value) else cast(bytes, message_value)
) )
porcelain.add(str(self._workspace), paths=self._staging_paths(*self._tracked_files))
sha_bytes = porcelain.commit( sha_bytes = porcelain.commit(
str(self._workspace), str(self._workspace),
message=msg_bytes, message=msg_bytes,
+12
View File
@@ -40,6 +40,13 @@ LENGTH_RECOVERY_PROMPT = (
"existing text, recap, or apologize." "existing text, recap, or apologize."
) )
SUSTAINED_GOAL_CONTINUE_PROMPT = (
"You have an active sustained goal. Please continue working toward the "
"objective using your tools, or call update_goal with action='complete' "
"if the work is truly finished."
)
def empty_tool_result_message(tool_name: str) -> str: def empty_tool_result_message(tool_name: str) -> str:
"""Short prompt-safe marker for tools that completed without visible output.""" """Short prompt-safe marker for tools that completed without visible output."""
return f"({tool_name} completed with no output)" return f"({tool_name} completed with no output)"
@@ -90,6 +97,11 @@ def build_length_recovery_message(content: str) -> dict[str, str]:
return {"role": "user", "content": prompt} return {"role": "user", "content": prompt}
def build_goal_continue_message(custom: str | None = None) -> dict[str, str]:
"""Prompt the model to continue when a sustained goal is still active."""
return {"role": "user", "content": custom or SUSTAINED_GOAL_CONTINUE_PROMPT}
def external_lookup_signature(tool_name: str, arguments: Any) -> str | None: def external_lookup_signature(tool_name: str, arguments: Any) -> str | None:
"""Stable signature for repeated external lookups we want to throttle.""" """Stable signature for repeated external lookups we want to throttle."""
if not isinstance(arguments, dict): if not isinstance(arguments, dict):
+43 -44
View File
@@ -2,15 +2,15 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import re
import uuid import uuid
from collections.abc import Mapping from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Protocol from typing import TYPE_CHECKING, Any, TypeGuard
from loguru import logger
from nanobot.session.manager import SessionManager from nanobot.session.manager import SessionManager
from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title from nanobot.session.webui_turns import WEBUI_TITLE_METADATA_KEY, clean_generated_title
from nanobot.webui.session_identity import is_valid_webui_chat_id, webui_session_key from nanobot.utils.cancellation import shield_and_drain
from nanobot.webui.transcript import ( from nanobot.webui.transcript import (
append_fork_marker, append_fork_marker,
delete_webui_transcript, delete_webui_transcript,
@@ -21,25 +21,13 @@ from nanobot.webui.transcript import (
if TYPE_CHECKING: if TYPE_CHECKING:
from websockets.asyncio.server import ServerConnection from websockets.asyncio.server import ServerConnection
from nanobot.webui.gateway_services import GatewayServices from nanobot.channels.websocket.runtime import WebSocketChannel
_WEBUI_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$")
class WebUIForkHost(Protocol): def _valid_webui_chat_id(value: Any) -> TypeGuard[str]:
gateway: GatewayServices return isinstance(value, str) and _WEBUI_CHAT_ID_RE.match(value) is not None
async def send_webui_protocol_error(
self,
connection: ServerConnection,
detail: str,
) -> None: ...
async def attach_webui_fork(
self,
connection: ServerConnection,
*,
fork_id: str,
fork_key: str,
) -> None: ...
def create_webui_chat_fork( def create_webui_chat_fork(
@@ -51,8 +39,8 @@ def create_webui_chat_fork(
) -> tuple[str, str] | None: ) -> tuple[str, str] | None:
"""Return ``(chat_id, session_key)`` for a new fork, or ``None`` for bad input.""" """Return ``(chat_id, session_key)`` for a new fork, or ``None`` for bad input."""
new_id = str(uuid.uuid4()) new_id = str(uuid.uuid4())
source_key = webui_session_key(source_chat_id) source_key = f"websocket:{source_chat_id}"
target_key = webui_session_key(new_id) target_key = f"websocket:{new_id}"
try: try:
forked = session_manager.fork_session_before_user_index( forked = session_manager.fork_session_before_user_index(
source_key, source_key,
@@ -83,7 +71,7 @@ def create_webui_chat_fork(
async def handle_webui_fork_chat( async def handle_webui_fork_chat(
channel: WebUIForkHost, channel: WebSocketChannel,
connection: ServerConnection, connection: ServerConnection,
envelope: Mapping[str, Any], envelope: Mapping[str, Any],
) -> None: ) -> None:
@@ -95,7 +83,7 @@ async def handle_webui_fork_chat(
""" """
source_chat_id = envelope.get("source_chat_id") source_chat_id = envelope.get("source_chat_id")
raw_index = envelope.get("before_user_index") raw_index = envelope.get("before_user_index")
if not is_valid_webui_chat_id(source_chat_id): if not _valid_webui_chat_id(source_chat_id):
await channel.send_webui_protocol_error(connection, "invalid source_chat_id") await channel.send_webui_protocol_error(connection, "invalid source_chat_id")
return return
if isinstance(raw_index, bool) or not isinstance(raw_index, int) or raw_index < 0: if isinstance(raw_index, bool) or not isinstance(raw_index, int) or raw_index < 0:
@@ -107,24 +95,35 @@ async def handle_webui_fork_chat(
await channel.send_webui_protocol_error(connection, "session_manager_unavailable") await channel.send_webui_protocol_error(connection, "session_manager_unavailable")
return return
try: async def create_and_attach() -> None:
forked = create_webui_chat_fork( try:
session_manager, forked = await asyncio.to_thread(
source_chat_id=source_chat_id, create_webui_chat_fork,
before_user_index=raw_index, session_manager,
title=envelope.get("title") if isinstance(envelope.get("title"), str) else None, source_chat_id=source_chat_id,
) before_user_index=raw_index,
if forked is None: title=(
await channel.send_webui_protocol_error(connection, "invalid fork source or index") envelope.get("title")
if isinstance(envelope.get("title"), str)
else None
),
)
if forked is None:
await channel.send_webui_protocol_error(
connection,
"invalid fork source or index",
)
return
fork_id, fork_key = forked
except Exception as exc:
channel.logger.warning("fork_chat failed: {}", exc)
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
return return
fork_id, fork_key = forked
except Exception as exc:
logger.warning("fork_chat failed: {}", exc)
await channel.send_webui_protocol_error(connection, "fork_chat_failed")
return
await channel.attach_webui_fork( await channel.attach_webui_fork(
connection, connection,
fork_id=fork_id, fork_id=fork_id,
fork_key=fork_key, fork_key=fork_key,
) )
await shield_and_drain(create_and_attach())
-114
View File
@@ -1,114 +0,0 @@
"""HTTP and handshake composition for the WebUI gateway listener."""
from __future__ import annotations
import hmac
from collections.abc import Callable
from typing import TYPE_CHECKING, Any
from websockets.asyncio.server import ServerConnection
from websockets.http11 import Request as WsRequest
from nanobot.webui.gateway_tokens import GatewayTokenStore
from nanobot.webui.http_utils import (
is_trusted_proxy_authenticated_request,
normalize_config_path,
parse_request_path,
query_first,
)
from nanobot.webui.ws_http import GatewayHTTPHandler
if TYPE_CHECKING:
from nanobot.channels.websocket.runtime import WebSocketConfig
def is_websocket_upgrade(request: WsRequest) -> bool:
"""Return whether a request contains a complete WebSocket upgrade handshake."""
upgrade = request.headers.get("Upgrade") or request.headers.get("upgrade")
connection = request.headers.get("Connection") or request.headers.get("connection")
return bool(
upgrade
and "websocket" in upgrade.lower()
and connection
and "upgrade" in connection.lower()
)
class WebUIGatewayEndpoint:
"""Compose HTTP routing and WebSocket authentication on one listener."""
def __init__(
self,
*,
config: WebSocketConfig,
http: GatewayHTTPHandler,
tokens: GatewayTokenStore,
) -> None:
self._config = config
self._http = http
self._tokens = tokens
self.webui_connections: set[ServerConnection] = set()
async def process_request(
self,
connection: ServerConnection,
request: WsRequest,
*,
is_allowed: Callable[[str], bool],
) -> Any:
"""Route one listener request to a WS handshake or the HTTP application."""
got, query = parse_request_path(request.path)
expected_ws = normalize_config_path(self._config.path)
if got == expected_ws and is_websocket_upgrade(request):
client_id = query_first(query, "client_id") or ""
if len(client_id) > 128:
client_id = client_id[:128]
if not is_allowed(client_id):
return connection.respond(403, "Forbidden")
return self.authorize_websocket_handshake(connection, query, request.headers)
return await self._http.dispatch(connection, request)
def authorize_websocket_handshake(
self,
connection: ServerConnection,
query: dict[str, list[str]],
headers: Any = None,
) -> Any:
"""Authorize a WebSocket upgrade and remember trusted WebUI connections."""
if is_trusted_proxy_authenticated_request(connection, headers or {}, self._config):
self.webui_connections.add(connection)
return None
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.consume_issued_token(connection, supplied):
return None
return connection.respond(401, "Unauthorized")
if self._config.websocket_requires_token:
if supplied and self.consume_issued_token(connection, supplied):
return None
return connection.respond(401, "Unauthorized")
if supplied:
self.consume_issued_token(connection, supplied)
return None
def consume_issued_token(self, connection: ServerConnection, token: str) -> bool:
"""Consume one issued token and record its WebUI audience when present."""
audience = self._tokens.take_issued_token_audience(token)
if audience == "webui":
self.webui_connections.add(connection)
return audience is not None
def is_webui_connection(self, connection: ServerConnection) -> bool:
return connection in self.webui_connections
def discard_connection(self, connection: ServerConnection) -> None:
self.webui_connections.discard(connection)
def clear(self) -> None:
self.webui_connections.clear()
-8
View File
@@ -10,11 +10,9 @@ from typing import TYPE_CHECKING, Any, Callable
from loguru import logger as default_logger from loguru import logger as default_logger
from nanobot.config.loader import get_config_path from nanobot.config.loader import get_config_path
from nanobot.webui.gateway_endpoint import WebUIGatewayEndpoint
from nanobot.webui.gateway_tokens import GatewayTokenStore from nanobot.webui.gateway_tokens import GatewayTokenStore
from nanobot.webui.ingress_policy import DEFAULT_WEBUI_INGRESS_POLICY, WebUIIngressPolicy from nanobot.webui.ingress_policy import DEFAULT_WEBUI_INGRESS_POLICY, WebUIIngressPolicy
from nanobot.webui.media_gateway import WebUIMediaGateway from nanobot.webui.media_gateway import WebUIMediaGateway
from nanobot.webui.session_projection import WebUISessionProjection
from nanobot.webui.settings_services import WebUISettingsServices from nanobot.webui.settings_services import WebUISettingsServices
from nanobot.webui.temporary_chats import WebUITemporaryChats from nanobot.webui.temporary_chats import WebUITemporaryChats
from nanobot.webui.transcript import WebUITranscriptRecorder from nanobot.webui.transcript import WebUITranscriptRecorder
@@ -34,7 +32,6 @@ class GatewayServices:
"""Explicit dependencies shared by WebSocket transport and HTTP routes.""" """Explicit dependencies shared by WebSocket transport and HTTP routes."""
http: GatewayHTTPHandler http: GatewayHTTPHandler
endpoint: WebUIGatewayEndpoint
settings: WebUISettingsServices settings: WebUISettingsServices
tokens: GatewayTokenStore tokens: GatewayTokenStore
media: WebUIMediaGateway media: WebUIMediaGateway
@@ -42,7 +39,6 @@ class GatewayServices:
transcripts: WebUITranscriptRecorder transcripts: WebUITranscriptRecorder
workspaces: WebUIWorkspaceController workspaces: WebUIWorkspaceController
temporary_chats: WebUITemporaryChats temporary_chats: WebUITemporaryChats
session_projection: WebUISessionProjection
session_manager: SessionManager | None session_manager: SessionManager | None
cron_service: CronService | None cron_service: CronService | None
local_trigger_store: LocalTriggerStore | None local_trigger_store: LocalTriggerStore | None
@@ -112,7 +108,6 @@ def build_gateway_services(
workspaces=workspaces, workspaces=workspaces,
logger=logger, logger=logger,
) )
session_projection = WebUISessionProjection(session_manager, log=logger)
http = GatewayHTTPHandler( http = GatewayHTTPHandler(
config=config, config=config,
session_manager=session_manager, session_manager=session_manager,
@@ -140,10 +135,8 @@ def build_gateway_services(
recovery_action=recovery_action, recovery_action=recovery_action,
log=logger, log=logger,
) )
endpoint = WebUIGatewayEndpoint(config=config, http=http, tokens=tokens)
return GatewayServices( return GatewayServices(
http=http, http=http,
endpoint=endpoint,
settings=settings, settings=settings,
tokens=tokens, tokens=tokens,
media=media, media=media,
@@ -151,7 +144,6 @@ def build_gateway_services(
transcripts=transcripts, transcripts=transcripts,
workspaces=workspaces, workspaces=workspaces,
temporary_chats=temporary_chats, temporary_chats=temporary_chats,
session_projection=session_projection,
session_manager=session_manager, session_manager=session_manager,
cron_service=cron_service, cron_service=cron_service,
local_trigger_store=local_trigger_store, local_trigger_store=local_trigger_store,
-978
View File
@@ -1,978 +0,0 @@
"""Application orchestration for typed WebUI WebSocket commands."""
from __future__ import annotations
import asyncio
import hashlib
import json
import re
import time
import uuid
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Protocol, cast
from loguru import logger
from websockets.asyncio.server import ServerConnection
from nanobot.bus.events import INBOUND_META_USER_SHELL
from nanobot.command.builtin import USER_SHELL_COMMAND, builtin_command_starts_agent_turn
from nanobot.runtime_context import (
RUNTIME_CONTEXT_INPUT_META,
WEBUI_QUOTE_METADATA,
RuntimeContextBlock,
webui_quote_runtime_context,
)
from nanobot.security.workspace_access import (
WORKSPACE_SCOPE_METADATA_KEY,
WorkspaceScopeError,
)
from nanobot.session.webui_turns import (
clear_websocket_turn_if_current,
clear_websocket_turns,
register_queued_websocket_turn_if_idle,
websocket_turn_id,
websocket_turn_wall_started_at,
)
from nanobot.utils.helpers import safe_filename
from nanobot.webui.cli_apps_api import normalize_cli_app_mentions
from nanobot.webui.forking import handle_webui_fork_chat
from nanobot.webui.gateway_services import GatewayServices
from nanobot.webui.mcp_presets_api import normalize_mcp_preset_mentions
from nanobot.webui.metadata import WEBSOCKET_TURN_OWNER_METADATA_KEY
from nanobot.webui.session_access import (
SessionMention,
WebuiSessionAccess,
session_mentions_runtime_context,
)
from nanobot.webui.session_identity import is_valid_webui_chat_id, webui_session_key
from nanobot.webui.sidebar_state import write_webui_sidebar_state
from nanobot.webui.temporary_chats import TemporaryChatError
from nanobot.webui.transcription_ws import webui_transcription_event
_WEBUI_REQUEST_CACHE_TTL_S = 5 * 60.0
_WEBUI_REQUEST_CACHE_MAX = 256
@dataclass(frozen=True)
class WebUIRequestResult:
result: Any = None
status: int | None = None
message: str | None = None
@dataclass
class WebUIRequestOperation:
action: str
payload_digest: bytes
task: asyncio.Task[WebUIRequestResult]
completed_at: float | None = None
class WebUICommandTransport(Protocol):
"""Typed transport capabilities consumed by WebUI command orchestration."""
def is_allowed(self, sender_id: str) -> bool: ...
def webui_subscribers(self, chat_id: str) -> tuple[ServerConnection, ...]: ...
def webui_connection_chats(self, connection: ServerConnection) -> tuple[str, ...]: ...
def webui_attach(self, connection: ServerConnection, chat_id: str) -> None: ...
def webui_detach(self, connection: ServerConnection, chat_id: str) -> None: ...
def webui_clear_connection_default(self, connection: ServerConnection) -> None: ...
def webui_clear_stream_buffers(self, chat_id: str) -> None: ...
async def webui_hydrate(self, chat_id: str) -> None: ...
async def webui_send_event(
self,
connection: ServerConnection,
event: str,
**fields: Any,
) -> None: ...
async def webui_send_raw(
self,
connection: ServerConnection,
raw: str,
*,
label: str = "",
) -> None: ...
async def webui_dispatch_message(
self,
*,
sender_id: str,
chat_id: str,
content: str,
media: list[str] | None,
metadata: dict[str, Any],
is_dm: bool,
session_key: str | None,
require_existing_session: bool,
) -> None: ...
async def send_session_updated(
self,
chat_id: str,
*,
scope: str | None = None,
) -> None: ...
class WebUICommandRouter:
"""Own WebUI command semantics while a transport host owns raw connections."""
def __init__(self, transport: WebUICommandTransport, gateway: GatewayServices) -> None:
self._transport = transport
self.gateway = gateway
self._http_router = gateway.http
self._media = gateway.media
self._ingress = gateway.ingress
self._transcripts = gateway.transcripts
self._workspaces = gateway.workspaces
self._temporary_chats = gateway.temporary_chats
self._session_projection = gateway.session_projection
self._webui_connections = gateway.endpoint.webui_connections
self._session_access = (
WebuiSessionAccess(gateway.session_manager)
if gateway.session_manager is not None
else None
)
self.request_tasks: dict[
tuple[ServerConnection, str],
asyncio.Task[None],
] = {}
self.request_operations: dict[str, WebUIRequestOperation] = {}
self.request_locks: dict[ServerConnection, asyncio.Lock] = {}
def workspace_controls_available(self, connection: ServerConnection) -> bool:
return self._http_router.workspace_controls_available(connection)
async def send_webui_protocol_error(
self,
connection: ServerConnection,
detail: str,
) -> None:
await self._transport.webui_send_event(connection, "error", detail=detail)
async def attach_webui_fork(
self,
connection: ServerConnection,
*,
fork_id: str,
fork_key: str,
) -> None:
scope = self._workspaces.scope_for_session_key(fork_key)
self._transport.webui_attach(connection, fork_id)
await self._transport.webui_send_event(
connection,
"attached",
chat_id=fork_id,
**self._session_projection.attach_fields(fork_key),
)
await self._transport.webui_send_event(
connection,
"session_updated",
chat_id=fork_id,
scope="metadata",
workspace_scope=scope.payload(),
)
await self._transport.webui_hydrate(fork_id)
async def discard_owned_chat(
self,
connection: ServerConnection,
chat_id: str,
) -> None:
await self._temporary_chats.discard(connection, chat_id)
self._transport.webui_detach(connection, chat_id)
clear_websocket_turns(chat_id)
self._transport.webui_clear_stream_buffers(chat_id)
async def cleanup_connection(self, connection: ServerConnection) -> None:
"""Release command-owned state associated with one transport connection."""
chat_ids = self._transport.webui_connection_chats(connection)
for chat_id in chat_ids:
if self._temporary_chats.owns(connection, chat_id):
await self.discard_owned_chat(connection, chat_id)
else:
self._transport.webui_detach(connection, chat_id)
for chat_id in self._temporary_chats.chat_ids_for_owner(connection):
await self.discard_owned_chat(connection, chat_id)
self._transport.webui_clear_connection_default(connection)
self.gateway.endpoint.discard_connection(connection)
self.discard_request_lock_if_idle(connection)
async def broadcast_webui_event(self, event: str, **fields: Any) -> None:
for connection in tuple(self._webui_connections):
await self._transport.webui_send_event(connection, event, **fields)
async def broadcast_user_message(
self,
origin: ServerConnection,
chat_id: str,
text: str,
*,
turn_id: str | None,
starts_turn: bool,
media_paths: list[str],
media_names: list[str | None],
cli_apps: list[dict[str, Any]],
mcp_presets: list[dict[str, Any]],
session_mentions: list[SessionMention],
) -> None:
body: dict[str, Any] = {
"event": "user_message",
"chat_id": chat_id,
"text": text,
"starts_turn": starts_turn,
}
if turn_id is not None:
body["turn_id"] = turn_id
media = self._media.augment_transcript_user_media(media_paths)
for attachment, name in zip(media, media_names, strict=False):
if name:
attachment["name"] = name
if media:
body["media_urls"] = media
if cli_apps:
body["cli_apps"] = cli_apps
if mcp_presets:
body["mcp_presets"] = mcp_presets
if session_mentions:
body["session_mentions"] = session_mentions
active_turn_id = websocket_turn_id(chat_id)
if active_turn_id is not None:
body["active_turn_id"] = active_turn_id
started_at = websocket_turn_wall_started_at(chat_id)
if active_turn_id is not None and started_at is not None:
body["started_at"] = started_at
raw = json.dumps(body, ensure_ascii=False)
for connection in self._transport.webui_subscribers(chat_id):
if connection is not origin:
await self._transport.webui_send_raw(connection, raw, label=" user_message ")
async def workspace_scope_or_error(
self,
connection: ServerConnection,
resolver: Callable[[], Any],
*,
chat_id: str | None = None,
turn_id: str | None = None,
) -> Any | None:
try:
return resolver()
except WorkspaceScopeError as exc:
await self._transport.webui_send_event(
connection,
"error",
detail="workspace_scope_rejected",
reason=exc.message,
**({"chat_id": chat_id} if chat_id else {}),
**({"turn_id": turn_id} if turn_id else {}),
)
return None
async def dispatch(
self,
connection: ServerConnection,
client_id: str,
envelope: dict[str, Any],
) -> None:
"""Execute one typed WebUI command."""
command_type = envelope.get("type")
if command_type == "webui_request":
await self.start_webui_request(connection, envelope)
return
if command_type == "new_chat":
new_id = str(uuid.uuid4())
scope = await self.workspace_scope_or_error(
connection,
lambda: self._workspaces.scope_for_new_chat(
envelope,
controls_available=self.workspace_controls_available(connection),
),
)
if scope is None:
return
self._workspaces.stage_scope(new_id, scope)
self._transport.webui_attach(connection, new_id)
await self._transport.webui_send_event(
connection,
"attached",
chat_id=new_id,
**self._session_projection.attach_fields(webui_session_key(new_id)),
)
await self._transport.webui_send_event(
connection,
"session_updated",
chat_id=new_id,
scope="metadata",
workspace_scope=scope.payload(),
)
await self._transport.webui_hydrate(new_id)
return
if command_type == "new_temporary_chat":
try:
new_id = self._temporary_chats.create(
connection,
trusted_webui=connection in self._webui_connections,
)
except TemporaryChatError as exc:
await self._transport.webui_send_event(connection, "error", detail=exc.detail)
return
self._transport.webui_attach(connection, new_id)
await self._transport.webui_send_event(
connection,
"attached",
chat_id=new_id,
temporary=True,
)
return
if command_type == "fork_chat":
await handle_webui_fork_chat(self, connection, envelope)
return
if command_type == "discard_temporary_chat":
chat_id = envelope.get("chat_id")
if not is_valid_webui_chat_id(chat_id):
await self._transport.webui_send_event(
connection,
"error",
detail="invalid temporary chat_id",
)
return
try:
await self.discard_owned_chat(connection, chat_id)
except TemporaryChatError as exc:
await self._transport.webui_send_event(
connection,
"error",
detail=exc.detail,
chat_id=chat_id,
)
return
if command_type == "attach":
chat_id = envelope.get("chat_id")
if not is_valid_webui_chat_id(chat_id):
await self._transport.webui_send_event(
connection,
"error",
detail="invalid chat_id",
)
return
try:
self._temporary_chats.validate_attach(chat_id)
except TemporaryChatError as exc:
await self._transport.webui_send_event(
connection,
"error",
detail=exc.detail,
chat_id=chat_id,
)
return
self._transport.webui_attach(connection, chat_id)
await self._transport.webui_send_event(
connection,
"attached",
chat_id=chat_id,
**self._session_projection.attach_fields(webui_session_key(chat_id)),
)
await self._transport.webui_hydrate(chat_id)
return
if command_type == "set_sidebar_state":
if connection not in self._webui_connections:
await self._transport.webui_send_event(connection, "error", detail="access_denied")
return
state = envelope.get("state")
if not isinstance(state, dict):
await self._transport.webui_send_event(
connection,
"error",
detail="invalid_sidebar_state",
)
return
try:
saved_state = await asyncio.to_thread(
write_webui_sidebar_state,
cast(dict[str, Any], state),
)
except (OSError, ValueError):
await self._transport.webui_send_event(
connection,
"error",
detail="invalid_sidebar_state",
)
return
await self.broadcast_webui_event("sidebar_state_updated", state=saved_state)
return
if command_type == "set_workspace_scope":
chat_id = envelope.get("chat_id")
if not is_valid_webui_chat_id(chat_id):
await self._transport.webui_send_event(
connection,
"error",
detail="invalid chat_id",
)
return
try:
self._temporary_chats.validate_workspace_update(chat_id)
except TemporaryChatError as exc:
await self._transport.webui_send_event(
connection,
"error",
detail=exc.detail,
chat_id=chat_id,
)
return
scope = await self.workspace_scope_or_error(
connection,
lambda: self._workspaces.scope_for_set_request(
envelope,
chat_id=chat_id,
chat_running=websocket_turn_wall_started_at(chat_id) is not None,
controls_available=self.workspace_controls_available(connection),
),
chat_id=chat_id,
)
if scope is None:
return
self._workspaces.stage_scope(chat_id, scope)
await self._transport.send_session_updated(chat_id, scope="metadata")
await self._transport.webui_send_event(
connection,
"session_updated",
chat_id=chat_id,
scope="metadata",
workspace_scope=scope.payload(),
)
return
if command_type == "transcribe_audio":
event, payload = await webui_transcription_event(
envelope,
config_path=self.gateway.settings.config.path,
)
await self._transport.webui_send_event(connection, event, **payload)
return
if command_type == "message":
await self._dispatch_message(connection, client_id, envelope)
return
await self._transport.webui_send_event(
connection,
"error",
detail=f"unknown type: {command_type!r}",
)
async def _dispatch_message(
self,
connection: ServerConnection,
client_id: str,
envelope: dict[str, Any],
) -> None:
chat_id = envelope.get("chat_id")
content = envelope.get("content")
if not is_valid_webui_chat_id(chat_id):
await self._transport.webui_send_event(connection, "error", detail="invalid chat_id")
return
raw_turn_id = envelope.get("turn_id")
turn_id = raw_turn_id if isinstance(raw_turn_id, str) and raw_turn_id else None
rejection_fields = {
"chat_id": chat_id,
**({"turn_id": turn_id} if turn_id else {}),
}
if not self._transport.is_allowed(client_id):
await self._transport.webui_send_event(
connection,
"error",
detail="access_denied",
**rejection_fields,
)
return
if not isinstance(content, str):
await self._transport.webui_send_event(
connection,
"error",
detail="missing content",
**rejection_fields,
)
return
message_rejection = self._ingress.validate_text(content)
if message_rejection is not None:
await self._transport.webui_send_event(
connection,
"error",
detail="message_rejected",
reason=message_rejection,
**rejection_fields,
)
return
try:
temporary_policy = self._temporary_chats.message_policy(
connection,
chat_id,
content,
)
except TemporaryChatError as exc:
await self._transport.webui_send_event(
connection,
"error",
detail=exc.detail,
**rejection_fields,
)
return
raw_media = envelope.get("media")
media_paths: list[str] = []
media_names: list[str | None] = []
if raw_media is not None:
if not isinstance(raw_media, list):
await self._transport.webui_send_event(
connection,
"error",
detail="attachment_rejected",
reason="malformed",
**rejection_fields,
)
return
media_paths, reason = self._media.store_inbound_attachments(
cast(list[Any], raw_media)
)
if reason is not None:
await self._transport.webui_send_event(
connection,
"error",
detail="attachment_rejected",
reason=reason,
**rejection_fields,
)
return
for item in cast(list[Any], raw_media):
attachment = cast(dict[str, Any], item) if isinstance(item, dict) else {}
name = attachment.get("name")
media_names.append((safe_filename(name) or None) if isinstance(name, str) else None)
if temporary_policy is not None:
self._temporary_chats.register_media(connection, chat_id, media_paths)
if not content.strip() and not media_paths:
await self._transport.webui_send_event(
connection,
"error",
detail="missing content",
**rejection_fields,
)
return
self._transport.webui_attach(connection, chat_id)
if temporary_policy is None or temporary_policy.hydrate_transcript:
await self._transport.webui_hydrate(chat_id)
scope = await self.workspace_scope_or_error(
connection,
lambda: (
temporary_policy.workspace_scope
if temporary_policy is not None
else self._workspaces.scope_for_message(
envelope,
chat_id=chat_id,
chat_running=websocket_turn_wall_started_at(chat_id) is not None,
controls_available=self.workspace_controls_available(connection),
)
),
chat_id=chat_id,
turn_id=turn_id,
)
if scope is None:
return
if not self._transport.is_allowed(client_id):
await self._transport.webui_send_event(
connection,
"error",
detail="access_denied",
**rejection_fields,
)
return
metadata: dict[str, Any] = {
"remote": getattr(connection, "remote_address", None)
}
if envelope.get("webui") is True:
metadata["webui"] = True
metadata.update(self._transcripts.client_turn_metadata(envelope.get("turn_id")))
trusted_webui = metadata.get("webui") is True and connection in self._webui_connections
is_user_shell = (
trusted_webui
and envelope.get("user_shell") is True
and content.startswith("!")
)
if is_user_shell:
metadata[INBOUND_META_USER_SHELL] = True
dispatch_content = (
f"{USER_SHELL_COMMAND} {content[1:].lstrip()}" if is_user_shell else content
)
cli_apps = normalize_cli_app_mentions(envelope.get("cli_apps"))
if cli_apps:
metadata["cli_apps"] = cli_apps
mcp_presets = normalize_mcp_preset_mentions(
envelope.get("mcp_presets"),
config_path=self.gateway.settings.config.path,
)
if mcp_presets:
metadata["mcp_presets"] = mcp_presets
session_mentions: list[SessionMention] = []
if trusted_webui and self._session_access is not None:
session_mentions = await asyncio.to_thread(
self._session_access.normalize_mentions,
envelope.get("session_mentions"),
exclude_session_key=webui_session_key(chat_id),
)
if session_mentions:
metadata["session_mentions"] = session_mentions
metadata[WORKSPACE_SCOPE_METADATA_KEY] = scope.metadata()
is_webui = metadata.get("webui") is True
queued_owner = None
if is_webui and not is_user_shell and builtin_command_starts_agent_turn(content):
queued_owner = register_queued_websocket_turn_if_idle(chat_id, turn_id)
if queued_owner is not None:
metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] = queued_owner
accepted = False
try:
if is_webui and (
temporary_policy is None or temporary_policy.persist_transcript
):
self._transcripts.append_user_message(
chat_id,
content,
metadata=metadata,
media_paths=media_paths or None,
cli_apps=cli_apps or None,
mcp_presets=mcp_presets or None,
session_mentions=session_mentions or None,
)
if trusted_webui:
context_blocks: list[RuntimeContextBlock] = []
quote = webui_quote_runtime_context(
{WEBUI_QUOTE_METADATA: envelope.get("quoted_context")}
)
if quote is not None:
context_blocks.append(quote)
session_context = session_mentions_runtime_context(session_mentions)
if session_context is not None:
context_blocks.append(session_context)
if context_blocks:
metadata[RUNTIME_CONTEXT_INPUT_META] = context_blocks
await self._transport.webui_dispatch_message(
sender_id=client_id,
chat_id=chat_id,
content=dispatch_content,
media=media_paths or None,
metadata=metadata,
is_dm=False,
session_key=(
temporary_policy.session_key if temporary_policy is not None else None
),
require_existing_session=(
temporary_policy.require_existing_session
if temporary_policy is not None
else False
),
)
self._workspaces.persist_scope(chat_id, scope)
accepted = True
finally:
if not accepted and queued_owner is not None:
clear_websocket_turn_if_current(chat_id, queued_owner)
if is_webui:
await self.broadcast_user_message(
connection,
chat_id,
content,
turn_id=turn_id,
starts_turn=queued_owner is not None,
media_paths=media_paths,
media_names=media_names,
cli_apps=cli_apps,
mcp_presets=mcp_presets,
session_mentions=session_mentions,
)
if is_webui and turn_id:
active_turn_id = websocket_turn_id(chat_id)
started_at = websocket_turn_wall_started_at(chat_id)
await self._transport.webui_send_event(
connection,
"message_accepted",
chat_id=chat_id,
turn_id=turn_id,
starts_turn=queued_owner is not None,
**(
{"active_turn_id": active_turn_id}
if active_turn_id is not None
else {}
),
**(
{"started_at": started_at}
if active_turn_id is not None and started_at is not None
else {}
),
)
async def start_webui_request(
self,
connection: ServerConnection,
envelope: dict[str, Any],
) -> None:
request_id = envelope.get("request_id")
if not isinstance(request_id, str) or re.fullmatch(
r"[A-Za-z0-9._:-]{1,128}",
request_id,
) is None:
await self._transport.webui_send_event(
connection,
"error",
detail="invalid webui request_id",
)
return
if connection not in self._webui_connections:
await self.send_webui_response(
connection,
request_id,
status=403,
message="access_denied",
)
return
action = envelope.get("action")
payload = envelope.get("payload")
if not isinstance(action, str) or re.fullmatch(
r"[a-z][a-z0-9_.]{0,127}",
action,
) is None:
await self.send_webui_response(
connection,
request_id,
status=400,
message="invalid WebUI mutation action",
)
return
if not isinstance(payload, dict):
await self.send_webui_response(
connection,
request_id,
status=400,
message="WebUI mutation payload must be an object",
)
return
payload_digest = hashlib.sha256(
json.dumps(
payload,
ensure_ascii=False,
sort_keys=True,
separators=(",", ":"),
).encode("utf-8")
).digest()
self.prune_request_operations()
operation = self.request_operations.get(request_id)
is_replay = operation is not None
if operation is not None and (
operation.action != action or operation.payload_digest != payload_digest
):
await self.send_webui_response(
connection,
request_id,
status=409,
message="request_id was already used for a different WebUI mutation",
)
return
if operation is None:
operation_task = asyncio.create_task(
self.execute_webui_request(
connection,
action,
cast(dict[str, Any], payload),
)
)
new_operation = WebUIRequestOperation(
action=action,
payload_digest=payload_digest,
task=operation_task,
)
operation = new_operation
self.request_operations[request_id] = new_operation
def mark_complete(_task: asyncio.Task[WebUIRequestResult]) -> None:
current = self.request_operations.get(request_id)
if current is not new_operation:
return
new_operation.completed_at = time.monotonic()
self.prune_request_operations()
operation_task.add_done_callback(mark_complete)
key = (connection, request_id)
if key in self.request_tasks:
return
delivery_task = asyncio.create_task(
self.deliver_webui_request(
connection,
request_id,
operation.task,
sequence=is_replay,
)
)
self.request_tasks[key] = delivery_task
def prune_request_operations(self) -> None:
now = time.monotonic()
for request_id, operation in tuple(self.request_operations.items()):
if (
operation.completed_at is not None
and now - operation.completed_at >= _WEBUI_REQUEST_CACHE_TTL_S
):
self.request_operations.pop(request_id, None)
completed = sorted(
(
(operation.completed_at, request_id)
for request_id, operation in self.request_operations.items()
if operation.completed_at is not None
),
key=lambda item: item[0],
)
for _, request_id in completed[:-_WEBUI_REQUEST_CACHE_MAX]:
self.request_operations.pop(request_id, None)
def discard_request_lock_if_idle(self, connection: ServerConnection) -> None:
if connection in self._webui_connections:
return
if any(task_connection is connection for task_connection, _ in self.request_tasks):
return
self.request_locks.pop(connection, None)
async def deliver_webui_request(
self,
connection: ServerConnection,
request_id: str,
operation_task: asyncio.Task[WebUIRequestResult],
*,
sequence: bool = False,
) -> None:
try:
if sequence:
lock = self.request_locks.setdefault(connection, asyncio.Lock())
async with lock:
result = await asyncio.shield(operation_task)
await self.send_webui_response(
connection,
request_id,
result=result.result,
status=result.status,
message=result.message,
)
return
result = await asyncio.shield(operation_task)
await self.send_webui_response(
connection,
request_id,
result=result.result,
status=result.status,
message=result.message,
)
finally:
self.request_tasks.pop((connection, request_id), None)
self.discard_request_lock_if_idle(connection)
async def execute_webui_request(
self,
connection: ServerConnection,
action: str,
payload: dict[str, Any],
) -> WebUIRequestResult:
try:
lock = self.request_locks.setdefault(connection, asyncio.Lock())
async with lock:
response = await self._http_router.dispatch_webui_mutation(
connection,
action,
payload,
)
status = response.status_code
body = bytes(response.body).decode("utf-8", errors="replace").strip()
if 200 <= status < 300:
try:
result = json.loads(body)
except json.JSONDecodeError:
return WebUIRequestResult(
status=502,
message="WebUI mutation returned an invalid response",
)
if action == "sidebar.update" and isinstance(result, dict):
await self.broadcast_webui_event(
"sidebar_state_updated",
state=result,
)
return WebUIRequestResult(result=result)
return WebUIRequestResult(
status=status,
message=body or response.reason_phrase,
)
except asyncio.CancelledError:
raise
except Exception:
logger.exception("WebUI mutation '{}' failed", action)
return WebUIRequestResult(
status=500,
message="WebUI mutation failed",
)
async def send_webui_response(
self,
connection: ServerConnection,
request_id: str,
*,
result: Any = None,
status: int | None = None,
message: str | None = None,
) -> None:
if status is None:
await self._transport.webui_send_event(
connection,
"webui_response",
request_id=request_id,
ok=True,
result=result,
)
return
await self._transport.webui_send_event(
connection,
"webui_response",
request_id=request_id,
ok=False,
error={
"status": status,
"message": message or "WebUI mutation failed",
},
)
async def close(self) -> None:
"""Cancel command work and release application-owned gateway state."""
delivery_tasks = tuple(self.request_tasks.values())
operation_tasks = tuple(operation.task for operation in self.request_operations.values())
for task in (*delivery_tasks, *operation_tasks):
task.cancel()
if delivery_tasks:
await asyncio.gather(*delivery_tasks, return_exceptions=True)
if operation_tasks:
await asyncio.gather(*operation_tasks, return_exceptions=True)
self.request_tasks.clear()
self.request_locks.clear()
self.request_operations.clear()
self.gateway.tokens.clear()
self.gateway.endpoint.clear()
self._temporary_chats.close()
+2 -2
View File
@@ -16,7 +16,7 @@ from nanobot.agent.tools.mcp import MCPConnection, connect_mcp_servers
from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH, MCPOAuthHandlers from nanobot.agent.tools.mcp_oauth import MCP_OAUTH_CALLBACK_PATH, MCPOAuthHandlers
from nanobot.agent.tools.registry import ToolRegistry from nanobot.agent.tools.registry import ToolRegistry
from nanobot.config.schema import MCPServerConfig from nanobot.config.schema import MCPServerConfig
from nanobot.security.network import validate_url_target from nanobot.security.network import async_validate_url_target
from nanobot.webui.http_utils import is_loopback_host from nanobot.webui.http_utils import is_loopback_host
McpReload = Callable[[], Awaitable[dict[str, Any]]] McpReload = Callable[[], Awaitable[dict[str, Any]]]
@@ -259,7 +259,7 @@ class McpOAuthManager:
): ):
flow.error = "The MCP server returned an unsafe authorization URL." flow.error = "The MCP server returned an unsafe authorization URL."
raise McpOAuthError(flow.error) raise McpOAuthError(flow.error)
ok, _error = validate_url_target(authorization_url) ok, _error = await async_validate_url_target(authorization_url)
if not ok: if not ok:
flow.error = "The MCP server returned an unsafe authorization URL." flow.error = "The MCP server returned an unsafe authorization URL."
raise McpOAuthError(flow.error) raise McpOAuthError(flow.error)
-245
View File
@@ -1,245 +0,0 @@
"""Project agent runtime events onto the WebUI wire protocol."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Protocol
from loguru import logger
from nanobot.bus.events import OutboundMessage
from nanobot.bus.outbound_events import (
GoalStateSyncEvent,
GoalStatusEvent,
ProgressEvent,
RecoveryStateEvent,
RuntimeModelUpdatedEvent,
SessionUpdatedEvent,
TurnEndEvent,
TurnModelUpdatedEvent,
UserInputEvent,
outbound_event_from_message,
)
from nanobot.session.webui_turns import clear_websocket_turn_if_current
from nanobot.webui.metadata import (
WEBSOCKET_TURN_OWNER_METADATA_KEY,
WEBUI_SYSTEM_COMMAND_TURN_PREFIX,
WEBUI_TURN_METADATA_KEY,
)
from nanobot.webui.session_identity import webui_session_key
from nanobot.webui.session_projection import WebUISessionProjection
if TYPE_CHECKING:
from websockets.asyncio.server import ServerConnection
from nanobot.providers.base import LLMUsage
class WebUIOutboundTransport(Protocol):
"""Wire operations required by the outbound application projector."""
def webui_subscribers(self, chat_id: str) -> tuple[ServerConnection, ...]: ...
async def send_runtime_model_updated(
self,
*,
model_name: str | None,
model_preset: str | None = None,
) -> None: ...
async def send_turn_model_updated(
self,
chat_id: str,
*,
model_name: str,
model_preset: str | None = None,
context_window_tokens: int | None = None,
fallback: bool = False,
) -> None: ...
async def send_user_input(
self,
chat_id: str,
*,
content: str,
created_at_ms: int,
provenance: dict[str, Any],
) -> None: ...
async def send_recovery_state(self, chat_id: str, event: RecoveryStateEvent) -> None: ...
async def send_goal_state(self, chat_id: str, blob: dict[str, Any]) -> None: ...
async def send_goal_status(
self,
chat_id: str,
status: str,
*,
started_at: float | None = None,
turn_id: str | None = None,
) -> None: ...
async def send_turn_end(
self,
chat_id: str,
latency_ms: int | None = None,
*,
goal_state: dict[str, Any] | None = None,
usage: LLMUsage | None = None,
context_window_tokens: int | None = None,
metadata: dict[str, Any] | None = None,
turn_owner: str | None = None,
) -> None: ...
async def send_session_updated(self, chat_id: str, *, scope: str | None = None) -> None: ...
async def send_file_edit_events(
self,
chat_id: str,
edits: list[dict[str, Any]],
metadata: dict[str, Any] | None = None,
) -> None: ...
async def send_projected_message(
self,
msg: OutboundMessage,
progress_event: ProgressEvent | None,
) -> None: ...
class WebUIOutboundProjector:
"""Interpret runtime events without coupling that state machine to the channel."""
def __init__(
self,
transport: WebUIOutboundTransport,
session_projection: WebUISessionProjection,
) -> None:
self._transport = transport
self._session_projection = session_projection
async def hydrate(self, chat_id: str) -> None:
"""Replay reconnect state through the existing stable wire operations."""
for event in self._session_projection.hydration_events(
webui_session_key(chat_id),
chat_id,
):
if event["event"] == "goal_state":
await self._transport.send_goal_state(chat_id, event["goal_state"])
continue
await self._transport.send_goal_status(
chat_id,
"running",
started_at=event["started_at"],
turn_id=event.get("turn_id"),
)
async def send(self, msg: OutboundMessage) -> None:
event = outbound_event_from_message(msg)
progress_event = event if isinstance(event, ProgressEvent) else None
if isinstance(event, RuntimeModelUpdatedEvent):
await self._transport.send_runtime_model_updated(
model_name=event.model,
model_preset=event.model_preset,
)
return
conns = list(self._transport.webui_subscribers(msg.chat_id))
if not conns:
quiet_events = (
ProgressEvent,
UserInputEvent,
TurnEndEvent,
SessionUpdatedEvent,
GoalStatusEvent,
GoalStateSyncEvent,
)
log = (
logger.debug
if isinstance(event, quiet_events)
else logger.warning
)
log("no active subscribers for chat_id={}", msg.chat_id)
if isinstance(event, TurnModelUpdatedEvent):
if conns:
await self._transport.send_turn_model_updated(
msg.chat_id,
model_name=event.model,
model_preset=event.model_preset,
context_window_tokens=event.context_window_tokens,
fallback=event.fallback,
)
return
if isinstance(event, UserInputEvent):
if conns:
await self._transport.send_user_input(
msg.chat_id,
content=event.content,
created_at_ms=event.created_at_ms,
provenance=event.provenance,
)
return
if isinstance(event, RecoveryStateEvent):
if conns:
await self._transport.send_recovery_state(msg.chat_id, event)
return
if isinstance(event, GoalStateSyncEvent):
if conns:
await self._transport.send_goal_state(
msg.chat_id,
event.goal_state or {"active": False},
)
return
if isinstance(event, GoalStatusEvent):
turn_id = (msg.metadata or {}).get(WEBUI_TURN_METADATA_KEY)
current_turn_id = turn_id if isinstance(turn_id, str) else None
turn_owner = (msg.metadata or {}).get(WEBSOCKET_TURN_OWNER_METADATA_KEY)
current_turn_owner = turn_owner if isinstance(turn_owner, str) else None
try:
if conns and event.status in ("running", "idle"):
await self._transport.send_goal_status(
msg.chat_id,
event.status,
started_at=event.started_at,
turn_id=current_turn_id,
)
finally:
if event.status == "idle":
clear_websocket_turn_if_current(
msg.chat_id,
current_turn_owner,
preserve_persistence_failure=True,
)
return
if isinstance(event, TurnEndEvent):
turn_id = (msg.metadata or {}).get(WEBUI_TURN_METADATA_KEY)
session_update_scope = (
"metadata"
if isinstance(turn_id, str)
and turn_id.startswith(WEBUI_SYSTEM_COMMAND_TURN_PREFIX)
else "thread"
)
turn_owner = (msg.metadata or {}).get(WEBSOCKET_TURN_OWNER_METADATA_KEY)
await self._transport.send_turn_end(
msg.chat_id,
latency_ms=event.latency_ms,
goal_state=event.goal_state,
usage=event.usage,
context_window_tokens=event.context_window_tokens,
metadata=msg.metadata,
turn_owner=turn_owner if isinstance(turn_owner, str) else None,
)
await self._transport.send_session_updated(msg.chat_id, scope=session_update_scope)
return
if isinstance(event, SessionUpdatedEvent):
if conns:
await self._transport.send_session_updated(msg.chat_id, scope=event.scope)
return
if progress_event and progress_event.file_edit_events:
await self._transport.send_file_edit_events(
msg.chat_id,
progress_event.file_edit_events,
msg.metadata,
)
return
await self._transport.send_projected_message(msg, progress_event)
+1 -1
View File
@@ -44,7 +44,7 @@ def session_context_payload(session: Session) -> dict[str, Any]:
"schema_version": 1, "schema_version": 1,
"session_key": session.key, "session_key": session.key,
"total_messages": len(session.messages), "total_messages": len(session.messages),
"archived_messages": min(session.last_archived, len(session.messages)), "archived_messages": min(session.last_consolidated, len(session.messages)),
"replay_messages": len(replay), "replay_messages": len(replay),
"estimated_replay_tokens": replay_tokens, "estimated_replay_tokens": replay_tokens,
"estimated_summary_tokens": summary_tokens, "estimated_summary_tokens": summary_tokens,
-32
View File
@@ -1,32 +0,0 @@
"""Stable mapping between public WebUI chat IDs and persisted session keys."""
from __future__ import annotations
import re
from typing import Any, TypeGuard
WEBUI_SESSION_STORAGE_PREFIX = "websocket:"
_WEBUI_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$")
def is_valid_webui_chat_id(value: Any) -> TypeGuard[str]:
"""Validate the compact chat IDs accepted by the WebUI protocol."""
return isinstance(value, str) and _WEBUI_CHAT_ID_RE.fullmatch(value) is not None
def webui_session_key(chat_id: str) -> str:
"""Return the backward-compatible persisted key for a WebUI chat."""
return f"{WEBUI_SESSION_STORAGE_PREFIX}{chat_id}"
def is_webui_session_key(session_key: str) -> bool:
"""Return whether *session_key* belongs to the WebUI session namespace."""
return session_key.startswith(WEBUI_SESSION_STORAGE_PREFIX)
def webui_chat_id(session_key: str) -> str | None:
"""Extract a non-empty WebUI chat ID from a persisted session key."""
if not is_webui_session_key(session_key):
return None
chat_id = session_key.removeprefix(WEBUI_SESSION_STORAGE_PREFIX)
return chat_id or None
+7 -15
View File
@@ -32,12 +32,6 @@ from nanobot.session.manager import (
) )
from nanobot.session.model_selection import model_preset_from_metadata from nanobot.session.model_selection import model_preset_from_metadata
from nanobot.session.recovery import recovery_state_from_metadata from nanobot.session.recovery import recovery_state_from_metadata
from nanobot.webui.session_identity import (
WEBUI_SESSION_STORAGE_PREFIX,
is_webui_session_key,
webui_chat_id,
webui_session_key,
)
_INDEX_VERSION = 8 _INDEX_VERSION = 8
_INDEX_FILENAME = ".webui_session_index.json" _INDEX_FILENAME = ".webui_session_index.json"
@@ -56,7 +50,7 @@ _WEBUI_ACTIVITY_MTIME_NS = "webui_activity_mtime_ns"
_WEBUI_ACTIVITY_SIZE = "webui_activity_size" _WEBUI_ACTIVITY_SIZE = "webui_activity_size"
_WEBUI_ACTIVITY_FILES = "webui_activity_files" _WEBUI_ACTIVITY_FILES = "webui_activity_files"
_VISIBLE_TRANSCRIPT_ROLES = {"user", "assistant"} _VISIBLE_TRANSCRIPT_ROLES = {"user", "assistant"}
_WEBUI_SESSION_STEM_PREFIX = SessionManager.safe_key(WEBUI_SESSION_STORAGE_PREFIX) _WEBUI_SESSION_STEM_PREFIX = SessionManager.safe_key("websocket:")
_WEBUI_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$") _WEBUI_CHAT_ID_RE = re.compile(r"^[A-Za-z0-9_:-]{1,64}$")
_TRANSCRIPT_SEGMENTS_SUFFIX = ".segments" _TRANSCRIPT_SEGMENTS_SUFFIX = ".segments"
_TRANSCRIPT_NON_ANSWER_KINDS = {"progress", "reasoning", "tool_hint"} _TRANSCRIPT_NON_ANSWER_KINDS = {"progress", "reasoning", "tool_hint"}
@@ -96,7 +90,7 @@ def _reconcile_index(session_manager: SessionManager) -> tuple[list[dict[str, An
session_keys_by_stem = { session_keys_by_stem = {
SessionManager.safe_key(key): key SessionManager.safe_key(key): key
for key in session_paths for key in session_paths
if is_webui_session_key(key) if key.startswith("websocket:")
} }
rows: list[dict[str, Any]] = [] rows: list[dict[str, Any]] = []
changed = existing_rows is None changed = existing_rows is None
@@ -381,9 +375,9 @@ def _transcript_record(line: str) -> dict[str, Any] | None:
def _valid_transcript_session_key(key: str, stem: str) -> bool: def _valid_transcript_session_key(key: str, stem: str) -> bool:
chat_id = webui_chat_id(key) if not key.startswith("websocket:"):
if chat_id is None:
return False return False
chat_id = key.split(":", 1)[1]
return _WEBUI_CHAT_ID_RE.fullmatch(chat_id) is not None and SessionManager.safe_key(key) == stem return _WEBUI_CHAT_ID_RE.fullmatch(chat_id) is not None and SessionManager.safe_key(key) == stem
@@ -541,9 +535,7 @@ def _scan_transcript_row(
paths: tuple[Path, ...], paths: tuple[Path, ...],
webui_dir: Path, webui_dir: Path,
) -> dict[str, Any] | None: ) -> dict[str, Any] | None:
path_key = session_key or webui_session_key( path_key = session_key or f"websocket:{stem.removeprefix(_WEBUI_SESSION_STEM_PREFIX)}"
stem.removeprefix(_WEBUI_SESSION_STEM_PREFIX)
)
signature = _webui_activity_signature(path_key, webui_dir) signature = _webui_activity_signature(path_key, webui_dir)
activity_updated_at = _webui_activity_updated_at(signature) activity_updated_at = _webui_activity_updated_at(signature)
if activity_updated_at is None: if activity_updated_at is None:
@@ -568,7 +560,7 @@ def _scan_transcript_row(
saw_record = True saw_record = True
chat_id = record.get("chat_id") chat_id = record.get("chat_id")
if isinstance(chat_id, str) and chat_id.strip(): if isinstance(chat_id, str) and chat_id.strip():
candidate = webui_session_key(chat_id.strip()) candidate = f"websocket:{chat_id.strip()}"
if _valid_transcript_session_key(candidate, stem): if _valid_transcript_session_key(candidate, stem):
session_key = candidate session_key = candidate
if created_at is None: if created_at is None:
@@ -594,7 +586,7 @@ def _scan_transcript_row(
if not saw_record: if not saw_record:
return None return None
if session_key is None: if session_key is None:
fallback = webui_session_key(stem.removeprefix(_WEBUI_SESSION_STEM_PREFIX)) fallback = f"websocket:{stem.removeprefix(_WEBUI_SESSION_STEM_PREFIX)}"
if not _valid_transcript_session_key(fallback, stem): if not _valid_transcript_session_key(fallback, stem):
return None return None
session_key = fallback session_key = fallback

Some files were not shown because too many files have changed in this diff Show More