Compare commits

...
Author SHA1 Message Date
chengyongruandchengyongru 3c25f826ea fix(telegram): detect unsupported rich stream edits 2026-08-31 18:39:35 +08:00
Nolanandchengyongru 195e4c281d fix(telegram): preserve final-edit retry contract in rich stream upgrade
Address review feedback on the rich edit error classification:

- Transport, rate-limit, and unexpected errors now propagate instead of
  returning False. Returning False made send_delta fall through to an
  immediate legacy edit_message_text, which under connection-pool
  exhaustion doubled demand and discarded the buffered retry state that
  ChannelManager relies on.
- 'Message is not modified' is treated as success: when the rich edit is
  applied server-side but its response times out, the retry inside
  _call_with_retry reports the edit as already applied, and the previous
  BadRequest branch would have let the legacy edit overwrite the
  successful rich result.
- False is now returned only for capability errors (pre-10.1 servers,
  which still trip the rich latch) and content-shaped rejections, where
  the legacy HTML path is the intended fallback.

Adds regression coverage for a direct NetworkError (propagates, buffer
kept for manager retry) and TimedOut followed by Message-is-not-modified
(treated as success, no legacy overwrite).
2026-08-31 18:39:35 +08:00
Nolanandchengyongru 37663ac947 fix(telegram): upgrade streaming preview to rich in place at stream end
The rich branch in send_delta(stream_end=True) was unreachable: it was
guarded by 'not buf.message_id' after an earlier return had already
ensured buf.message_id is set, so sendRichMessage never fired with
streaming enabled and the final message always went through the legacy
HTML editMessageText path.

Bot API 10.1 added a rich_message parameter to editMessageText, which
upgrades an existing message to rich in place. Use it at stream end via
do_api_request so the streaming preview keeps its identity — no
delete-and-resend, so none of the flickering or dropped line breaks
that made rich-at-stream-end fail in #4470.

Capability errors (server older than 10.1) trip the existing
_rich_send_disabled latch and fall back to the legacy HTML path, which
is unchanged.

Fixes #5516
2026-08-31 18:39:35 +08:00
chengyongruandGitHub e111b83af6 refactor(agent): unify runner request fitting (#5612)
* refactor(agent): unify runner request fitting

* fix(agent): fit every runner model request

* fix(agent): count resumed state during request fitting

* refactor(agent): consolidate request fitting state
2026-08-31 18:08:21 +08:00
chengyongruandchengyongru 6d6d58d329 style(tui): simplify runtime header controls 2026-08-31 17:03:24 +08:00
chengyongruandchengyongru e69159cdae style(tui): replace header dots with spacing 2026-08-31 17:03:24 +08:00
feb33e1f99 docs(tools): clarify edit_file selector exclusivity (#5598)
* docs(tools): clarify edit selector exclusivity

* docs(tools): streamline edit_file guidance

---------

Co-authored-by: chengyongru <chengyongru.ai@gmail.com>
2026-08-31 14:28:41 +08:00
chengyongruandGitHub bb34b58f47 refactor(agent): make memory summaries cumulative (#5610)
* refactor(agent): make memory summaries cumulative

Treat the latest session summary as a replacement checkpoint, preserve it through bounded raw fallbacks, and reserve history.jsonl for Dream ingestion.

* fix(agent): preserve cumulative checkpoint context

* fix(agent): preserve memory archive prompt cache

* refactor(agent): state archive prompt positively

* refactor(agent): remove checkpoint version migration

* refactor(agent): summarize full archive context

* test(agent): align cumulative archive prompt assertion

* refactor(agent): clarify memory checkpoint contract
2026-08-31 13:37:20 +08:00
chengyongruandGitHub 6cd7063682 refactor(agent): defer transcript assembly to runner (#5608)
* refactor(agent): defer transcript assembly to runner

Keep persisted history and the fresh turn as explicit inputs until the Runner assembles the provider transcript. Preserve ContextBuilder and direct AgentRunner compatibility while making the save boundary structural.

Refs NAN-81.

* fix(providers): preserve mixed adjacent user content
2026-08-31 00:06:15 +08:00
Xubin Ren d019658501 fix(dingtalk): stop late inbound task creation 2026-08-30 17:56:34 +08:00
yu-xin-candXubin Ren e8385d9257 fix(dingtalk): drain inbound background tasks 2026-08-30 17:56:34 +08:00
Xubin Ren 5c71ef6e49 test(email): assert header-first fetch contract 2026-08-30 17:38:44 +08:00
Till AdamandXubin Ren f573ecfe56 perf(email): fetch headers before body, use UID SEARCH to skip re-fetch
The IMAP poll loop previously downloaded the entire message body for
every UNSEEN message before running any filter (self-sent, SPF/DKIM,
allow-list), and only learned the UID by parsing it back out of that
fetch response. Rejected messages stay unseen and get re-fetched in
full on every subsequent poll.

Switch to UID SEARCH (UIDs come back directly, no per-message fetch
needed to learn them) so already-processed UIDs are skipped before any
network fetch, then fetch headers only to evaluate every filter — the
full body is downloaded only for messages that pass every check and
are actually delivered. No behavior change: filter outcomes and
\Seen semantics are unchanged.
2026-08-30 17:38:44 +08:00
Xubin Ren 1ac1b35c84 fix(cron): sanitize migrated legacy metadata 2026-08-30 17:24:05 +08:00
Oxygen56andXubin Ren 679a07460e fix(cron): sanitize persisted origin metadata 2026-08-30 17:24:05 +08:00
Oxygen56andXubin Ren 919e3d341e fix(agent): add retry hint to tool exceptions 2026-08-30 17:03:51 +08:00
yu-xin-candXubin Ren 2c55934198 fix(agent): bound session message rate-limit state 2026-08-30 16:46:36 +08:00
Xubin Ren 5afdffff51 fix(agent): settle reasoning close before cancellation 2026-08-30 16:32:08 +08:00
KDBandXubin Ren bfe041def7 fix(agent): close reasoning stream on cancellation 2026-08-30 16:32:08 +08:00
Kail Tianandchengyongru 1c1b13a3a9 fix(tui): preserve cursor position on Windows exit 2026-08-30 12:17:40 +08:00
Xubin Ren 2c87143f77 fix(cli): preserve gateway log stream boundaries 2026-08-29 23:03:23 +08:00
Xubin Ren d7df2726de fix(cli): stream gateway logs in WebUI launcher 2026-08-29 23:03:23 +08:00
Xubin Ren c02f013b17 fix(providers): harden OAuth model discovery 2026-08-29 21:22:20 +08:00
Xubin Ren 1c6483147e refactor(providers): localize OAuth model discovery 2026-08-29 21:22:20 +08:00
Xubin Ren 7941450a5d fix(providers): recover incomplete Grok searches 2026-08-29 21:22:20 +08:00
Xubin Ren e6c839ee37 refactor(webui): simplify preset save label 2026-08-29 21:22:20 +08:00
Xubin Ren 97fb9aaf72 fix(providers): complete OAuth model discovery 2026-08-29 21:22:20 +08:00
Xubin Ren bc4de246a4 feat(providers): discover OAuth model catalogs online 2026-08-29 21:22:20 +08:00
Xubin Ren 2389ab1f5a feat(providers): add Grok 4.6 subscription model 2026-08-29 21:22:20 +08:00
Xubin Ren 65f2a6dbf5 fix(webui): hide SkillHub install counts 2026-08-29 15:29:18 +08:00
Xubin Ren caab883f9f fix(webui): preserve named pane groups 2026-08-29 14:22:19 +08:00
Xubin Ren 1fe14f2ee6 fix(cli): preserve root shell completion 2026-08-29 11:26:50 +08:00
Xubin Ren 7fc90ca6aa feat(cli): make nanobot launch the terminal agent 2026-08-29 11:26:50 +08:00
chengyongruandchengyongru 559b2d2e5d test(tui): avoid clipboard status race 2026-08-28 17:14:39 +08:00
chengyongruandchengyongru a339966543 fix(tui): preserve full UI in Herdr panes 2026-08-28 16:36:07 +08:00
chengyongruandGitHub e73cce706c refactor(agent): extract tool execution boundary (#5569)
* refactor(agent): extract tool execution boundary

* test(agent): use extracted tool execution boundary
2026-08-28 13:52:03 +08:00
chengyongruandGitHub cace42af14 refactor(memory): remove consolidation ratio (#5575)
* refactor(memory): remove consolidation ratio

* docs(memory): document fixed consolidation policy

* docs(memory): simplify consolidation overview

* docs(memory): rely on soft wrapping
2026-08-28 13:20:09 +08:00
chengyongruandGitHub 29025f5a8b fix(agent): default request concurrency to unlimited (#5572)
* fix(agent): default request concurrency to unlimited

* test(agent): clarify session serialization coverage
2026-08-27 23:31:38 +08:00
chengyongruandGitHub 3c61fef7e8 refactor(memory): decouple archival from provider state (#5565)
* refactor(memory): decouple archival from provider state

* test(memory): remove obsolete consolidation offset coverage
2026-08-27 21:21:15 +08:00
chengyongruandGitHub 4d204ba077 feat(tui): support pasting clipboard images (#5563)
* feat(tui): support pasting clipboard images

* fix(tui): keep image placeholders atomic

* fix(tui): reconcile duplicate image placeholders

* fix(tui): retain highlighted image placeholders

* fix(tui): preserve image placeholder layout

* fix(tui): keep image display state local

* fix(tui): reject images in commands
2026-08-27 20:37:43 +08:00
chengyongruandGitHub b9e7c7f6fe fix: queue concurrent subagents (#5566)
* fix: queue concurrent subagents

* chore: keep spawn schema concise
2026-08-27 17:53:28 +08:00
chengyongruandGitHub 39de4594d7 refactor(agent): decouple loop from message tool state (#5559)
* refactor(agent): decouple loop from message tool state

* refactor(agent): scope message delivery tracking per run

* refactor(agent): clarify message delivery scope name
2026-08-27 17:23:16 +08:00
chengyongruandchengyongru d6c112ab74 refactor(agent): load MyTool through tool loader 2026-08-27 14:05:21 +08:00
chengyongruandchengyongru 91f5a85db0 fix(agent): complete native reasoning lifecycle 2026-08-27 13:20:04 +08:00
chengyongruandGitHub 3a62b0b744 fix(tui): surface chat connection failures (#5543)
* fix(tui): surface chat connection failures

* fix(tui): keep connection copy implementation-agnostic

* test(tui): avoid millisecond timing boundary

* fix(tui): use gateway health for connection status

* fix(tui): use product-level readiness copy
2026-08-27 11:37:30 +08:00
chengyongruandchengyongru 55f85b3c1f refactor(agent): remove duplicate progress streaming path 2026-08-27 10:56:07 +08:00
Xubin Ren 2113870e27 fix(cli): foreground the WebUI browser 2026-08-27 10:45:40 +08:00
Xubin Ren b632186b5f perf(webui): smooth paged history scrolling 2026-08-27 10:45:40 +08:00
Xubin Ren d7d03b25ef fix(webui): preserve semantic history anchors 2026-08-27 10:45:40 +08:00
Xubin Ren 6a3f53a917 fix(webui): stabilize fast history scrolling 2026-08-27 10:45:40 +08:00
Xubin Ren 25e20a1458 feat(webui): link model picker to settings 2026-08-27 10:45:40 +08:00
Xubin Ren 5678f83290 fix(webui): tighten conversation top spacing 2026-08-27 10:45:40 +08:00
Xubin Ren cb7b640d36 fix(webui): soften composer scroll boundary 2026-08-27 10:45:40 +08:00
Xubin Ren 701926eba1 fix(webui): preserve file edit activity order 2026-08-27 10:45:40 +08:00
Xubin Ren bbbfacbc64 fix(webui): compact single-pane chat header 2026-08-27 10:45:40 +08:00
150 changed files with 9150 additions and 4308 deletions
+4 -4
View File
@@ -146,7 +146,7 @@ Activate it with `source .venv/bin/activate` on macOS/Linux or
python -m pip install -e .
```
After that, the normal commands are identical to a stable install. `nanobot agent` runs the TUI
After that, the normal commands are identical to a stable install. `nanobot` runs the TUI
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
`python -m pip install -e .` when Python dependencies change. Contributors should also read
@@ -206,10 +206,10 @@ Use `nanobot gateway --background` for the same direct entry point without keepi
**Prefer to work entirely in the terminal?**
```bash
nanobot agent
nanobot
```
This opens the native terminal client with the launch directory as its workspace. It shares saved conversations and the local gateway with the WebUI.
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.
- 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).
@@ -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:
```bash
nanobot agent -m "Hello!"
nanobot -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.
+13 -13
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 |
| 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 |
| Send one test message | `nanobot agent -m "Hello!"` | First proof that install, config, provider, model, and workspace all work |
| Chat in the terminal | `nanobot agent` | Interactive local chat; exit with `exit`, `/exit`, `:q`, or `Ctrl+D` |
| Send one test message | `nanobot -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 |
| 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 |
| 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 |
|---|---|
| `nanobot agent -m "Hello!"` | Send one message and exit |
| `nanobot agent` | Start interactive terminal chat |
| `nanobot agent --session <id>` | Use a WebSocket session key; add `--classic` for another channel |
| `nanobot agent --workspace <path>` | Override workspace |
| `nanobot agent --config <path>` | Use a specific config file |
| `nanobot agent --classic` | Use the classic Python prompt instead of the native terminal UI |
| `nanobot agent --theme auto\|dark\|light` | Auto-detect the terminal appearance or force a TUI palette |
| `nanobot agent --no-markdown` | Use the classic prompt and print plain text instead of Markdown |
| `nanobot agent --logs` | Use the classic prompt and show runtime logs while chatting |
| `nanobot -m "Hello!"` | Send one message and exit |
| `nanobot` | Start interactive terminal chat |
| `nanobot --session <id>` | Use a WebSocket session key; add `--classic` for another channel |
| `nanobot --workspace <path>` | Override workspace |
| `nanobot --config <path>` | Use a specific config file |
| `nanobot --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 --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 |
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
@@ -139,7 +139,7 @@ Interactive mode exits with `exit`, `quit`, `/exit`, `/quit`, `:q`, or `Ctrl+D`.
| Command | Description |
|---|---|
| `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` | 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 --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 --no-open` | Prepare and start the WebUI without opening a browser |
@@ -344,7 +344,7 @@ remain accepted as no-op compatibility aliases.
| Command | Description |
|---|---|
| `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.5; 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.6; 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 logout openai-codex` | Remove OpenAI Codex OAuth state |
| `nanobot provider logout xai-grok --config <path>` | Remove the selected nanobot instance's xAI OAuth state |
+24 -16
View File
@@ -188,7 +188,7 @@ These variables are process-level switches. Set them in the same terminal, servi
| Variable | Default | Description |
|----------|---------|-------------|
| `NANOBOT_MAX_CONCURRENT_REQUESTS` | `3` | Maximum concurrently running inbound agent requests. Must be an integer; set `0` or a negative value for unlimited. |
| `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_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_OPENAI_COMPAT_TIMEOUT_S` | `120` | HTTP request timeout, in seconds, for OpenAI-compatible providers. Invalid or non-positive values are ignored. |
@@ -729,6 +729,11 @@ Then run:
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:
```json
@@ -764,11 +769,14 @@ nanobot provider login xai-grok --set-main
nanobot agent -m "Hello from Grok."
```
The default model is `xai-grok/grok-4.5` with a 500,000-token context window.
The provider reads xAI's model catalog and includes the server-hosted `x_search`
tool only when the selected model advertises `supportsBackendSearch`. 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.
The default model is `xai-grok/grok-4.6` with a 500,000-token context window.
The provider reads and caches xAI's online model catalog for both WebUI model
selection and runtime capabilities. Newly available models appear automatically;
when discovery fails, the last successful catalog or built-in fallback remains
available. The server-hosted `x_search` tool is included only when the selected
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
WebUI provider settings or with `providers.xaiGrok.extraBody.tools: []`.
@@ -805,6 +813,10 @@ 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.
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:
```bash
export NANOBOT_GITHUB_COPILOT_CLIENT_ID="your-enterprise-client-id"
@@ -2213,7 +2225,7 @@ The notification gate runs on a built-in system prompt. Advanced users can overr
## Subagent Concurrency
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:
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:
```json
{
@@ -2229,7 +2241,7 @@ The deprecated `agents.defaults.failOnToolError` field is silently ignored when
| Option | Default | Description |
|--------|---------|-------------|
| `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.maxConcurrentSubagents` | `4` | Maximum number of subagents that may run at the same time. Additional tasks wait for capacity. |
## Auto Compact
@@ -2256,16 +2268,12 @@ When a user is idle for longer than a configured threshold, nanobot **proactivel
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.
2. **Background compaction**: Idle sessions summarize the older live prefix via LLM and keep the most recent legal suffix (currently 8 messages).
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 summary is also mirrored into session metadata so it can still be recovered after a process restart.
2. **Background compaction**: Older context is summarized while the most recent messages remain available.
3. **Session preservation**: The complete session history remains stored for later inspection and reuse.
4. **Restart-safe resume**: The compacted context remains available after a process restart.
> [!NOTE]
> 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.
> Auto compact shortens the context sent to the model without deleting the session's structured message history.
## Timezone
+1 -3
View File
@@ -29,9 +29,7 @@ Memory moves through nanobot in two stages.
### Stage 1: Consolidator
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`.
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.
This file is:
+15 -3
View File
@@ -572,15 +572,23 @@ For OpenAI Codex:
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:
```bash
nanobot provider login xai-grok --set-main
```
This selects `xai-grok/grok-4.5`. The provider reads xAI's model catalog and
exposes the hosted `x_search` tool only when the selected model advertises
`supportsBackendSearch`; otherwise the model runs without hosted X Search.
This selects `xai-grok/grok-4.6`. The WebUI model selector reads xAI's online
model catalog, so newly available subscription models appear without a nanobot
release. Online metadata is cached and enriched with nanobot's curated labels;
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
without invoking a local nanobot tool. Credentials are stored under the
active instance's `auth/xai.json` (normally `~/.nanobot/auth/xai.json`), not in
@@ -599,6 +607,10 @@ For GitHub Copilot:
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.
## Provider Resolution
+3 -3
View File
@@ -103,13 +103,13 @@ 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:
```bash
nanobot agent -m "Hello!"
nanobot -m "Hello!"
```
Then start an interactive terminal chat with:
```bash
nanobot agent
nanobot
```
In interactive mode, `Enter` sends and `Shift+Enter` inserts a newline (`Ctrl+J` is the
@@ -173,7 +173,7 @@ nanobot webui
```
The source path follows current `main` and can be newer than the published package. The editable
install keeps Python pointed at the checkout; `nanobot agent` runs `tui/` with Bun, and
install keeps Python pointed at the checkout; `nanobot` runs `tui/` with Bun, and
`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
[`../CONTRIBUTING.md`](../CONTRIBUTING.md).
+3 -1
View File
@@ -23,7 +23,9 @@ one is missing, starts or joins the same on-demand gateway used by the native
TUI, and opens the browser. With a fresh config,
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
it is not available from other devices on your LAN.
it is not available from other devices on your LAN. While the launcher remains
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:
+1 -1
View File
@@ -48,7 +48,7 @@ class AutoCompact:
def _has_unarchived_messages(self, key: str) -> bool:
session = self.sessions.get_or_create(key)
return session.last_consolidated < len(session.messages)
return session.last_archived < len(session.messages)
@classmethod
def _is_internal_session(cls, key: str) -> bool:
+67 -79
View File
@@ -30,11 +30,7 @@ from nanobot.security.workspace_access import WorkspaceScopeResolver
from nanobot.session.keys import last_channel_from_metadata
from nanobot.session.manager import Session
from nanobot.session.summary import SessionSummary
from nanobot.utils.helpers import (
detect_image_mime,
load_bundled_template,
truncate_text_to_tokens,
)
from nanobot.utils.helpers import detect_image_mime, load_bundled_template
from nanobot.utils.prompt_templates import render_template
@@ -75,14 +71,29 @@ class PersistedPromptContextResolver:
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:
"""Builds the context (system prompt + messages) for the agent."""
BOOTSTRAP_FILES = ["AGENTS.md", "SOUL.md", "USER.md"]
_SKIPPABLE_DEFAULTS = {"AGENTS.md", "USER.md"}
_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
def __init__(self, workspace: Path, timezone: str | None = None, disabled_skills: list[str] | None = None):
@@ -98,9 +109,6 @@ class ContextBuilder:
session_summary: SessionSummary | None = None,
workspace: Path | None = None,
include_memory: bool = True,
include_memory_recent_history: bool = True,
session_key: str | None = None,
unified_session: bool = False,
) -> str:
"""Build the system prompt from identity, bootstrap files, memory, and skills."""
root = workspace or self.workspace
@@ -138,29 +146,6 @@ class ContextBuilder:
if 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:
parts.append(
"[Archived Context Summary]\n\n"
@@ -170,25 +155,6 @@ class ContextBuilder:
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:
"""Get the core identity section."""
root = workspace or self.workspace
@@ -278,46 +244,68 @@ class ContextBuilder:
runtime_context_blocks: Sequence[RuntimeContextBlock] | None = None,
workspace: Path | None = None,
include_memory: bool = True,
include_memory_recent_history: bool = True,
session_key: str | None = None,
unified_session: bool = False,
) -> list[dict[str, Any]]:
"""Build the complete message list for an LLM call."""
"""Compatibility wrapper for callers that need merged adjacent roles."""
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
messages: list[dict[str, Any]] = [
{
"role": "system",
"content": self.build_system_prompt(
channel=channel,
session_summary=session_summary,
session_summary=transcript.session_summary,
workspace=root,
include_memory=include_memory,
include_memory_recent_history=include_memory_recent_history,
session_key=session_key,
unified_session=unified_session,
),
},
*history,
*transcript.history,
]
current = self.build_current_message(
current_message,
media=media,
current_role=current_role,
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
if transcript.current_message is None:
return messages
current = self.build_current_message(
transcript.current_message,
media=list(transcript.media) if transcript.media else None,
current_role=transcript.current_role,
runtime_context_blocks=transcript.runtime_context_blocks,
)
messages.append(current)
return messages
+101 -134
View File
@@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, cast
from loguru import logger
from nanobot.providers.base import LLMUsage
from nanobot.utils.helpers import (
estimate_message_tokens,
estimate_prompt_tokens_chain,
@@ -27,12 +28,6 @@ if TYPE_CHECKING:
from nanobot.providers.base import LLMProvider
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.
TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"})
BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]"
@@ -41,6 +36,27 @@ 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:
"""Whether a persisted OpenAI-style tool_call carries a usable name.
@@ -67,7 +83,6 @@ class ContextGovernanceConfig:
context_window_tokens: int | None = None
context_block_limit: int | None = None
max_tokens: int | None = None
inflight_start_index: int = 0
class ContextGovernor:
@@ -77,17 +92,85 @@ class ContextGovernor:
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
compacted_tool_call_ids: set[str],
) -> list[dict[str, Any]]:
updated = self.strip_placeholder_assistant_messages(messages)
updated = self.strip_malformed_tool_calls(updated)
updated = self.drop_orphan_tool_results(updated)
updated = self.backfill_missing_tool_results(updated)
updated = self.apply_tool_result_budget(config, updated)
updated = self.compact_inflight_overflow(config, updated, compacted_tool_call_ids)
updated = self.snip_history(config, updated)
return self.apply_tool_result_budget(config, updated)
def fit_to_budget(
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)
return self.backfill_missing_tool_results(updated)
updated = 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
def input_budget(config: ContextGovernanceConfig) -> int:
@@ -326,71 +409,13 @@ class ContextGovernor:
updated[idx]["content"] = normalized
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(
self,
config: ContextGovernanceConfig,
messages: list[dict[str, Any]],
*,
tool_definitions: list[dict[str, Any]] | None,
force: bool = False,
) -> list[dict[str, Any]]:
if not messages or not config.context_window_tokens:
return messages
@@ -399,14 +424,13 @@ class ContextGovernor:
if budget <= 0:
return messages
tools = config.tools.get_definitions()
estimate, _ = estimate_prompt_tokens_chain(
config.provider,
config.model,
messages,
tools,
tool_definitions,
)
if estimate <= budget:
if not force and estimate <= budget:
return messages
system_messages = [dict(msg) for msg in messages if msg.get("role") == "system"]
@@ -419,7 +443,7 @@ class ContextGovernor:
config.provider,
config.model,
system_messages,
tools,
tool_definitions,
)
remaining_budget = max(0, budget - max(system_tokens, fixed_tokens))
kept: list[dict[str, Any]] = []
@@ -434,16 +458,6 @@ class ContextGovernor:
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(
self,
kept: list[dict[str, Any]],
@@ -462,50 +476,3 @@ class ContextGovernor:
if messages[idx].get("role") == "user":
return messages[idx:]
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])
+54 -64
View File
@@ -14,6 +14,7 @@ from collections.abc import Coroutine, Iterable, Mapping
from contextlib import AbstractContextManager, ExitStack, nullcontext, suppress
from dataclasses import dataclass, field
from enum import Enum, auto
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, Awaitable, Callable, TypeVar, cast
@@ -23,7 +24,7 @@ from nanobot.agent import context as agent_context
from nanobot.agent import model_presets as preset_helpers
from nanobot.agent.autocompact import AutoCompact
from nanobot.agent.automation_turns import publish_next_deferred_turn
from nanobot.agent.context import ContextBuilder, PersistedPromptContextResolver
from nanobot.agent.context import ContextBuilder, PersistedPromptContextResolver, TranscriptInput
from nanobot.agent.cron_turns import CronTurnCoordinator
from nanobot.agent.hook import AgentHook, AgentTurnHookFactory
from nanobot.agent.memory import Consolidator
@@ -38,10 +39,9 @@ from nanobot.agent.subagent import SubagentManager
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.file_state import FileStateStore, bind_file_states, reset_file_states
from nanobot.agent.tools.message import MessageTool
from nanobot.agent.tools.message import capture_message_deliveries
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.agent.tools.runtime_control import AgentRuntimeControl
from nanobot.agent.tools.self import MyTool
from nanobot.agent.turn_delivery import (
TurnDelivery,
TurnDeliveryFactory,
@@ -136,7 +136,7 @@ class TurnContext:
session: Session | None = None
history: list[dict[str, Any]] = field(default_factory=list)
initial_messages: list[dict[str, Any]] = field(default_factory=list)
transcript_input: TranscriptInput | None = None
provider_state: ProviderConversationState | None = field(default=None, repr=False)
request_context: RequestContext | None = None
runtime_context_blocks: list[RuntimeContextBlock] = field(default_factory=list)
@@ -145,7 +145,6 @@ class TurnContext:
final_content: str | None = None
all_messages: list[dict[str, Any]] = field(default_factory=list)
stop_reason: str = ""
had_injections: bool = False
streamed_content: bool = False
input_persisted_early: bool = False
@@ -275,7 +274,6 @@ class AgentLoop:
channels_config: ChannelsConfig | None = None,
timezone: str | None = None,
session_ttl_minutes: int = 0,
consolidation_ratio: float = 0.5,
hooks: list[AgentHook] | None = None,
hook_factories: list[AgentTurnHookFactory] | None = None,
unified_session: bool = False,
@@ -432,8 +430,8 @@ class AgentLoop:
("cron", self._cron_turns),
("local trigger", self._local_trigger_turns),
)
# NANOBOT_MAX_CONCURRENT_REQUESTS: <=0 means unlimited; default 3.
_max = int(os.environ.get("NANOBOT_MAX_CONCURRENT_REQUESTS", "3"))
# NANOBOT_MAX_CONCURRENT_REQUESTS: unset or <=0 means unlimited.
_max = int(os.environ.get("NANOBOT_MAX_CONCURRENT_REQUESTS", "0"))
self._concurrency_gate: asyncio.Semaphore | None = (
asyncio.Semaphore(_max) if _max > 0 else None
)
@@ -446,8 +444,6 @@ class AgentLoop:
workspace_scopes=self.workspace_scopes,
unified_session=unified_session,
),
consolidation_ratio=consolidation_ratio,
unified_session=unified_session,
)
self.auto_compact = AutoCompact(
sessions=self.sessions,
@@ -519,7 +515,6 @@ class AgentLoop:
disabled_skills=defaults.disabled_skills,
session_ttl_minutes=defaults.session_ttl_minutes,
idle_compact_check_interval_seconds=defaults.idle_compact_check_interval_seconds,
consolidation_ratio=defaults.consolidation_ratio,
tools_config=config.tools,
model_presets=preset_helpers.configured_model_presets(config),
model_preset=defaults.model_preset,
@@ -644,20 +639,11 @@ class AgentLoop:
timezone=self.context.timezone or "UTC",
workspace_sandbox=self.workspace_scopes.sandbox_status,
runtime_events=self.runtime_events,
runtime_control=AgentRuntimeControl(self),
)
loader = ToolLoader()
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)
def register_runtime_context_provider(
@@ -737,22 +723,15 @@ class AgentLoop:
return True
return False
def _build_initial_messages(self, ctx: TurnContext) -> list[dict[str, Any]]:
"""Build the initial message list for the LLM turn."""
def _build_transcript_input(self, ctx: TurnContext) -> TranscriptInput:
"""Capture the persisted history and fresh input as separate transcript parts."""
assert ctx.session is not None
scope = self.workspace_scopes.for_message(ctx.msg, ctx.session.metadata)
return self.context.build_messages(
return TranscriptInput(
history=ctx.history,
current_message=ctx.msg.content,
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,
workspace=scope.project_path,
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:
@@ -943,7 +922,7 @@ class AgentLoop:
async def _run_agent_loop(
self,
initial_messages: list[dict[str, Any]],
transcript_input: TranscriptInput,
on_progress: Callable[..., Awaitable[None]] | None = None,
on_stream: Callable[[str], Awaitable[None]] | None = None,
on_stream_end: Callable[..., Awaitable[None]] | None = None,
@@ -1124,6 +1103,12 @@ class AgentLoop:
message_metadata=request_metadata,
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,
@@ -1170,11 +1155,13 @@ class AgentLoop:
run_extra_hooks_for_ephemeral=run_extra_hooks_for_ephemeral,
))
result = await self.runner.run(AgentRunSpec(
initial_messages=initial_messages,
initial_messages=None,
tools=effective_tools,
runtime=runtime,
max_iterations=self.max_iterations,
max_tool_result_chars=self.max_tool_result_chars,
transcript_input=transcript_input,
transcript_builder=transcript_builder,
hook=hook,
concurrent_tools=True,
workspace=effective_scope.project_path,
@@ -1733,18 +1720,12 @@ class AgentLoop:
msg: InboundMessage,
final_content: str,
stop_reason: str,
had_injections: bool,
streamed_content: bool,
*,
log_content: bool = True,
turn_latency_ms: int | None = None,
) -> OutboundMessage | None:
"""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:
preview = final_content[:120] + "..." if len(final_content) > 120 else final_content
logger.info("Response to {}:{}: {}", msg.channel, msg.sender_id, preview)
@@ -1898,12 +1879,15 @@ class AgentLoop:
session,
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"
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] = {
"max_tokens": self._replay_token_budget(runtime),
"extend_to_user": is_subagent,
@@ -1992,7 +1976,7 @@ class AgentLoop:
# Upgrade the replay-safe baseline to the resumable state before
# prompt assembly and the first model checkpoint.
self.sessions.save(session)
ctx.initial_messages = self._build_initial_messages(ctx)
ctx.transcript_input = self._build_transcript_input(ctx)
if ctx.on_progress is None:
ctx.on_progress = ctx.delivery.progress_callback()
@@ -2004,28 +1988,35 @@ class AgentLoop:
if ctx.visible_run_started_at is None:
ctx.visible_run_started_at = time.time()
await ctx.delivery.running(started_at=ctx.visible_run_started_at)
result = await self._run_agent_loop(
ctx.initial_messages,
runtime=runtime,
on_progress=ctx.on_progress,
on_stream=ctx.on_stream,
on_stream_end=ctx.on_stream_end,
on_retry_wait=ctx.on_retry_wait,
session=ctx.session,
pending_queue=ctx.pending_queue,
ephemeral=ctx.ephemeral,
run_extra_hooks_for_ephemeral=ctx.run_extra_hooks_for_ephemeral,
hooks=ctx.hooks,
hook_factories=ctx.hook_factories,
turn_scopes=ctx.turn_scopes,
tools=ctx.tools,
request_context=ctx.request_context,
provider_state=ctx.provider_state,
)
assert ctx.transcript_input is not None
with capture_message_deliveries() as message_sends:
result = await self._run_agent_loop(
ctx.transcript_input,
runtime=runtime,
on_progress=ctx.on_progress,
on_stream=ctx.on_stream,
on_stream_end=ctx.on_stream_end,
on_retry_wait=ctx.on_retry_wait,
session=ctx.session,
pending_queue=ctx.pending_queue,
ephemeral=ctx.ephemeral,
run_extra_hooks_for_ephemeral=ctx.run_extra_hooks_for_ephemeral,
hooks=ctx.hooks,
hook_factories=ctx.hook_factories,
turn_scopes=ctx.turn_scopes,
tools=ctx.tools,
request_context=ctx.request_context,
provider_state=ctx.provider_state,
)
ctx.final_content = result.final_content
ctx.all_messages = result.messages
ctx.stop_reason = result.stop_reason
ctx.had_injections = result.had_injections
if (
ctx.kind is TurnKind.USER
and (ctx.delivery.route.channel, ctx.delivery.route.chat_id) in message_sends
and (not result.had_injections or result.stop_reason == "empty_final_response")
):
ctx.suppress_response = True
ctx.usage = result.usage
ctx.delivery.record_usage(ctx.usage)
if ctx.kind is TurnKind.USER:
@@ -2094,7 +2085,6 @@ class AgentLoop:
ctx.delivery.delivery_message,
cast(str, ctx.final_content),
ctx.stop_reason,
ctx.had_injections,
ctx.streamed_content,
log_content=ctx.require_session().policy.log_content,
turn_latency_ms=ctx.turn_latency_ms,
+319 -297
View File
@@ -1,4 +1,4 @@
"""Memory system: pure file I/O store and lightweight Consolidator."""
"""Memory storage, transcript archiving, and legacy consolidation coordination."""
# Tool schemas are installed by the ``@tool_parameters`` class decorator at
# runtime; static analyzers cannot observe that it clears ``parameters`` from
@@ -32,10 +32,10 @@ from nanobot.utils.gitstore import GitStore
from nanobot.utils.helpers import (
content_with_media_breadcrumbs,
ensure_dir,
estimate_message_tokens,
estimate_prompt_tokens_chain,
strip_think,
truncate_text,
truncate_text_to_tokens,
)
from nanobot.utils.prompt_templates import render_template
from nanobot.utils.workspace_prompts import (
@@ -66,8 +66,6 @@ class MemoryStore:
# durable files are tiny in practice (~5 KB total), but a runaway file must
# not unbounded the prompt.
_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_TIMESTAMP_RE = re.compile(r"^\[(\d{4}-\d{2}-\d{2} \d{2}:\d{2})\]\s*")
_LEGACY_RAW_MESSAGE_RE = re.compile(
@@ -261,6 +259,29 @@ class MemoryStore:
# -- 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(
self,
entry: str,
@@ -275,27 +296,16 @@ class MemoryStore:
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
to the raw leak otherwise `strip_think`'s guarantees would be
undone by history replay / consolidation downstream.
undone when Dream consumes the journal entry.
A defensive cap (*max_chars*, default ``_HISTORY_ENTRY_HARD_CAP``) is
applied as a final safety net: individual callers should cap their own
content more tightly; this default only exists to catch unintentional
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")
raw = entry.rstrip()
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)
content = self._normalize_history_entry(entry, max_chars=max_chars)
# Cursor allocation and the append must be atomic: concurrent writers
# could otherwise read the same current cursor and emit duplicates.
with self._append_lock:
@@ -303,7 +313,7 @@ class MemoryStore:
if raw and not content:
logger.debug(
"history entry {} stripped to empty (likely template leak); "
"persisting empty content to avoid re-polluting context",
"persisting empty content to avoid re-polluting Dream input",
cursor,
)
record = {"cursor": cursor, "timestamp": ts, "content": content}
@@ -393,36 +403,6 @@ class MemoryStore:
"""Return history entries with a valid cursor > *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:
"""Drop oldest processed entries without discarding pending Dream input."""
if self.max_history_entries <= 0:
@@ -719,21 +699,28 @@ class MemoryStore:
*,
max_chars: int | None = None,
session_key: str | None = None,
) -> None:
"""Fallback: dump raw messages to history.jsonl without LLM summarization."""
limit = max_chars if max_chars is not None else _RAW_ARCHIVE_MAX_CHARS
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,
)
) -> str:
"""Persist and return a bounded raw checkpoint when summarization degrades."""
checkpoint = self._build_raw_checkpoint(messages, max_chars=max_chars)
self.append_history(checkpoint, session_key=session_key)
logger.warning(
"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
@@ -785,21 +772,215 @@ class MemoryStore:
# ---------------------------------------------------------------------------
# Consolidator — lightweight token-budget triggered consolidation
# Memory ingestion and legacy context-pressure coordination
# ---------------------------------------------------------------------------
# Individual history.jsonl writers cap their own payloads tightly; the
# _HISTORY_ENTRY_HARD_CAP at append_history() is a belt-and-suspenders default
# that catches any new caller that forgot to set its own cap.
_RAW_ARCHIVE_MAX_CHARS = 16_000 # fallback dump (LLM failed)
_ARCHIVE_SUMMARY_MAX_CHARS = 8_000 # LLM-produced consolidation summary
_HISTORY_ENTRY_HARD_CAP = 64_000 # emergency cap in append_history
# Raw fallbacks use a tighter cap. Completed model summaries may scale with the
# configured generation budget, while append_history() still enforces the
# emergency hard cap against pathological provider output.
_RAW_ARCHIVE_MAX_CHARS = 16_000 # fallback dump (LLM failed)
_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:
"""Summarize compacted messages into history.jsonl."""
_MAX_CONSOLIDATION_ROUNDS = 5
"""Legacy context-pressure coordinator backed by a MemoryArchiver."""
_SAFETY_BUFFER = 1024 # extra headroom for tokenizer estimation drift
@@ -810,16 +991,17 @@ class Consolidator:
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,
consolidation_ratio: float = 0.5,
unified_session: bool = False,
):
self.store = store
self.sessions = sessions
self.consolidation_ratio = consolidation_ratio
self.unified_session = unified_session
self._build_messages = build_messages
self._get_tool_definitions = get_tool_definitions
self._resolve_prompt_context = resolve_prompt_context
self.archiver = MemoryArchiver(
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] = (
weakref.WeakValueDictionary()
)
@@ -831,24 +1013,19 @@ class Consolidator:
def pick_consolidation_boundary(
self,
session: Session,
tokens_to_remove: int,
) -> tuple[int, int] | None:
"""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:
) -> int | None:
"""Return the fixed user-led boundary before the recent replay tail."""
if not session.messages:
return None
removed_tokens = 0
last_boundary: tuple[int, int] | None = None
for idx in range(start, len(session.messages)):
message = session.messages[idx]
if idx > start and message.get("role") == "user":
last_boundary = (idx, removed_tokens)
if removed_tokens >= tokens_to_remove:
return last_boundary
removed_tokens += estimate_message_tokens(message)
return last_boundary
boundary = max(0, len(session.messages) - MIN_COMPACTED_REPLAY_MESSAGES)
while boundary > 0 and session.messages[boundary].get("role") != "user":
boundary -= 1
if (
boundary <= session.last_archived
or session.messages[boundary].get("role") != "user"
):
return None
return boundary
@staticmethod
def _full_replay_history(
@@ -859,13 +1036,18 @@ class Consolidator:
return []
return session.get_history()
def _persist_last_summary(self, session: Session, summary: str | None) -> None:
if summary and summary != "(nothing)":
@staticmethod
def _set_last_summary(
session: Session,
summary: str,
*,
last_active: datetime | None = None,
) -> None:
if summary != "(nothing)":
session.metadata["_last_summary"] = {
"text": summary,
"last_active": session.updated_at.isoformat(),
"last_active": (last_active or session.updated_at).isoformat(),
}
self.sessions.save(session)
def estimate_session_prompt_tokens(
self,
@@ -885,8 +1067,6 @@ class Consolidator:
current_message="[token-probe]",
channel=channel,
session_summary=summary,
session_key=session.key,
unified_session=self.unified_session,
)
return estimate_prompt_tokens_chain(
runtime.provider,
@@ -903,58 +1083,6 @@ class Consolidator:
- 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(
self,
session: Session,
@@ -962,82 +1090,12 @@ class Consolidator:
archive_end: int,
runtime: LLMRuntime,
) -> str | None:
"""Archive a session prefix by appending a consolidation instruction."""
messages = list(session.messages[session.last_consolidated:archive_end])
if not messages:
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,
"""Compatibility wrapper for the extracted MemoryArchiver."""
return await self.archiver.archive_session(
session,
archive_end=archive_end,
runtime=runtime,
session_key=session.key,
request_messages=request_messages,
request_tools=tools,
input_token_budget=self._input_token_budget(runtime),
)
async def maybe_consolidate_by_tokens(
@@ -1046,104 +1104,71 @@ class Consolidator:
*,
runtime: LLMRuntime,
) -> None:
"""Loop: archive old messages until prompt fits within safe budget.
"""Archive one fixed old prefix when the prompt exceeds the safe budget.
The budget reserves space for completion tokens and a safety buffer
so the LLM request never exceeds the context window.
"""
if runtime.context_window_tokens <= 0:
return
lock = self.get_lock(session.key)
async with lock:
# Refresh session reference: AutoCompact may have replaced it.
fresh = self.sessions.get_or_create(session.key)
if fresh is not session:
session = fresh
if runtime.context_window_tokens <= 0:
return
if not session.messages:
return
budget = self._input_token_budget(runtime)
target = int(budget * self.consolidation_ratio)
last_summary: str | None = None
estimated, source = self.estimate_session_prompt_tokens(
session,
runtime=runtime,
)
if estimated <= 0:
self._persist_last_summary(session, last_summary)
return
if estimated < budget:
unconsolidated_count = len(session.messages) - session.last_consolidated
unarchived_count = len(session.messages) - session.last_archived
logger.debug(
"Token consolidation idle {}: {}/{} via {}, msgs={}",
session.key,
estimated,
runtime.context_window_tokens,
source,
unconsolidated_count,
unarchived_count,
)
self._persist_last_summary(session, last_summary)
return
for round_num in range(self._MAX_CONSOLIDATION_ROUNDS):
if estimated <= target:
break
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,
end_idx = self.pick_consolidation_boundary(session)
if end_idx is None:
logger.debug(
"Token consolidation: no safe fixed boundary for {}",
session.key,
estimated,
runtime.context_window_tokens,
source,
len(chunk),
)
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
self.sessions.save(session)
if not summary:
# LLM is degraded — stop hammering it this call;
# the next invocation can retry a fresh chunk.
break
return
estimated, source = self.estimate_session_prompt_tokens(
session,
runtime=runtime,
)
if estimated <= 0:
break
chunk = session.messages[session.last_archived:end_idx]
if not chunk:
return
# Persist the last summary to session metadata so it can be injected
# into the runtime context on the next prepare_session() call, aligning
# the summary injection strategy with AutoCompact._archive().
self._persist_last_summary(session, last_summary)
logger.info(
"Token consolidation for {}: {}/{} via {}, chunk={} msgs",
session.key,
estimated,
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(
self,
@@ -1170,7 +1195,7 @@ class Consolidator:
self.sessions.invalidate(session_key)
session = self.sessions.get_or_create(session_key)
archive_start = session.last_consolidated
archive_start = session.last_archived
messages_to_archive = list(session.messages[archive_start:])
if not messages_to_archive:
return ""
@@ -1182,17 +1207,14 @@ class Consolidator:
archive_end=archive_end,
runtime=runtime,
)
if summary is None:
return None
if summary and summary != "(nothing)":
session.metadata["_last_summary"] = {
"text": summary,
"last_active": last_active.isoformat(),
}
self._set_last_summary(session, summary, last_active=last_active)
# A turn can append while the provider call is in flight. Advance only
# through the captured batch so new messages remain eligible next time.
session.last_consolidated = archive_end
session.provider_state = None
session.last_archived = archive_end
self.sessions.save(session)
visible = session.get_history(
+202 -392
View File
@@ -14,12 +14,14 @@ from typing import Any, cast
from loguru import logger
from nanobot.agent.context import TranscriptInput
from nanobot.agent.context_governance import (
ContextGovernanceConfig,
ContextGovernor,
)
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
from nanobot.agent.tools.execution import execute_tool_calls
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.llm_usage.context import (
LLMUsageSource,
bind_llm_usage_source,
@@ -32,7 +34,6 @@ from nanobot.providers.base import (
LLMUsage,
ProviderCallContext,
ProviderConversationState,
ToolCallRequest,
)
from nanobot.providers.conversation_state import (
ProviderConversationStateController,
@@ -46,13 +47,11 @@ from nanobot.runtime_context import (
from nanobot.session.history_visibility import is_hidden_history_message
from nanobot.session.recovery import PENDING_FOLLOWUP_ID_KEY
from nanobot.utils.helpers import (
IncrementalThinkExtractor,
build_assistant_message,
estimate_message_tokens,
estimate_prompt_tokens_chain,
extract_reasoning,
strip_reasoning_tags,
strip_think,
)
from nanobot.utils.llm_runtime import LLMRuntime
from nanobot.utils.prompt_templates import render_template
@@ -62,15 +61,13 @@ from nanobot.utils.runtime import (
build_finalization_retry_message,
build_length_recovery_message,
is_blank_text,
repeated_external_lookup_error,
repeated_workspace_violation_error,
)
ContinuationCallback = Callable[[], str | None]
ProgressCallback = Callable[[str], Awaitable[None]]
RetryWaitCallback = Callable[[str], Awaitable[None]]
CheckpointCallback = Callable[[dict[str, Any]], Awaitable[None]]
InjectionCallback = Callable[..., Awaitable[Iterable[Any] | None]]
TranscriptBuilder = Callable[[TranscriptInput], list[dict[str, Any]]]
_DEFAULT_ERROR_MESSAGE = "Sorry, I encountered an error calling the AI model."
_ARREARAGE_ERROR_MESSAGE = (
@@ -99,11 +96,13 @@ def _restore_outer_whitespace(content: str, original: str | None) -> str:
class AgentRunSpec:
"""Configuration for a single agent execution."""
initial_messages: list[dict[str, Any]]
initial_messages: list[dict[str, Any]] | None
tools: ToolRegistry
runtime: LLMRuntime
max_iterations: int
max_tool_result_chars: int
transcript_input: TranscriptInput | None = None
transcript_builder: TranscriptBuilder | None = None
hook: AgentHook | None = None
error_message: str | None = _DEFAULT_ERROR_MESSAGE
max_iterations_message: str | None = None
@@ -112,7 +111,6 @@ class AgentRunSpec:
session_key: str | None = None
context_block_limit: int | None = None
provider_retry_mode: str = "standard"
progress_callback: ProgressCallback | None = None
retry_wait_callback: RetryWaitCallback | None = None
checkpoint_callback: CheckpointCallback | None = None
injection_callback: InjectionCallback | None = None
@@ -141,6 +139,17 @@ class AgentRunResult:
provider_state: ProviderConversationState | None = field(default=None, repr=False)
@dataclass(slots=True)
class _ModelRequestState:
"""Per-run state used to govern the next provider request."""
config: ContextGovernanceConfig
conversation: ProviderConversationStateController
usage: LLMUsage | None = None
messages: list[dict[str, Any]] | None = None
tool_definitions: list[dict[str, Any]] | None = None
class AgentRunner:
"""Run a tool-capable LLM loop without product-layer concerns."""
@@ -416,7 +425,7 @@ class AgentRunner:
async def run(self, spec: AgentRunSpec) -> AgentRunResult:
hook = spec.hook or AgentHook()
messages = list(spec.initial_messages)
messages = self._initial_transcript(spec)
context = AgentRunHookContext(messages=deepcopy(messages))
llm_usage_source_token = bind_llm_usage_source(
spec.llm_usage_source or source_from_session_key(spec.session_key)
@@ -468,6 +477,19 @@ class AgentRunner:
finally:
reset_llm_usage_source(llm_usage_source_token)
@staticmethod
def _initial_transcript(spec: AgentRunSpec) -> list[dict[str, Any]]:
"""Resolve exactly one supported source for the initial model transcript."""
if spec.transcript_input is not None:
if spec.initial_messages is not None:
raise ValueError("provide either transcript_input or initial_messages, not both")
if spec.transcript_builder is None:
raise ValueError("transcript_builder is required with transcript_input")
return list(spec.transcript_builder(spec.transcript_input))
if spec.initial_messages is None:
raise ValueError("initial_messages is required without transcript_input")
return list(spec.initial_messages)
async def _run_core(
self,
spec: AgentRunSpec,
@@ -489,7 +511,6 @@ class AgentRunner:
length_recovery_parts: list[str] = []
had_injections = False
injection_cycles = 0
compacted_tool_call_ids: set[str] = set()
pending_stream_content: str | None = None
conversation_state = ProviderConversationStateController(
provider=spec.runtime.provider,
@@ -508,39 +529,29 @@ class AgentRunner:
context_window_tokens=spec.runtime.context_window_tokens,
context_block_limit=spec.context_block_limit,
max_tokens=spec.runtime.generation.max_tokens,
inflight_start_index=len(spec.initial_messages),
)
request_state = _ModelRequestState(
config=governance_config,
conversation=conversation_state,
)
for iteration in range(spec.max_iterations):
# Keep the persisted conversation untouched. Context governance
# may repair or compact historical messages for the model, but
# those synthetic edits must not shift the append boundary used
# later when the caller saves only the new turn. A governance
# failure must stop the run instead of sending an ungoverned copy.
messages_for_model = self.context_governor.prepare_for_model(
governance_config,
messages,
compacted_tool_call_ids,
)
context = AgentHookContext(
iteration=iteration,
messages=messages,
session_key=spec.session_key,
)
await hook.before_iteration(context)
provider_context = conversation_state.prepare_request(
messages,
context_window_tokens=spec.runtime.context_window_tokens,
model_messages=messages_for_model,
)
response = await self._request_model(
spec,
messages_for_model,
messages,
hook,
context,
conversation_state=conversation_state,
provider_context=provider_context,
request_state=request_state,
transcript=messages,
)
assert request_state.messages is not None
messages_for_model = request_state.messages
conversation_state.observe_response(response, messages)
context.response = response
context.tool_calls = list(response.tool_calls)
@@ -552,7 +563,7 @@ class AgentRunner:
response.content,
)
response.content = cleaned_content
raw_usage = self._usage_or_estimate(spec, messages_for_model, response)
raw_usage = self._record_request_usage(spec, request_state, response)
context.usage = raw_usage
usage = self._merge_usage(usage, raw_usage)
if reasoning_text and not context.streamed_reasoning:
@@ -590,13 +601,14 @@ class AgentRunner:
await hook.before_execute_tools(context)
results, new_events = await self._execute_tools(
spec,
results, new_events = await execute_tool_calls(
spec.tools,
response.tool_calls,
external_lookup_counts,
workspace_violation_counts,
hook,
context,
concurrent=spec.concurrent_tools,
external_lookup_counts=external_lookup_counts,
workspace_violation_counts=workspace_violation_counts,
hook=hook,
context=context,
)
tool_events.extend(new_events)
tools_used.extend(
@@ -625,7 +637,6 @@ class AgentRunner:
self.context_governor.prepare_for_model(
governance_config,
messages,
compacted_tool_call_ids,
)
if response.provider_state is not None
else None
@@ -691,14 +702,13 @@ class AgentRunner:
)
if hook.wants_streaming():
await hook.on_stream_end(context, resuming=False)
retry_messages = self._finalization_retry_messages(messages_for_model)
response = await self._request_finalization_retry(
spec,
messages_for_model,
request_state=request_state,
transcript=messages,
conversation_state=conversation_state,
)
retry_usage = self._usage_or_estimate(spec, retry_messages, response)
retry_usage = self._record_request_usage(spec, request_state, response)
usage = self._merge_usage(usage, retry_usage)
raw_usage = self._merge_usage(raw_usage, retry_usage)
context.response = response
@@ -885,7 +895,7 @@ class AgentRunner:
hook,
messages,
usage,
conversation_state,
request_state=request_state,
)
if terminal_content is None:
terminal_content = self._max_iterations_fallback(spec)
@@ -932,6 +942,60 @@ class AgentRunner:
kwargs["reasoning_effort"] = generation.reasoning_effort
return kwargs
def _prepare_model_request(
self,
state: _ModelRequestState,
messages: list[dict[str, Any]],
*,
tool_definitions: list[dict[str, Any]] | None,
transcript: list[dict[str, Any]] | None = None,
) -> tuple[list[dict[str, Any]], ProviderCallContext | None]:
"""Prepare, fit, and record the exact payload sent to a provider."""
prepared = self.context_governor.prepare_for_model(state.config, messages)
supplemental_messages = (
[prepared[-1]] if transcript is not None and tool_definitions is None else None
)
model_messages = None if supplemental_messages is not None else prepared
request_context_tokens = (
state.conversation.estimate_request_context_tokens(
transcript,
model_messages=model_messages,
supplemental_messages=supplemental_messages,
tool_definitions=tool_definitions,
)
if transcript is not None
else None
)
usage_matches_messages = (
state.messages is not None
and prepared == state.messages
and tool_definitions == state.tool_definitions
)
prepared, fitted = self.context_governor.fit_request(
state.config,
prepared,
state.usage,
usage_matches_messages=usage_matches_messages,
tool_definitions=tool_definitions,
request_context_tokens=request_context_tokens,
)
provider_context = (
state.conversation.prepare_request(
transcript,
context_window_tokens=state.config.context_window_tokens,
model_messages=model_messages,
supplemental_messages=supplemental_messages,
resume_state=not fitted,
)
if transcript is not None
else state.conversation.independent_request_context(
context_window_tokens=state.config.context_window_tokens,
)
)
state.messages = deepcopy(prepared)
state.tool_definitions = deepcopy(tool_definitions)
return prepared, provider_context
async def _request_model(
self,
spec: AgentRunSpec,
@@ -939,27 +1003,29 @@ class AgentRunner:
hook: AgentHook,
context: AgentHookContext,
*,
request_state: _ModelRequestState,
malformed_retry: bool = False,
conversation_state: ProviderConversationStateController,
provider_context: ProviderCallContext | None = None,
transcript: list[dict[str, Any]] | None,
) -> LLMResponse:
timeout_s = self._resolve_llm_timeout_s(spec)
tool_definitions = spec.tools.get_definitions()
messages, provider_context = self._prepare_model_request(
request_state,
messages,
tool_definitions=tool_definitions,
transcript=transcript,
)
kwargs = self._build_request_kwargs(
spec,
messages,
tools=spec.tools.get_definitions(),
tools=tool_definitions,
)
wants_streaming = hook.wants_streaming()
progress_callback = spec.progress_callback
wants_progress_streaming = (
not wants_streaming
and progress_callback is not None
and getattr(spec.runtime.provider, "supports_progress_deltas", False) is True
)
progress_state: dict[str, bool] | None = None
active_hosted_tools: dict[str, dict[str, Any]] = {}
native_reasoning_open = False
native_reasoning_close_task: asyncio.Task[None] | None = None
request_started_at = 0.0
first_output_at: float | None = None
generation_started_at: float | None = None
@@ -982,9 +1048,35 @@ class AgentRunner:
generation_elapsed_s += max(0.0, time.perf_counter() - generation_started_at)
generation_started_at = None
async def _close_native_reasoning() -> None:
nonlocal native_reasoning_open, native_reasoning_close_task
if native_reasoning_close_task is None:
if not native_reasoning_open:
return
native_reasoning_open = False
native_reasoning_close_task = asyncio.create_task(
hook.emit_reasoning_end()
)
close_task = native_reasoning_close_task
cancellation: asyncio.CancelledError | None = None
while not close_task.done():
try:
await asyncio.shield(close_task)
except asyncio.CancelledError as exc:
cancellation = cancellation or exc
try:
close_task.result()
finally:
if native_reasoning_close_task is close_task:
native_reasoning_close_task = None
if cancellation is not None:
raise cancellation
async def _provider_tool_event(event: dict[str, Any]) -> None:
if event.get("kind") != "hosted_tool":
return
await _close_native_reasoning()
await hook.on_provider_tool_event(context, event)
call_id = event.get("call_id")
if not call_id:
@@ -1002,10 +1094,11 @@ class AgentRunner:
_generation_delta(delta)
if delta:
context.streamed_content = True
await _close_native_reasoning()
await hook.on_stream(context, delta)
async def _thinking(delta: str) -> None:
nonlocal thinking_buf
nonlocal native_reasoning_open, thinking_buf
if not delta:
return
_generation_delta(delta)
@@ -1015,10 +1108,12 @@ class AgentRunner:
incremental = new_clean[len(prev_clean):]
if incremental:
context.streamed_reasoning = True
native_reasoning_open = True
await hook.emit_reasoning(incremental)
async def _stream_recover() -> None:
_pause_generation()
await _close_native_reasoning()
await hook.on_stream_end(context, resuming=True)
coro = spec.runtime.provider.chat_stream_with_retry(
@@ -1029,40 +1124,6 @@ class AgentRunner:
on_tool_call_delta=_provider_tool_event,
on_stream_recover=_stream_recover,
)
elif wants_progress_streaming:
stream_buf = ""
think_extractor = IncrementalThinkExtractor()
progress_state = {"reasoning_open": False}
async def _stream_progress(delta: str) -> None:
nonlocal stream_buf
if not delta:
return
_generation_delta(delta)
prev_clean = strip_think(stream_buf)
stream_buf += delta
new_clean = strip_think(stream_buf)
incremental = new_clean[len(prev_clean):]
if await think_extractor.feed(stream_buf, hook.emit_reasoning):
context.streamed_reasoning = True
progress_state["reasoning_open"] = True
if incremental:
if progress_state["reasoning_open"]:
await hook.emit_reasoning_end()
progress_state["reasoning_open"] = False
context.streamed_content = True
callback = progress_callback
if callback is not None:
await callback(incremental)
coro = spec.runtime.provider.chat_stream_with_retry(
**kwargs,
provider_context=provider_context,
on_content_delta=_stream_progress,
on_tool_call_delta=_provider_tool_event,
)
else:
coro = spec.runtime.provider.chat_with_retry(
**kwargs,
@@ -1074,10 +1135,9 @@ class AgentRunner:
# very slow deltas can still run forever. Use a more generous wall-clock
# timeout for streaming while preserving NANOBOT_LLM_TIMEOUT_S=0 as an
# opt-out for all LLM wall-clock timeouts.
is_streaming_request = wants_streaming or wants_progress_streaming
outer_timeout_s = (
max(300.0, timeout_s * 2)
if is_streaming_request and timeout_s is not None
if wants_streaming and timeout_s is not None
else timeout_s
)
request_started_at = time.perf_counter()
@@ -1086,6 +1146,10 @@ class AgentRunner:
await coro if outer_timeout_s is None
else await asyncio.wait_for(coro, timeout=outer_timeout_s)
)
except asyncio.CancelledError:
_pause_generation()
await _close_native_reasoning()
raise
except asyncio.TimeoutError:
if outer_timeout_s is None:
response = LLMResponse(
@@ -1100,6 +1164,7 @@ class AgentRunner:
error_kind="timeout",
)
_pause_generation()
await _close_native_reasoning()
if first_output_at is not None:
response.ttft_ms = max(0, round((first_output_at - request_started_at) * 1000))
if generation_elapsed_s > 0:
@@ -1115,8 +1180,6 @@ class AgentRunner:
"error": response.content
or "Model request failed before the provider-hosted tool completed.",
})
if progress_state and progress_state.get("reasoning_open"):
await hook.emit_reasoning_end()
dropped, all_dropped, original_finish_reason = (
self._drop_malformed_tool_calls(response)
)
@@ -1134,11 +1197,9 @@ class AgentRunner:
)
return await self._request_model(
spec, retry_messages, hook, context,
request_state=request_state,
malformed_retry=True,
conversation_state=conversation_state,
provider_context=conversation_state.independent_request_context(
context_window_tokens=spec.runtime.context_window_tokens,
),
transcript=None,
)
if (
all_dropped
@@ -1154,9 +1215,7 @@ class AgentRunner:
return await self._request_no_tools(
spec,
fallback_messages,
provider_context=conversation_state.independent_request_context(
context_window_tokens=spec.runtime.context_window_tokens,
),
request_state=request_state,
)
return response
@@ -1224,21 +1283,17 @@ class AgentRunner:
spec: AgentRunSpec,
messages: list[dict[str, Any]],
*,
request_state: _ModelRequestState,
transcript: list[dict[str, Any]],
conversation_state: ProviderConversationStateController,
) -> LLMResponse:
retry_messages = self._finalization_retry_messages(messages)
provider_context = conversation_state.prepare_request(
transcript,
context_window_tokens=spec.runtime.context_window_tokens,
supplemental_messages=[retry_messages[-1]],
)
response = await self._request_no_tools(
spec,
retry_messages,
provider_context=provider_context,
request_state=request_state,
transcript=transcript,
)
conversation_state.observe_response(
request_state.conversation.observe_response(
response,
transcript,
adopt_candidate_state=False,
@@ -1257,16 +1312,15 @@ class AgentRunner:
hook: AgentHook,
messages: list[dict[str, Any]],
usage: LLMUsage | None,
conversation_state: ProviderConversationStateController,
*,
request_state: _ModelRequestState,
) -> tuple[str | None, LLMUsage | None]:
retry_messages = self._budget_exhausted_finalization_messages(messages)
try:
response = await self._request_no_tools(
spec,
retry_messages,
provider_context=conversation_state.independent_request_context(
context_window_tokens=spec.runtime.context_window_tokens,
),
request_state=request_state,
)
except Exception:
logger.exception(
@@ -1275,7 +1329,7 @@ class AgentRunner:
)
return None, usage
raw_usage = self._usage_or_estimate(spec, retry_messages, response)
raw_usage = self._record_request_usage(spec, request_state, response)
usage = self._merge_usage(usage, raw_usage)
if response.finish_reason == "error" or response.has_tool_calls:
logger.warning(
@@ -1304,8 +1358,15 @@ class AgentRunner:
spec: AgentRunSpec,
messages: list[dict[str, Any]],
*,
provider_context: ProviderCallContext | None = None,
request_state: _ModelRequestState,
transcript: list[dict[str, Any]] | None = None,
) -> LLMResponse:
messages, provider_context = self._prepare_model_request(
request_state,
messages,
tool_definitions=None,
transcript=transcript,
)
kwargs = self._build_request_kwargs(
spec,
messages,
@@ -1317,17 +1378,18 @@ class AgentRunner:
)
timeout_s = self._resolve_llm_timeout_s(spec)
try:
return (
response = (
await coro
if timeout_s is None
else await asyncio.wait_for(coro, timeout=timeout_s)
)
except asyncio.TimeoutError:
return LLMResponse(
response = LLMResponse(
content=f"Error calling LLM: timed out after {timeout_s:g}s",
finish_reason="error",
error_kind="timeout",
)
return response
@staticmethod
def _resolve_llm_timeout_s(spec: AgentRunSpec) -> float | None:
@@ -1369,33 +1431,53 @@ class AgentRunner:
spec: AgentRunSpec,
messages: list[dict[str, Any]],
response: LLMResponse,
*,
tool_definitions: list[dict[str, Any]] | None,
) -> LLMUsage | None:
usage = response.usage
if response.finish_reason == "error":
if usage is None or usage.total_tokens == 0:
usage = LLMUsage.empty_request()
elif usage is None or usage.total_tokens == 0:
usage = self._estimate_response_usage(spec, messages, response)
usage = self._estimate_response_usage(
spec,
messages,
response,
tool_definitions=tool_definitions,
)
return usage.with_timing(
generation_ms=response.generation_ms,
ttft_ms=response.ttft_ms,
)
def _record_request_usage(
self,
spec: AgentRunSpec,
state: _ModelRequestState,
response: LLMResponse,
) -> LLMUsage | None:
assert state.messages is not None
state.usage = self._usage_or_estimate(
spec,
state.messages,
response,
tool_definitions=state.tool_definitions,
)
return state.usage
def _estimate_response_usage(
self,
spec: AgentRunSpec,
messages: list[dict[str, Any]],
response: LLMResponse,
*,
tool_definitions: list[dict[str, Any]] | None,
) -> LLMUsage:
try:
tools = spec.tools.get_definitions()
except Exception:
tools = None
prompt_tokens, _ = estimate_prompt_tokens_chain(
spec.runtime.provider,
spec.runtime.model,
messages,
tools,
tool_definitions,
)
assistant_message = build_assistant_message(
response.content or "",
@@ -1420,253 +1502,6 @@ class AgentRunner:
return left
return left + right
async def _execute_tools(
self,
spec: AgentRunSpec,
tool_calls: list[ToolCallRequest],
external_lookup_counts: dict[str, int],
workspace_violation_counts: dict[str, int],
hook: AgentHook | None = None,
context: AgentHookContext | None = None,
) -> tuple[list[Any], list[dict[str, str]]]:
hook = hook or AgentHook()
context = context or AgentHookContext(iteration=0, messages=[])
batches = self._partition_tool_batches(spec, tool_calls)
tool_results: list[tuple[Any, dict[str, str]]] = []
for batch in batches:
if spec.concurrent_tools and len(batch) > 1:
batch_results = await asyncio.gather(*(
self._run_tool(
spec,
tool_call,
external_lookup_counts,
workspace_violation_counts,
hook,
context,
)
for tool_call in batch
))
tool_results.extend(batch_results)
else:
batch_results: list[tuple[Any, dict[str, str]]] = []
for tool_call in batch:
result = await self._run_tool(
spec,
tool_call,
external_lookup_counts,
workspace_violation_counts,
hook,
context,
)
tool_results.append(result)
batch_results.append(result)
results: list[Any] = []
events: list[dict[str, str]] = []
for result, event in tool_results:
results.append(result)
events.append(event)
return results, events
async def _run_tool(
self,
spec: AgentRunSpec,
tool_call: ToolCallRequest,
external_lookup_counts: dict[str, int],
workspace_violation_counts: dict[str, int],
hook: AgentHook | None = None,
context: AgentHookContext | None = None,
) -> tuple[Any, dict[str, str]]:
hook = hook or AgentHook()
context = context or AgentHookContext(iteration=0, messages=[])
hint = "\n\n[Analyze the error above and try a different approach.]"
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 lookup_error + hint, event
prepare_call = cast(
Callable[[str, Any], object] | None,
getattr(spec.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:
event = {
"name": tool_call.name,
"status": "error",
"detail": prep_error.split(": ", 1)[-1][:120],
}
handled = self._classify_violation(
raw_text=prep_error,
soft_payload=prep_error + hint,
event=event,
tool_call=tool_call,
workspace_violation_counts=workspace_violation_counts,
)
if handled is not None:
return handled
return prep_error + hint, 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 spec.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 = f"Error: {type(exc).__name__}: {exc}"
handled = self._classify_violation(
raw_text=str(exc),
# Preserve legacy exception payloads without the retry hint.
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)
event = {
"name": tool_call.name,
"status": "error",
"detail": result.replace("\n", " ").strip()[:120],
}
handled = self._classify_violation(
raw_text=result,
soft_payload=result + hint,
event=event,
tool_call=tool_call,
workspace_violation_counts=workspace_violation_counts,
)
if handled is not None:
return handled
return result + hint, 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}
# 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: str = (
"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 LLM 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",
)
@classmethod
def _is_ssrf_violation(cls, text: str) -> bool:
if not text:
return False
lowered = text.lower()
return any(marker in lowered for marker in cls._SSRF_MARKERS)
@classmethod
def _is_workspace_violation(cls, text: str) -> bool:
"""True when *text* looks like any policy boundary rejection."""
if not text:
return False
lowered = text.lower()
if cls._is_ssrf_violation(lowered):
return True
return any(marker in lowered for marker in cls._WORKSPACE_VIOLATION_MARKERS)
def _classify_violation(
self,
*,
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:
"""Classify safety-boundary failures, or return ``None`` to pass through."""
if self._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"] = self._event_detail("ssrf_violation: ", raw_text)
return self._ssrf_soft_payload(raw_text), event
if self._is_workspace_violation(raw_text):
escalation = repeated_workspace_violation_error(
tool_call.name,
tool_call.arguments,
workspace_violation_counts,
)
event["detail"] = self._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"] = self._event_detail(
"workspace_violation_escalated: ",
raw_text,
)
return escalation, event
return soft_payload, event
return None
@classmethod
def _ssrf_soft_payload(cls, raw_text: str) -> str:
text = raw_text.strip() or "Error: request blocked by SSRF guard"
return f"{text}\n\n{cls._SSRF_BOUNDARY_NOTE}"
@staticmethod
def _event_detail(prefix: str, text: str, limit: int = 160) -> str:
return (prefix + text.replace("\n", " ").strip())[:limit]
async def _emit_checkpoint(
self,
spec: AgentRunSpec,
@@ -1696,28 +1531,3 @@ class AgentRunner:
if messages and messages[-1].get("role") == "assistant" and not messages[-1].get("tool_calls"):
return
messages.append(build_assistant_message(_PERSISTED_MODEL_ERROR_PLACEHOLDER))
def _partition_tool_batches(
self,
spec: AgentRunSpec,
tool_calls: list[ToolCallRequest],
) -> list[list[ToolCallRequest]]:
if not spec.concurrent_tools:
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(spec.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
+32 -1
View File
@@ -55,7 +55,8 @@ class SubagentStatus:
label: str
task_description: str
started_at: float # time.monotonic()
phase: str = "initializing" # initializing | awaiting_tools | tools_completed | final_response | done | error
# queued | initializing | awaiting_tools | tools_completed | final_response | done | error
phase: str = "initializing"
iteration: int = 0
tool_events: list[dict[str, str]] = field(default_factory=list)
usage: LLMUsage | None = None
@@ -147,6 +148,7 @@ class SubagentManager:
if max_concurrent_subagents is not None
else defaults.max_concurrent_subagents
)
self._run_slots = asyncio.Semaphore(self.max_concurrent_subagents)
self.runner = AgentRunner()
self._exec_session_manager = ExecSessionManager()
self._llm_wall_timeout_for_session = llm_wall_timeout_for_session
@@ -363,6 +365,35 @@ class SubagentManager:
workspace_scope: WorkspaceScope | None = None,
*,
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:
"""Execute the subagent task and announce the result."""
logger.info("Subagent [{}] starting task: {}", task_id, label)
+2
View File
@@ -11,6 +11,7 @@ if TYPE_CHECKING:
from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.exec_session import ExecSessionManager
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.runtime_events import RuntimeEventBus
from nanobot.config.schema import ProviderConfig, ToolsConfig
@@ -90,3 +91,4 @@ class ToolContext:
timezone: str = "UTC"
workspace_sandbox: WorkspaceSandboxStatus | None = None
runtime_events: RuntimeEventBus | None = None
runtime_control: RuntimeControl | None = None
+293
View File
@@ -0,0 +1,293 @@
"""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
+7 -11
View File
@@ -861,8 +861,10 @@ def _best_window(old_text: str, content: str) -> tuple[float, int, list[str], li
@tool_parameters(
tool_parameters_schema(
path=StringSchema("The file path to edit"),
old_text=StringSchema("The text to find and replace"),
new_text=StringSchema("The text to replace with"),
old_text=StringSchema("The text to find and replace; copy it from read_file."),
new_text=StringSchema(
"The replacement text; must differ from old_text for an existing file."
),
replace_all=BooleanSchema(description="Replace all occurrences (default false)"),
occurrence=IntegerSchema(
description="Optional 1-based occurrence to replace when old_text appears multiple times.",
@@ -899,15 +901,9 @@ class EditFileTool(_FsTool):
@property
def description(self) -> str:
return (
"Perform a small, exact replacement in one file by replacing "
"old_text with new_text. When replacing text in an existing file, "
"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."
"Perform a small, exact replacement in one file. "
"Prefer apply_patch for multi-file, structural, or generated edits. "
"occurrence, line_hint, and replace_all=true are mutually exclusive."
)
@staticmethod
+22 -16
View File
@@ -2,9 +2,11 @@
# pyright: reportIncompatibleMethodOverride=false
from collections.abc import Awaitable, Callable, Generator
from contextlib import contextmanager
from contextvars import ContextVar, Token
from pathlib import Path
from typing import Any, Awaitable, Callable, cast
from typing import Any, cast
from loguru import logger
@@ -16,6 +18,22 @@ from nanobot.bus.events import OutboundMessage
from nanobot.config.paths import get_workspace_path
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_schema(
@@ -68,7 +86,6 @@ class MessageTool(Tool):
self._fallback_chat_id = default_chat_id
self._fallback_message_id = default_message_id
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(
"message_suppress_delivery",
default=False,
@@ -87,10 +104,6 @@ class MessageTool(Tool):
"""Set the callback for sending messages."""
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]:
"""Acknowledge but don't deliver tool sends (heartbeat internal check)."""
return self._suppress_delivery_var.set(active)
@@ -99,14 +112,6 @@ class MessageTool(Tool):
"""Restore previous delivery-suppression state."""
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
def name(self) -> str:
return "message"
@@ -244,8 +249,9 @@ class MessageTool(Tool):
try:
await self._send_callback(msg)
if channel == default_channel and chat_id == default_chat_id:
self._sent_in_turn = True
sends = _CURRENT_MESSAGE_SENDS.get()
if sends is not None:
sends.add((channel, chat_id))
media_info = f" with {len(media)} attachments" if media else ""
button_info = (
f" with {sum(len(row) for row in button_rows)} button(s)"
+10 -2
View File
@@ -58,7 +58,6 @@ def _is_string_mapping(value: object) -> TypeGuard[Mapping[str, object]]:
class MyTool(Tool):
"""Check and set the agent loop's runtime configuration."""
_plugin_discoverable = False # Requires AgentLoop reference; registered manually
config_key = "my"
@classmethod
@@ -67,7 +66,16 @@ class MyTool(Tool):
@classmethod
def enabled(cls, ctx: ToolContext) -> bool:
return ctx.config.my.enable
return ctx.runtime_control is not None and 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({
# Core infrastructure
+16 -3
View File
@@ -7,7 +7,7 @@ from __future__ import annotations
import asyncio
import json
import time
from collections import deque
from collections import OrderedDict, deque
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Protocol
@@ -127,7 +127,7 @@ class SendSessionMessageTool(Tool):
self._max_messages_per_minute = max_messages_per_minute
self._schedule_later = schedule_later
self._clock = clock or time.monotonic
self._sent_at: dict[str, deque[float]] = {}
self._sent_at: OrderedDict[str, deque[float]] = OrderedDict()
self._pending_replies: dict[tuple[str, str], _PendingReply] = {}
self._expiry_tasks: set[asyncio.Task[None]] = set()
self._send_lock = asyncio.Lock()
@@ -240,8 +240,11 @@ class SendSessionMessageTool(Tool):
async with self._send_lock:
now = self._clock()
sent_at = self._sent_at.setdefault(source.session_key, deque())
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:
sent_at.popleft()
if len(sent_at) >= self._max_messages_per_minute:
@@ -259,6 +262,8 @@ class SendSessionMessageTool(Tool):
input_role="user",
))
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)
if timeout_seconds is not None:
self._cancel_pending_reply(wait_key)
@@ -271,6 +276,14 @@ class SendSessionMessageTool(Tool):
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
def _validate_reply_timeout(
expect_reply: bool,
+5 -8
View File
@@ -73,6 +73,11 @@ class SpawnTool(Tool):
"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(
self,
task: str,
@@ -82,14 +87,6 @@ class SpawnTool(Tool):
**kwargs: Any,
) -> str:
"""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()
if request_ctx is None or request_ctx.runtime is None:
return ToolResult.error("Error: spawn requires an active model runtime")
+24 -2
View File
@@ -182,6 +182,12 @@ 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)
# Forward to Nanobot via _on_message (non-blocking).
@@ -196,7 +202,7 @@ class NanobotDingTalkHandler(_CallbackHandlerBase):
)
)
self.channel._background_tasks.add(task)
task.add_done_callback(self.channel._background_tasks.discard)
task.add_done_callback(self.channel._on_background_task_done)
return AckMessage.STATUS_OK, "OK"
@@ -256,6 +262,17 @@ class DingTalkChannel(BaseChannel):
# Hold references to background tasks to prevent GC
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:
"""Start the DingTalk bot with Stream Mode."""
@@ -272,6 +289,7 @@ class DingTalkChannel(BaseChannel):
self.logger.error("client_id and client_secret not configured")
return
self._accepting_inbound_tasks = True
self._running = True
self._http = httpx.AsyncClient(
timeout=httpx.Timeout(10.0, connect=10.0, read=30.0, write=30.0, pool=10.0)
@@ -309,6 +327,7 @@ class DingTalkChannel(BaseChannel):
async def stop(self) -> None:
"""Stop the DingTalk bot."""
self._accepting_inbound_tasks = False
self._running = False
await self._close_stream_client()
start_task = self._start_task
@@ -326,8 +345,11 @@ class DingTalkChannel(BaseChannel):
await self._http.aclose()
self._http = None
# Cancel outstanding background tasks
for task in self._background_tasks:
background_tasks = tuple(self._background_tasks)
for task in background_tasks:
task.cancel()
if background_tasks:
await asyncio.gather(*background_tasks, return_exceptions=True)
self._background_tasks.clear()
async def _close_stream_client(self) -> None:
@@ -3,7 +3,7 @@ import json
import zipfile
from io import BytesIO
from types import SimpleNamespace
from unittest.mock import AsyncMock
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
@@ -402,6 +402,61 @@ async def test_handler_uses_voice_recognition_text_when_text_is_empty(monkeypatc
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
async def test_handler_processes_file_message(monkeypatch) -> None:
"""Test that file messages are handled and forwarded with downloaded path."""
@@ -451,6 +506,72 @@ async def test_handler_processes_file_message(monkeypatch) -> None:
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):
class _FakeRichTextChatbotMessage:
text = None
@@ -650,6 +771,41 @@ async def test_stop_cancels_stream_client_after_sdk_swallows_first_cancel(monkey
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
async def test_download_dingtalk_file(tmp_path, monkeypatch) -> None:
"""Test the two-step file download flow (get URL then download content)."""
+58 -43
View File
@@ -430,7 +430,13 @@ class EmailChannel(BaseChannel):
skipped_uids: set[str],
cycle_uids: set[str],
) -> 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"
client = self._open_imap_client(mailbox=mailbox, missing_mailbox_ok=True)
@@ -438,29 +444,30 @@ class EmailChannel(BaseChannel):
return messages
try:
status, data = client.search(None, *search_criteria)
if status != "OK" or not data:
status, data = client.uid("SEARCH", None, *search_criteria)
if status != "OK" or not data or not data[0]:
return messages
ids = data[0].split()
if limit > 0 and len(ids) > limit:
ids = ids[-limit:]
for imap_id in ids:
status, fetched = client.fetch(imap_id, "(BODY.PEEK[] UID)")
uids = [raw.decode("ascii", errors="ignore") for raw in data[0].split()]
if limit > 0 and len(uids) > limit:
uids = uids[-limit:]
features: _ServerFeatures | None = None
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:
continue
raw_bytes = self._extract_message_bytes(fetched)
if raw_bytes is None:
header_bytes = self._extract_message_bytes(fetched)
if header_bytes is None:
continue
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)
parsed = BytesParser(policy=policy.default).parsebytes(header_bytes)
sender = parseaddr(parsed.get("From", ""))[1].strip().lower()
if not sender:
continue
@@ -468,9 +475,8 @@ class EmailChannel(BaseChannel):
self.logger.info("From {} ignored: matches bot-owned address", sender)
self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen:
client.store(imap_id, "+FLAGS", "\\Seen")
if uid:
skipped_uids.add(uid)
features = self._mark_seen_uid(client, uid, features)
skipped_uids.add(uid)
continue
# --- Anti-spoofing: verify Authentication-Results ---
@@ -482,8 +488,7 @@ class EmailChannel(BaseChannel):
sender,
)
self._remember_processed_uid(uid, dedupe, cycle_uids)
if uid:
skipped_uids.add(uid)
skipped_uids.add(uid)
continue
if self.config.verify_dkim and not dkim_pass:
self.logger.warning(
@@ -492,18 +497,26 @@ class EmailChannel(BaseChannel):
sender,
)
self._remember_processed_uid(uid, dedupe, cycle_uids)
if uid:
skipped_uids.add(uid)
skipped_uids.add(uid)
continue
if not self.is_allowed(sender):
self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen:
client.store(imap_id, "+FLAGS", "\\Seen")
if uid:
skipped_uids.add(uid)
features = self._mark_seen_uid(client, uid, features)
skipped_uids.add(uid)
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", ""))
date_value = parsed.get("Date", "")
message_id = parsed.get("Message-ID", "").strip()
@@ -556,10 +569,19 @@ class EmailChannel(BaseChannel):
self._remember_processed_uid(uid, dedupe, cycle_uids)
if mark_seen:
client.store(imap_id, "+FLAGS", "\\Seen")
features = self._mark_seen_uid(client, uid, features)
finally:
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:
if self.config.imap_use_ssl:
client: Any = imaplib.IMAP4_SSL(self.config.imap_host, self.config.imap_port)
@@ -714,11 +736,14 @@ class EmailChannel(BaseChannel):
return data[0].split()[0]
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
# sequence-number lookup. If this fails once for the session, remember it
# and use the sequence STORE fallback directly for remaining UIDs.
if features.uid_store is not False:
status, _ = client.uid("STORE", uid, "+FLAGS", "(\\Deleted)")
status, _ = client.uid("STORE", uid, "+FLAGS", f"({flag})")
if status == "OK":
features.uid_store = True
return True
@@ -728,12 +753,12 @@ class EmailChannel(BaseChannel):
# unreliable: resolve the current sequence number from UID and use STORE.
imap_id = self._lookup_imap_id_by_uid(client, uid)
if not imap_id:
self.logger.warning("Post-action skipped: UID {} not found", uid)
self.logger.warning("Could not locate UID {} to set flag {}", uid, flag)
return False
status, _ = client.store(imap_id, "+FLAGS", "\\Deleted")
status, _ = client.store(imap_id, "+FLAGS", flag)
if status != "OK":
self.logger.warning("Post-action failed: could not mark UID {} as deleted", uid)
self.logger.warning("Failed to set flag {} on UID {}", flag, uid)
return False
return True
@@ -773,16 +798,6 @@ class EmailChannel(BaseChannel):
return bytes(fetched_item[1])
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
def _decode_header_value(value: str) -> str:
if not value:
@@ -53,30 +53,7 @@ def _make_raw_email(
def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
raw = _make_raw_email(subject="Invoice", body="Please pay")
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()
fake = _make_fake_imap(raw, uid=b"123")
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(), MessageBus())
@@ -86,38 +63,25 @@ def test_fetch_new_messages_parses_unseen_and_marks_seen(monkeypatch) -> None:
assert items[0]["sender"] == "alice@example.com"
assert items[0]["subject"] == "Invoice"
assert "Please pay" in items[0]["content"]
assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls
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()
# Same UID should be deduped in-process.
items_again, skipped_again = channel._fetch_new_messages()
assert items_again == []
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:
raw = _make_raw_email(subject="Invoice", body="Please pay")
class FakeIMAP:
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())
fake = _make_fake_imap(raw, uid=b"123")
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(post_action="delete"), MessageBus())
items, skipped_uids = channel._fetch_new_messages()
@@ -130,26 +94,10 @@ def test_fetch_new_messages_returns_accepted_and_skipped_uids(monkeypatch) -> No
def test_fetch_new_messages_rejected_returns_skipped_uid(monkeypatch) -> None:
raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test")
class FakeIMAP:
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())
monkeypatch.setattr(
"nanobot.channels.email.runtime.imaplib.IMAP4_SSL",
lambda _h, _p: _make_fake_imap(raw, uid=b"123"),
)
channel_skip = EmailChannel(
_make_config(from_address="bot@example.com", post_action="delete", post_action_ignore_skipped=True),
@@ -545,30 +493,7 @@ 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:
raw = _make_raw_email(from_addr="Nanobot <bot@example.com>", subject="Loop test")
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()
fake = _make_fake_imap(raw, uid=b"123")
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(from_address="bot@example.com"), MessageBus())
@@ -576,7 +501,7 @@ def test_fetch_new_messages_skips_self_sent_email_and_marks_seen(monkeypatch) ->
assert items == []
assert skipped_uids == {"123"}
assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls
# Same UID should still be deduped after being ignored.
items_again, skipped_again = channel._fetch_new_messages()
@@ -614,37 +539,14 @@ def test_fetch_new_messages_skips_self_sent_across_identity_sources(
imap_username matches, and must be case-insensitive."""
raw = _make_raw_email(from_addr=from_header, subject="Loop test")
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()
fake = _make_fake_imap(raw, uid=b"123")
monkeypatch.setattr("nanobot.channels.email.runtime.imaplib.IMAP4_SSL", lambda _h, _p: fake)
channel = EmailChannel(_make_config(**config_override), MessageBus())
items, _ = channel._fetch_new_messages()
assert items == []
assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
assert ("STORE", "123", "+FLAGS", "(\\Seen)") in fake.uid_calls
def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeypatch) -> None:
@@ -662,15 +564,16 @@ def test_fetch_new_messages_retries_once_when_imap_connection_goes_stale(monkeyp
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
self.search_calls += 1
if fail_once["pending"]:
fail_once["pending"] = False
raise imaplib.IMAP4.abort("socket error")
return "OK", [b"1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
def uid(self, command: str, *args):
if command == "SEARCH":
self.search_calls += 1
if fail_once["pending"]:
fail_once["pending"] = False
raise imaplib.IMAP4.abort("socket error")
return "OK", [b"123"]
if command == "FETCH":
return "OK", [(b"1 (UID 123 BODY[] {200})", raw), b")"]
return "OK", [b""]
def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags))
@@ -700,10 +603,7 @@ 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:
raw_first = _make_raw_email(subject="First", body="First body")
raw_second = _make_raw_email(subject="Second", body="Second body")
mailbox_state = {
b"1": {"uid": b"123", "raw": raw_first, "seen": False},
b"2": {"uid": b"124", "raw": raw_second, "seen": False},
}
mailbox_state = {"123": raw_first, "124": raw_second}
fail_once = {"pending": True}
class FlakyIMAP:
@@ -713,20 +613,18 @@ def test_fetch_new_messages_keeps_messages_collected_before_stale_retry(monkeypa
def select(self, _mailbox: str):
return "OK", [b"2"]
def search(self, *_args):
unseen_ids = [imap_id for imap_id, item in mailbox_state.items() if not item["seen"]]
return "OK", [b" ".join(unseen_ids)]
def fetch(self, imap_id: bytes, _parts: str):
if imap_id == b"2" and fail_once["pending"]:
fail_once["pending"] = False
raise imaplib.IMAP4.abort("socket error")
item = mailbox_state[imap_id]
header = b"%s (UID %s BODY[] {200})" % (imap_id, item["uid"])
return "OK", [(header, item["raw"]), b")"]
def store(self, imap_id: bytes, _op: str, _flags: str):
mailbox_state[imap_id]["seen"] = True
def uid(self, command: str, *args):
if command == "SEARCH":
keys = " ".join(sorted(mailbox_state.keys(), key=int))
return "OK", [keys.encode()]
if command == "FETCH":
uid = args[0]
if uid == "124" and fail_once["pending"]:
fail_once["pending"] = False
raise imaplib.IMAP4.abort("socket error")
raw = mailbox_state[uid]
header = f"{uid} (UID {uid} BODY[] {{200}})".encode()
return "OK", [(header, raw), b")"]
return "OK", [b""]
def logout(self):
@@ -1044,12 +942,13 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
self.search_args = _args
return "OK", [b"5"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"5 (UID 999 BODY[] {200})", raw), b")"]
def uid(self, command: str, *args):
if command == "SEARCH":
self.search_args = args
return "OK", [b"999"]
if command == "FETCH":
return "OK", [(b"5 (UID 999 BODY[] {200})", raw), b")"]
return "OK", [b""]
def store(self, imap_id: bytes, op: str, flags: str):
self.store_calls.append((imap_id, op, flags))
@@ -1070,7 +969,7 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
assert len(items) == 1
assert items[0]["subject"] == "Status"
# search(None, "SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
# uid("SEARCH", None, "SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
assert fake.search_args is not None
assert fake.search_args[1:] == ("SINCE", "06-Feb-2026", "BEFORE", "07-Feb-2026")
assert fake.store_calls == []
@@ -1080,11 +979,12 @@ def test_fetch_messages_between_dates_uses_imap_since_before_without_mark_seen(m
# Security: Anti-spoofing tests for Authentication-Results verification
# ---------------------------------------------------------------------------
def _make_fake_imap(raw: bytes):
def _make_fake_imap(raw: bytes, uid: bytes = b"500"):
"""Return a FakeIMAP class pre-loaded with the given raw email."""
class FakeIMAP:
def __init__(self) -> None:
self.store_calls: list[tuple[bytes, str, str]] = []
self.uid_calls: list[tuple] = []
def login(self, _user: str, _pw: str):
return "OK", [b"logged in"]
@@ -1092,11 +992,16 @@ def _make_fake_imap(raw: bytes):
def select(self, _mailbox: str):
return "OK", [b"1"]
def search(self, *_args):
return "OK", [b"1"]
def capability(self):
return "OK", [b"IMAP4rev1"]
def fetch(self, _imap_id: bytes, _parts: str):
return "OK", [(b"1 (UID 500 BODY[] {200})", raw), b")"]
def uid(self, command: str, *args):
self.uid_calls.append((command, *args))
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):
self.store_calls.append((imap_id, op, flags))
@@ -1292,7 +1197,10 @@ def test_fetch_new_messages_ignores_unauthorized_sender_before_attachments(monke
assert channel._fetch_new_messages() == ([], {"500"})
assert called["attachments"] is False
assert fake.store_calls == [(b"1", "+FLAGS", "\\Seen")]
assert [call for call in fake.uid_calls if call[0] == "FETCH"] == [
("FETCH", "500", "(BODY.PEEK[HEADER])")
]
assert ("STORE", "500", "+FLAGS", "(\\Seen)") in fake.uid_calls
def test_extract_attachments_saves_pdf(tmp_path, monkeypatch) -> None:
+71 -19
View File
@@ -897,6 +897,68 @@ class TelegramChannel(BaseChannel):
self.logger.debug("sendRichMessage failed: {}", exc)
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:
"""Send a message through Telegram."""
app = await self._wait_for_app()
@@ -1136,26 +1198,16 @@ class TelegramChannel(BaseChannel):
thread_kwargs["message_thread_id"] = message_thread_id
raw_text = buf.text
# Try sendRichMessage for final output (Bot API 10.1).
# Skip when a streaming preview already exists to avoid the
# delete-and-resend pattern that causes flickering and drops
# line breaks (issue #4470).
if not buf.message_id and self.config.rich_messages and not getattr(self, "_rich_send_disabled", False):
reply_params = None
if reply_to_message_id := meta.get("message_id"):
reply_params = {"message_id": int(reply_to_message_id), "allow_sending_without_reply": True}
rich_ok = await self._try_send_rich(
int_chat_id, raw_text, reply_params, thread_kwargs, None,
)
# Try upgrading the streaming preview to rich in place (Bot API 10.1:
# editMessageText gained a rich_message parameter). Editing in place
# keeps the message identity, so there is no delete-and-resend and
# none of the flickering / dropped line breaks from issue #4470.
# The previous branch here was unreachable: it was guarded by
# ``not buf.message_id`` after an early return had already ensured
# ``buf.message_id`` is set (issue #5516).
if self.config.rich_messages and not getattr(self, "_rich_send_disabled", False):
rich_ok = await self._try_edit_rich(int_chat_id, buf.message_id, raw_text)
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)
return
@@ -2735,3 +2735,130 @@ def test_markdown_to_html_code_block_same_line_no_newline() -> None:
stripped = _strip_md_block(text)
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
+12 -3
View File
@@ -87,7 +87,12 @@ app = typer.Typer(
name="nanobot",
context_settings={"help_option_names": ["-h", "--help"]},
help=f"{__logo__} nanobot - Personal AI Assistant",
no_args_is_help=True,
epilog=(
"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()
@@ -98,7 +103,7 @@ def version_callback(value: bool):
raise typer.Exit()
@app.callback()
@app.callback(invoke_without_command=True)
def main(
ctx: typer.Context,
version: bool = typer.Option(
@@ -110,7 +115,11 @@ def main(
# imports this Typer app directly instead of ``nanobot.cli.entry``. Keep the
# role identity correct until that launcher is regenerated.
command = ctx.invoked_subcommand
set_cli_process_identity([command] if command else sys.argv[1:])
set_cli_process_identity([command] if command else ["agent"])
if command is None:
from nanobot.cli.entry import _run_agent
_run_agent([], prog_name="nanobot")
# ============================================================================
+48 -10
View File
@@ -8,6 +8,28 @@ from contextlib import suppress
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:
"""Return whether ``agent`` can start without the classic agent stack."""
@@ -34,19 +56,35 @@ def _configure_windows_console() -> None:
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:
"""Dispatch native TUI startup without importing the complete CLI graph."""
set_cli_process_identity(sys.argv[1:])
raw_args = 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()
if _native_tui_candidate(sys.argv[1:]):
import typer
from nanobot.cli.agent import agent
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")
root_agent_alias = agent_args is not None and raw_args[:1] != ["agent"]
if agent_args is not None and (
root_agent_alias or _native_tui_candidate(dispatch_args)
):
prog_name = "nanobot" if root_agent_alias else "nanobot agent"
_run_agent(agent_args, prog_name=prog_name)
return
from nanobot.cli.commands import app
+5 -3
View File
@@ -23,6 +23,7 @@ from nanobot.cli.webui_support import (
_gateway_health_bind_note,
_gateway_health_url,
_host_for_local_browser,
_launch_browser,
_prepare_webui_bundle_for_gateway,
_print_foreground_port_conflict,
_tcp_endpoint_reachable,
@@ -864,7 +865,6 @@ def _run_gateway(
"""Wait for the gateway to bind, then point the user's browser at the webui."""
if not open_browser_url:
return
import webbrowser
from urllib.parse import urlparse
# Channels start asynchronously. When the caller supplies a backend
@@ -896,8 +896,10 @@ def _run_gateway(
await asyncio.sleep(0.1)
display_url = _webui_display_url(open_browser_url)
try:
webbrowser.open(open_browser_url)
console.print(f"[green]✓[/green] Opened browser at {display_url}")
if _launch_browser(open_browser_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:
console.print(f"[yellow]Could not open browser ({e}); visit {display_url}[/yellow]")
+5 -2
View File
@@ -29,7 +29,7 @@ _PROVIDER_DISPLAY: dict[str, str] = {
_OAUTH_PROVIDER_DEFAULT_MODELS: dict[str, str] = {
"openai_codex": "openai-codex/gpt-5.6-sol",
"xai_grok": "xai-grok/grok-4.5",
"xai_grok": "xai-grok/grok-4.6",
"github_copilot": "github-copilot/gpt-5.4-mini",
}
@@ -134,7 +134,10 @@ def _set_oauth_provider_as_main(
config.agents.defaults.model_preset = None
config.agents.defaults.provider = provider_name
config.agents.defaults.model = selected_model
if provider_name == "xai_grok" and selected_model == "xai-grok/grok-4.5":
if provider_name == "xai_grok" and selected_model in {
"xai-grok/grok-4.5",
"xai-grok/grok-4.6",
}:
config.agents.defaults.context_window_tokens = 500_000
save_config(config, resolved_config_path)
+5
View File
@@ -21,6 +21,7 @@ from nanobot.cli.process_identity import named_executable
from nanobot.cli.runtime_config import _model_display
from nanobot.cli.webui_support import (
_gateway_health_ready,
_gateway_health_url,
_gateway_instance_command,
_host_for_local_browser,
_webui_endpoint_reachable,
@@ -96,6 +97,10 @@ def launch_tui(
env.update(
{
"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_MODEL": _model_display(config)[0],
"NANOBOT_TUI_MODEL_PRESET": config.agents.defaults.model_preset or "default",
+104 -8
View File
@@ -1,10 +1,14 @@
"""Shared WebUI setup, URL, health, and browser helpers."""
import os
import subprocess
import sys
import time
import webbrowser
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, BinaryIO
import typer
from pydantic import ValidationError
@@ -40,6 +44,7 @@ __all__ = [
"_gateway_instance_command",
"_host_for_local_browser",
"_load_webui_setup_config",
"_launch_browser",
"_open_webui_browser",
"_prepare_webui_bundle_for_gateway",
"_print_foreground_port_conflict",
@@ -60,6 +65,20 @@ __all__ = [
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:
"""Confirm a WebUI first-run mutation or fail clearly in non-interactive shells."""
if yes:
@@ -419,14 +438,14 @@ def _print_foreground_port_conflict(
def _open_webui_browser(url: str, *, wait: bool = True) -> None:
"""Open the WebUI in the user's default browser, with a copyable fallback."""
import webbrowser
if wait:
_wait_for_webui(url)
display_url = _webui_display_url(url)
try:
webbrowser.open(url)
console.print(f"[green]✓[/green] Opened WebUI: [cyan]{display_url}[/cyan]")
if _launch_browser(url):
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:
console.print(f"[yellow]Could not open browser ({exc}); visit {display_url}[/yellow]")
@@ -440,27 +459,104 @@ def _print_webui_foreground_lifecycle(*, attached: bool) -> None:
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]Press Ctrl+C to detach; the gateway stops only when the last local client exits.[/dim]"
"[dim]Following live gateway logs. Press Ctrl+C to detach; the gateway stops "
"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(
runtime: "GatewayRuntime",
*,
poll_hook: Callable[[], None] | None = None,
sleep: Callable[[float], None] = time.sleep,
) -> None:
"""Keep a WebUI launcher attached without taking ownership of the gateway."""
"""Keep the launcher attached and mirror this gateway's new log output."""
status = runtime.status()
log_path = status.log_path
cursor = _start_gateway_log_cursor(log_path)
_print_webui_foreground_lifecycle(attached=True)
try:
while runtime.status().running:
while 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:
poll_hook()
sleep(0.5)
status = runtime.status()
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]")
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]")
+1 -1
View File
@@ -311,7 +311,7 @@ async def cmd_new(ctx: CommandContext) -> OutboundMessage:
snapshot = list(session.messages)
archive_snapshot = None
runtime = None
if session.last_consolidated < len(snapshot):
if session.last_archived < len(snapshot):
runtime = ctx.runtime or loop.runtime_for_session(session)
archive_snapshot = replace(
session,
+1 -8
View File
@@ -128,7 +128,7 @@ class AgentDefaults(Base):
temperature: float = 0.1
fallback_models: list[FallbackCandidate] = Field(default_factory=list)
max_tool_iterations: int = 200
max_concurrent_subagents: int = Field(default=1, ge=1)
max_concurrent_subagents: int = Field(default=4, ge=1)
max_tool_result_chars: int = 16_000
provider_retry_mode: Literal["standard", "persistent"] = "standard"
tool_hint_max_length: int = Field(
@@ -155,13 +155,6 @@ class AgentDefaults(Base):
default=60,
ge=0,
) # 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)
@model_validator(mode="before")
+21 -3
View File
@@ -25,6 +25,7 @@ from nanobot.cron.types import (
CronSchedule,
CronStore,
)
from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META
from nanobot.utils.run_records import (
write_run_record as write_automation_run_record,
)
@@ -115,8 +116,21 @@ def _disable_malformed_legacy_job(job: CronJob) -> None:
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:
"""Migrate legacy user cron payloads into session-bound payloads.
"""Make routing metadata persistable and migrate legacy user cron payloads.
Pre-bound user cron jobs stored their delivery target in ``channel``/``to``.
Normal user-created legacy jobs always have those fields; if they are
@@ -124,8 +138,12 @@ def _normalize_agent_turn_job(job: CronJob) -> bool:
a runtime legacy execution path.
"""
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):
return False
return changed
if not payload.channel or not payload.to:
_disable_malformed_legacy_job(job)
@@ -135,7 +153,7 @@ def _normalize_agent_turn_job(job: CronJob) -> bool:
payload.origin_channel = payload.origin_channel or payload.channel
payload.origin_chat_id = payload.origin_chat_id or payload.to
if not payload.origin_metadata:
payload.origin_metadata = dict(payload.channel_meta or {})
payload.origin_metadata = _persistable_origin_metadata(payload.channel_meta or {})
payload.deliver = False
payload.channel = None
+21 -2
View File
@@ -603,8 +603,6 @@ _SYNTHETIC_USER_CONTENT = "(conversation continued)"
class LLMProvider(ABC):
"""Base class for LLM providers."""
supports_progress_deltas = False
_CHAT_RETRY_DELAYS = (1, 2, 4)
_PERSISTENT_MAX_DELAY = 60
_PERSISTENT_IDENTICAL_ERROR_LIMIT = 10
@@ -1031,6 +1029,20 @@ class LLMProvider(ABC):
# Unknown 429 defaults to WAIT+retry.
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
def _enforce_role_alternation(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Merge consecutive same-role messages and drop trailing assistant messages.
@@ -1065,6 +1077,13 @@ class LLMProvider(ABC):
curr_content = msg.get("content") or ""
if isinstance(prev_content, str) and isinstance(curr_content, str):
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:
merged[-1] = dict(msg)
else:
+42 -1
View File
@@ -11,6 +11,7 @@ from nanobot.providers.base import (
ProviderCallContext,
ProviderConversationState,
)
from nanobot.utils.helpers import estimate_prompt_tokens_chain
_PROVIDER_STATE_OUTPUT_META = "provider_state_output"
_PROVIDER_STATE_BOUNDARY_META = "provider_state_boundary"
@@ -69,6 +70,37 @@ class ProviderConversationStateController:
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(
self,
messages: list[dict[str, Any]],
@@ -76,11 +108,20 @@ class ProviderConversationStateController:
context_window_tokens: int | None,
model_messages: list[dict[str, Any]] | None = None,
supplemental_messages: list[dict[str, Any]] | None = None,
resume_state: bool = True,
) -> ProviderCallContext | None:
"""Build typed context for the next request and remember its durable delta."""
"""Build 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(
context_window_tokens=context_window_tokens,
)
if not resume_state:
self._state = None
self._request_messages = []
return independent_context
if self._state is None:
self._request_messages = []
return independent_context
-4
View File
@@ -157,10 +157,6 @@ class FallbackProvider(LLMProvider):
super().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(
self,
state: ProviderConversationState,
+183 -4
View File
@@ -5,6 +5,7 @@
from __future__ import annotations
import asyncio
import hashlib
import os
import time
import webbrowser
@@ -17,7 +18,12 @@ from oauth_cli_kit.models import OAuthToken
from oauth_cli_kit.storage import FileTokenStorage
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.registry import ProviderModelSpec, find_by_name
DEFAULT_GITHUB_DEVICE_CODE_URL = "https://github.com/login/device/code"
DEFAULT_GITHUB_ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token"
@@ -96,7 +102,9 @@ def login_github_copilot(
device_code = str(payload["device_code"])
user_code = str(payload["user_code"])
verify_url = str(payload.get("verification_uri") or payload.get("verification_uri_complete") or "")
verify_url = str(
payload.get("verification_uri") or payload.get("verification_uri_complete") or ""
)
verify_complete = str(payload.get("verification_uri_complete") or verify_url)
interval = max(1, int(payload.get("interval") or 5))
expires_in = int(payload.get("expires_in") or 900)
@@ -180,8 +188,6 @@ class GitHubCopilotProvider(OpenAICompatProvider):
*,
provider_name: str = "github_copilot",
):
from nanobot.providers.registry import find_by_name
self._copilot_access_token: str | None = None
self._copilot_expires_at: float = 0.0
self._copilot_token_lock: asyncio.Lock = asyncio.Lock()
@@ -217,7 +223,9 @@ class GitHubCopilotProvider(OpenAICompatProvider):
)
timeout = httpx.Timeout(20.0, connect=20.0)
async with httpx.AsyncClient(timeout=timeout, follow_redirects=True, trust_env=True) as client:
async with httpx.AsyncClient(
timeout=timeout, follow_redirects=True, trust_env=True
) as client:
response = await client.get(
_resolve("NANOBOT_COPILOT_TOKEN_URL", DEFAULT_COPILOT_TOKEN_URL),
headers=_copilot_headers(github_token.access),
@@ -296,3 +304,174 @@ class GitHubCopilotProvider(OpenAICompatProvider):
on_tool_call_delta=on_tool_call_delta,
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,
)
+224
View File
@@ -0,0 +1,224 @@
"""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()
+173 -30
View File
@@ -14,7 +14,10 @@ from typing import Any, cast
import httpx
from loguru import logger
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 (
LLMProvider,
LLMResponse,
@@ -22,6 +25,10 @@ from nanobot.providers.base import (
ProviderConversationState,
resolve_stream_idle_timeout_s,
)
from nanobot.providers.oauth_model_catalog import (
OAuthModelCatalog,
OAuthModelCatalogSnapshot,
)
from nanobot.providers.openai_responses import (
ResponsesStreamCapture,
build_responses_state,
@@ -35,8 +42,11 @@ from nanobot.providers.openai_responses import (
responses_state_items,
responses_state_matches,
)
from nanobot.providers.registry import ProviderModelSpec, find_by_name
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"
_COMPACTION_RETAINED_CHAR_BUDGET = 256_000
@@ -44,8 +54,6 @@ _COMPACTION_RETAINED_CHAR_BUDGET = 256_000
class OpenAICodexProvider(LLMProvider):
"""Use Codex OAuth to call the Responses API."""
supports_progress_deltas = True
def __init__(
self,
default_model: str = "openai-codex/gpt-5.6-sol",
@@ -89,9 +97,7 @@ class OpenAICodexProvider(LLMProvider):
model = model or self.default_model
sanitized_messages = self._sanitize_empty_content(messages)
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:
sanitized_state = sanitized_state.with_pending_messages(
@@ -170,11 +176,7 @@ class OpenAICodexProvider(LLMProvider):
)
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,
)
if (
@@ -238,8 +240,12 @@ class OpenAICodexProvider(LLMProvider):
return response
async def chat(
self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
model: str | None = None, max_tokens: int = 4096, temperature: float = 0.7,
self,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
model: str | None = None,
max_tokens: int = 4096,
temperature: float = 0.7,
reasoning_effort: str | None = None,
tool_choice: str | dict[str, Any] | None = None,
provider_context: ProviderCallContext | None = None,
@@ -266,8 +272,12 @@ class OpenAICodexProvider(LLMProvider):
)
async def chat_stream(
self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None,
model: str | None = None, max_tokens: int = 4096, temperature: float = 0.7,
self,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
model: str | None = None,
max_tokens: int = 4096,
temperature: float = 0.7,
reasoning_effort: str | None = None,
tool_choice: str | dict[str, Any] | None = None,
on_content_delta: Callable[[str], Awaitable[None]] | None = None,
@@ -346,11 +356,7 @@ def _without_response_item_ids(
sanitized_input.append(raw_item)
continue
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["input"] = sanitized_input
@@ -446,15 +452,12 @@ async def _request_codex(
raw = text.decode("utf-8", "ignore")
retry_after = LLMProvider._extract_retry_after_from_headers(response.headers)
error_type, error_code = LLMProvider._extract_error_type_code(raw)
compaction_unsupported = (
response.status_code in {400, 404, 422}
and any(
marker in raw.lower()
for marker in (
"context_management",
"compact_threshold",
"compaction_trigger",
)
compaction_unsupported = response.status_code in {400, 404, 422} and any(
marker in raw.lower()
for marker in (
"context_management",
"compact_threshold",
"compaction_trigger",
)
)
raise _CodexHTTPError(
@@ -463,7 +466,9 @@ async def _request_codex(
retry_after=retry_after,
error_type=error_type,
error_code=error_code,
should_retry=_should_retry_status(response.status_code, error_type, error_code, raw),
should_retry=_should_retry_status(
response.status_code, error_type, error_code, raw
),
compaction_unsupported=compaction_unsupported,
)
capture = ResponsesStreamCapture()
@@ -536,7 +541,9 @@ def _codex_error_response(exc: Exception) -> LLMResponse:
default_detail = "HTTP request failed"
if status_code is not None and should_retry is None:
retry_content = None if int(status_code) == 429 and isinstance(exc, _CodexHTTPError) else detail
retry_content = (
None if int(status_code) == 429 and isinstance(exc, _CodexHTTPError) else detail
)
should_retry = _should_retry_status(
int(status_code),
getattr(exc, "error_type", None),
@@ -594,3 +601,139 @@ def _should_retry_status(
)
)
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,
)
+41 -8
View File
@@ -20,12 +20,15 @@ from pydantic.alias_generators import to_snake
@dataclass(frozen=True)
class ProviderModelSpec:
"""A curated model exposed by providers without a model-list endpoint."""
"""Curated model metadata used for fixed catalogs or online fallback."""
id: str
label: str = ""
description: str = ""
owned_by: str = ""
context_window: int | None = None
reasoning_efforts: tuple[str, ...] = ()
supports_backend_search: bool = False
@dataclass(frozen=True)
@@ -42,7 +45,7 @@ class ProviderSpec:
keywords: tuple[str, ...] # model-name keywords for matching (lowercase)
env_key: str # env var for API key, e.g. "DASHSCOPE_API_KEY"
display_name: str = "" # shown in `nanobot status`
model_catalog: str = "auto" # WebUI model-list source
model_catalog: str = "auto" # WebUI model-list source, including builtin/hybrid
builtin_models: tuple[ProviderModelSpec, ...] = ()
settings_alias_for: str = "" # compatibility alias grouped under this provider in Settings
@@ -407,45 +410,56 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
keywords=("openai-codex",),
env_key="",
display_name="OpenAI Codex",
model_catalog="builtin",
model_catalog="hybrid",
builtin_models=(
ProviderModelSpec(
id="openai-codex/gpt-5.6-sol",
label="GPT-5.6-Sol",
description="Latest frontier agentic coding model.",
context_window=372000,
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max", "ultra"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.6-terra",
label="GPT-5.6-Terra",
description="Balanced agentic coding model for everyday work.",
context_window=372000,
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max", "ultra"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.6-luna",
label="GPT-5.6-Luna",
description="Fast and affordable agentic coding model.",
context_window=372000,
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh", "max"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.5",
label="GPT-5.5",
description="Frontier model for complex coding, research, and real-world work.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.4",
label="GPT-5.4",
description="Strong model for everyday coding.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.4-mini",
label="GPT-5.4-Mini",
description="Small, fast, and cost-efficient model for simpler coding tasks.",
context_window=272_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
ProviderModelSpec(
id="openai-codex/gpt-5.3-codex-spark",
label="GPT-5.3-Codex-Spark",
description="Ultra-fast coding model.",
context_window=128_000,
reasoning_efforts=("low", "medium", "high", "xhigh"),
),
),
backend="openai_codex",
@@ -459,13 +473,19 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
keywords=("xai-grok", "xai_grok"),
env_key="",
display_name="xAI Grok",
model_catalog="builtin",
model_catalog="hybrid",
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(
id="xai-grok/grok-4.5",
label="Grok 4.5",
description="Grok via xAI subscription; X Search is enabled when supported.",
context_window=500000,
context_window=500_000,
),
),
backend="xai_grok",
@@ -478,6 +498,19 @@ PROVIDERS: tuple[ProviderSpec, ...] = (
keywords=("github_copilot", "copilot"),
env_key="",
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",
default_api_base="https://api.githubcopilot.com",
strip_model_prefix=True,
+392 -188
View File
@@ -4,9 +4,9 @@ from __future__ import annotations
import asyncio
import base64
import hashlib
import json
import re
import time
import uuid
from collections.abc import Awaitable, Callable
from typing import Any, cast
@@ -22,21 +22,24 @@ from nanobot.providers.base import (
ToolCallRequest,
resolve_stream_idle_timeout_s,
)
from nanobot.providers.oauth_model_catalog import OAuthModelCatalog, OAuthModelCatalogSnapshot
from nanobot.providers.openai_responses import (
consume_sse_with_reasoning,
convert_messages,
convert_tools,
)
from nanobot.providers.registry import ProviderModelSpec, find_by_name
from nanobot.providers.xai_oauth import (
XAI_CLIENT_VERSION,
XAIToken,
get_xai_oauth_login_status,
get_xai_oauth_storage_path,
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_MODELS_URL = "https://cli-chat-proxy.grok.com/v1/models"
DEFAULT_XAI_GROK_MODEL = "xai-grok/grok-4.5"
_MODEL_CAPABILITIES_TTL_S = 5 * 60
_HOSTED_SEARCH_MAX_TURNS = 5
_MAX_ERROR_BODY_CHARS = 1000
_SENSITIVE_ERROR_KEYS = {
"accesstoken",
@@ -63,7 +66,9 @@ def _is_named_x_search_tool(value: object) -> bool:
class XAIGrokProvider(LLMProvider):
"""Call xAI's subscription proxy and expose supported hosted tools."""
supports_progress_deltas = True
# An incomplete hosted-tool stream can already have emitted answer text. Let the
# provider close that stream segment before its one bounded recovery attempt.
supports_stream_recover_callback = True
def __init__(
self,
@@ -77,37 +82,19 @@ class XAIGrokProvider(LLMProvider):
self.default_model = default_model
self.proxy = proxy or None
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, token: XAIToken, model: str) -> bool:
now = time.monotonic()
capabilities = self._model_capabilities
if (
capabilities is None
or now - self._model_capabilities_fetched_at >= _MODEL_CAPABILITIES_TTL_S
):
try:
capabilities = await _fetch_xai_model_capabilities(
DEFAULT_XAI_GROK_MODELS_URL,
_build_model_headers(token),
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 _supports_backend_search(self, model: str) -> bool:
catalog = await asyncio.to_thread(
get_xai_grok_model_catalog,
self.proxy,
)
if catalog.message:
logger.warning(
"xAI model catalog unavailable; hosted X Search disabled unless cached: {}",
catalog.message,
)
info = catalog.find(model)
return bool(info and info.supports_backend_search)
async def _call_xai(
self,
@@ -121,6 +108,7 @@ class XAIGrokProvider(LLMProvider):
on_content_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_stream_recover: Callable[[], Awaitable[None]] | None = None,
) -> LLMResponse:
wire_model = _strip_model_prefix(model or self.default_model)
system_prompt, input_items = convert_messages(messages)
@@ -130,17 +118,13 @@ class XAIGrokProvider(LLMProvider):
token = await asyncio.to_thread(get_xai_oauth_token, proxy=self.proxy)
configured_tools = self._extra_body.get("tools")
tools_are_explicit = "tools" in self._extra_body
configured_hosted_search = (
isinstance(configured_tools, list)
and any(
_is_hosted_x_search_tool(tool)
for tool in cast(list[object], configured_tools)
)
configured_hosted_search = isinstance(configured_tools, list) and any(
_is_hosted_x_search_tool(tool) for tool in cast(list[object], configured_tools)
)
supports_backend_search = False
if not tools_are_explicit:
stage = "model_capabilities"
supports_backend_search = await self._supports_backend_search(token, wire_model)
supports_backend_search = await self._supports_backend_search(wire_model)
converted_tools = convert_tools(tools or [])
if isinstance(configured_tools, list):
converted_tools.extend(cast(list[dict[str, Any]], configured_tools))
@@ -151,6 +135,8 @@ class XAIGrokProvider(LLMProvider):
if supports_backend_search:
converted_tools.append({"type": "x_search"})
hosted_search_enabled = supports_backend_search or configured_hosted_search
body: dict[str, Any] = {
"model": wire_model,
"store": False,
@@ -166,51 +152,65 @@ class XAIGrokProvider(LLMProvider):
"temperature": temperature,
"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:
body.update({
key: value
for key, value in self._extra_body.items()
if key != "tools"
})
body.update(
{key: value for key, value in self._extra_body.items() if key != "tools"}
)
if tools_are_explicit and not isinstance(configured_tools, list):
body["tools"] = configured_tools
headers = _build_headers(token.access, wire_model)
stage = "xai_request"
try:
result = await _request_xai(
DEFAULT_XAI_GROK_URL,
headers,
body,
proxy=self.proxy,
on_content_delta=on_content_delta,
on_thinking_delta=on_thinking_delta,
on_tool_call_delta=on_tool_call_delta,
)
except _XAIHTTPError as exc:
if exc.status_code != 401:
raise
stage = "oauth_refresh"
token = await asyncio.to_thread(
get_xai_oauth_token,
proxy=self.proxy,
force_refresh=True,
)
self._model_capabilities = None
self._model_capabilities_fetched_at = 0.0
headers = _build_headers(token.access, wire_model)
stage = "xai_request_retry"
result = await _request_xai(
DEFAULT_XAI_GROK_URL,
headers,
body,
proxy=self.proxy,
on_content_delta=on_content_delta,
on_thinking_delta=on_thinking_delta,
on_tool_call_delta=on_tool_call_delta,
)
auth_retried = False
hosted_tool_retried = False
retry_usage: LLMUsage | None = None
while True:
try:
result = await _request_xai(
DEFAULT_XAI_GROK_URL,
headers,
body,
proxy=self.proxy,
on_content_delta=on_content_delta,
on_thinking_delta=on_thinking_delta,
on_tool_call_delta=on_tool_call_delta,
)
break
except _XAIHTTPError as exc:
if exc.status_code != 401 or auth_retried:
raise
auth_retried = True
stage = "oauth_refresh"
token = await asyncio.to_thread(
get_xai_oauth_token,
proxy=self.proxy,
force_refresh=True,
)
headers = _build_headers(token.access, wire_model)
stage = "xai_request_after_oauth_refresh"
except _XAIIncompleteHostedToolError as exc:
retry_usage = _combine_usage(retry_usage, exc.usage)
cannot_recover_stream = exc.stream_output_emitted and on_stream_recover is None
if hosted_tool_retried or cannot_recover_stream:
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
usage = _combine_usage(retry_usage, usage)
return LLMResponse(
content=content,
tool_calls=tool_calls,
@@ -259,6 +259,7 @@ class XAIGrokProvider(LLMProvider):
on_content_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_stream_recover: Callable[[], Awaitable[None]] | None = None,
) -> LLMResponse:
return await self._call_xai(
messages,
@@ -271,6 +272,7 @@ class XAIGrokProvider(LLMProvider):
on_content_delta,
on_thinking_delta,
on_tool_call_delta,
on_stream_recover,
)
def get_default_model(self) -> str:
@@ -290,6 +292,14 @@ def _build_reasoning_options(reasoning_effort: str | None) -> dict[str, str]:
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]:
conversation_id = str(uuid.uuid4())
return {
@@ -310,44 +320,6 @@ 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):
def __init__(
self,
@@ -369,65 +341,25 @@ class _XAIHTTPError(RuntimeError):
self.response_body = response_body
async def _fetch_xai_model_capabilities(
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)
class _XAIIncompleteHostedToolError(RuntimeError):
"""A nominally successful xAI stream ended before a hosted tool did."""
should_retry = False # _call_xai already performs the one safe recovery attempt.
def _parse_xai_model_capabilities(payload: Any) -> dict[str, bool]:
if isinstance(payload, dict):
payload = cast(dict[str, Any], payload)
rows: object = payload.get("data")
if not isinstance(rows, list):
rows = payload.get("models")
else:
rows = payload
if not isinstance(rows, list):
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"),
def __init__(
self,
active_tools: list[dict[str, Any]],
*,
usage: LLMUsage | None,
stream_output_emitted: bool = False,
) -> None:
names = [str(event.get("name") or "hosted_tool") for event in active_tools]
super().__init__(
"xAI ended the response before its hosted tool completed: " + ", ".join(names)
)
for identifier in identifiers:
if isinstance(identifier, str) and identifier.strip():
capabilities[_strip_model_prefix(identifier.strip())] = supports_backend_search
return capabilities
self.tool_names = tuple(names)
self.usage = usage
self.stream_output_emitted = stream_output_emitted
async def _request_xai(
@@ -440,10 +372,39 @@ async def _request_xai(
on_thinking_delta: Callable[[str], 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]:
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:
hosted_event = _xai_hosted_tool_event(event)
if hosted_event is not None and on_tool_call_delta is not None:
await on_tool_call_delta(hosted_event)
if hosted_event is not None:
await _track_and_forward_tool_event(hosted_event)
client_kwargs: dict[str, Any] = {"timeout": resolve_stream_idle_timeout_s()}
if proxy:
@@ -454,13 +415,34 @@ async def _request_xai(
content = await response.aread()
raw = content.decode("utf-8", "ignore")
raise _build_xai_http_error(response.status_code, response.headers, raw)
return await consume_sse_with_reasoning(
result = await consume_sse_with_reasoning(
response,
on_content_delta=on_content_delta,
on_tool_call_delta=on_tool_call_delta,
on_reasoning_delta=on_thinking_delta,
on_response_event=_on_response_event if on_tool_call_delta else None,
on_content_delta=(_forward_content_delta if on_content_delta is not None else None),
# Always observe tool events so protocol validation also works for
# non-streaming callers that did not request UI progress callbacks.
on_tool_call_delta=_track_and_forward_tool_event,
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:
@@ -474,19 +456,33 @@ def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
"phase": "start",
"call_id": str(call_id),
"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,
}
if event_type != "response.output_item.done":
if event_type not in {"response.output_item.added", "response.output_item.done"}:
return None
item = event.get("item")
if not isinstance(item, dict):
return None
item = cast(dict[str, Any], item)
if item.get("type") != "custom_tool_call":
item_type = item.get("type")
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
tool_name = item.get("name")
if not isinstance(tool_name, str) or not tool_name.startswith("x_"):
@@ -499,9 +495,7 @@ def _xai_hosted_tool_event(event: dict[str, Any]) -> dict[str, Any] | None:
"phase": "end",
"call_id": str(call_id),
"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
# in WebUI activity messages. The model answer already carries citations.
"result": {"name": tool_name},
@@ -610,6 +604,8 @@ def _xai_error_response(exc: Exception) -> LLMResponse:
should_retry = True if should_retry is None else should_retry
elif isinstance(exc, _XAIHTTPError):
error_kind = "http"
elif isinstance(exc, _XAIIncompleteHostedToolError):
error_kind = "provider"
if status_code is not None and should_retry is None:
should_retry = _should_retry_status(
int(status_code),
@@ -619,9 +615,11 @@ def _xai_error_response(exc: Exception) -> LLMResponse:
)
message = str(exc).strip() or "unexpected error"
retry_after = getattr(exc, "retry_after", None)
usage = getattr(exc, "usage", None)
return LLMResponse(
content=f"Error calling xAI ({type(exc).__name__}): {message}",
finish_reason="error",
usage=usage if isinstance(usage, LLMUsage) else None,
retry_after=retry_after,
error_status_code=int(status_code) if status_code is not None else None,
error_kind=error_kind,
@@ -649,3 +647,209 @@ def _should_retry_status(
)
)
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,
)
+39 -25
View File
@@ -82,6 +82,15 @@ def _json_object(value: object) -> dict[str, Any]:
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"))
@@ -277,7 +286,10 @@ class Session:
created_at: datetime = field(default_factory=datetime.now)
updated_at: datetime = field(default_factory=datetime.now)
metadata: dict[str, Any] = field(default_factory=dict)
last_consolidated: int = 0 # Number of messages already consolidated to files
# Legacy storage name for the Memory ingestion watermark. New code should
# 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)
policy: SessionPolicy = field(default_factory=SessionPolicy, repr=False, compare=False)
@@ -295,6 +307,15 @@ class Session:
):
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:
"""Add a message to the session."""
msg = {
@@ -319,9 +340,9 @@ class Session:
A positive ``max_messages`` applies an explicit caller-owned count
limit. The normal model path relies on ``max_tokens`` instead.
"""
replay_start = self.last_consolidated
replay_start = self.last_archived
if replay_start:
# ``last_consolidated`` is archive progress, not a replay boundary.
# ``last_archived`` is archive progress, not a replay boundary.
# Keep a small raw suffix for continuity, extending back to the user
# that started an assistant/tool sequence when necessary.
recent_start = recent_message_start_index(
@@ -335,8 +356,8 @@ class Session:
if max_messages <= 0:
start_idx = 0
else:
unarchived_count = len(self.messages) - self.last_consolidated
if replay_start < self.last_consolidated and unarchived_count < max_messages:
unarchived_count = len(self.messages) - self.last_archived
if replay_start < self.last_archived and unarchived_count < max_messages:
# The archived replay suffix can exceed the nominal count when one
# tool-heavy turn spans the boundary. Preserve that complete turn.
start_idx = 0
@@ -459,7 +480,7 @@ class Session:
def clear(self) -> None:
"""Clear all messages and reset session to initial state."""
self.messages = []
self.last_consolidated = 0
self.last_archived = 0
self.provider_state = None
self.updated_at = datetime.now()
self.metadata.pop("_last_summary", None)
@@ -474,11 +495,11 @@ class Session:
Returns a RetentionResult with dropped messages and how many of those
were in the already-consolidated prefix. This method mutates
self.messages and self.last_consolidated in place.
self.messages and self.last_archived in place.
"""
if max_messages <= 0:
dropped = list(self.messages)
lc = self.last_consolidated
lc = self.last_archived
self.clear()
return RetentionResult(
dropped=dropped,
@@ -491,7 +512,7 @@ class Session:
)
original = list(self.messages)
before_lc = self.last_consolidated
before_lc = self.last_archived
start_idx = max(0, len(self.messages) - max_messages)
if extend_to_user:
@@ -551,7 +572,7 @@ class Session:
if i < before_lc and id(m) not in retained_ids
)
# New last_consolidated = count of retained messages that were inside
# New last_archived = count of retained messages that were inside
# the old consolidated prefix.
new_lc = sum(
1 for i, m in enumerate(original)
@@ -559,7 +580,7 @@ class Session:
)
self.messages = retained
self.last_consolidated = new_lc
self.last_archived = new_lc
if dropped:
self.provider_state = None
self.updated_at = datetime.now()
@@ -1167,12 +1188,7 @@ class JsonlSessionStore:
if isinstance(updated_at_value, str) and updated_at_value
else None
)
offset = cast(object, data.get("last_consolidated", 0))
last_consolidated = (
offset
if isinstance(offset, int) and not isinstance(offset, bool)
else 0
)
last_consolidated = _archive_offset(data)
elif record_type == _PROVIDER_STATE_RECORD_TYPE:
provider_state = ProviderConversationState.from_private_record(
data.get("state")
@@ -1254,12 +1270,7 @@ class JsonlSessionStore:
if isinstance(updated_at_value, str) and updated_at_value:
with suppress(ValueError):
updated_at = datetime.fromisoformat(updated_at_value)
offset = cast(object, data.get("last_consolidated", 0))
last_consolidated = (
offset
if isinstance(offset, int) and not isinstance(offset, bool)
else 0
)
last_consolidated = _archive_offset(data)
elif record_type == _PROVIDER_STATE_RECORD_TYPE:
candidate = ProviderConversationState.from_private_record(
data.get("state")
@@ -1419,6 +1430,9 @@ class JsonlSessionStore:
"created_at": session.created_at.isoformat(),
"updated_at": session.updated_at.isoformat(),
"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,
}
f.write(json.dumps(metadata_line, ensure_ascii=False) + "\n")
@@ -2011,8 +2025,8 @@ class SessionManager:
for key in _FORK_VOLATILE_METADATA_KEYS:
metadata.pop(key, None)
last_consolidated = min(source.last_consolidated, len(copied))
if source.last_consolidated > len(copied):
last_consolidated = min(source.last_archived, len(copied))
if source.last_archived > len(copied):
metadata.pop("_last_summary", None)
last_consolidated = 0
+3 -7
View File
@@ -147,10 +147,10 @@ def prepare_save_boundary(ctx: TurnContext) -> None:
if ctx.session is not None:
clear_internal_continuation_state(ctx.session.metadata)
assert ctx.transcript_input is not None
ctx.save_skip = _save_skip_for_turn(
message_metadata=ctx.msg.metadata,
initial_message_count=len(ctx.initial_messages),
history_count=len(ctx.history),
initial_message_count=ctx.transcript_input.message_count,
input_persisted_early=ctx.input_persisted_early,
)
@@ -185,7 +185,6 @@ def _save_skip_for_turn(
*,
message_metadata: Mapping[str, Any] | None,
initial_message_count: int,
history_count: int,
input_persisted_early: bool,
) -> int:
"""Return the persisted-message append boundary for this turn."""
@@ -193,10 +192,7 @@ def _save_skip_for_turn(
return initial_message_count
if internal_continuation_inbound(message_metadata):
return initial_message_count
# 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:
if not input_persisted_early:
return initial_message_count - 1
return initial_message_count
+35 -18
View File
@@ -1,25 +1,42 @@
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.
Create a compact replacement checkpoint for this session.
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
When `[Archived Context Summary]` appears in the system prompt, update that previous checkpoint to reflect the current conversation state.
Format each fact as:
- [mark] fact content
## Merge rules
Marks (choose the best match):
- [permanent] Core preferences, personal traits, habits — never becomes stale
- [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
- Use the latest correction or decision as the current version of a fact, and merge duplicates.
- Preserve exact names, identifiers, paths, commands, decisions, results, and unresolved blockers when they are needed to continue the session.
- Retain a fact already present in long-term memory when it is needed for session continuity.
Priority: user corrections and preferences > solutions > decisions > events > environment facts.
## What to retain
Do not output facts already present in the system prompt's Recent History.
Always retain a compact working-state handoff:
- active objective
- current status
- completed results that constrain later work
- unresolved blockers
- next action
- exact identifiers needed for that action
Do not mark something [skip] merely because it might already exist in long-term memory.
Mark working-state facts `[ephemeral]`.
Return only formatted fact lines, or `(nothing)` if nothing noteworthy happened.
For other facts, retain a candidate only when it meets all four SNIP criteria:
- 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.
+1 -1
View File
@@ -40,7 +40,7 @@
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 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`; when editing a specific numbered line, pass that exact line as `line_hint`; add `occurrence` or `expected_replacements` when ambiguity matters.
- Use `edit_file` only for small exact replacements in one file, with `old_text` copied from `read_file`.
- 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`.
+1 -1
View File
@@ -44,7 +44,7 @@ def session_context_payload(session: Session) -> dict[str, Any]:
"schema_version": 1,
"session_key": session.key,
"total_messages": len(session.messages),
"archived_messages": min(session.last_consolidated, len(session.messages)),
"archived_messages": min(session.last_archived, len(session.messages)),
"replay_messages": len(replay),
"estimated_replay_tokens": replay_tokens,
"estimated_summary_tokens": summary_tokens,
+32
View File
@@ -28,6 +28,10 @@ from nanobot.config.loader import resolve_config_env_vars
from nanobot.config.schema import Config, FallbackCandidate, ModelPresetConfig, ProviderConfig
from nanobot.providers.image_generation import get_image_gen_provider
from nanobot.providers.oauth_guidance import OAUTH_CLI_KIT_MISSING_MESSAGE
from nanobot.providers.oauth_model_catalog import (
get_oauth_model_catalog,
invalidate_oauth_model_catalog,
)
from nanobot.providers.registry import PROVIDERS, create_dynamic_spec, find_by_name
from nanobot.webui.settings_contracts import (
QueryParams,
@@ -661,6 +665,30 @@ def provider_models_payload(
"models": rows,
"model_count": len(rows),
}
if catalog_kind == "hybrid":
proxy = _resolve_env_placeholders(provider_config.proxy)
catalog = get_oauth_model_catalog(spec.name, proxy=proxy)
rows = [
{
"id": model.id,
"label": model.label or None,
"description": model.description or None,
"owned_by": model.owned_by or spec.label,
"context_window": model.context_window,
"reasoning_efforts": list(model.reasoning_efforts),
"supports_backend_search": model.supports_backend_search,
}
for model in catalog.models
]
return {
**base_payload,
"status": "available",
"source": catalog.source,
"models": rows,
"model_count": len(rows),
"message": catalog.message,
"fetched_at": catalog.fetched_at,
}
api_base = _resolve_env_placeholders(provider_config.api_base) or spec.default_api_base
if spec.name == "openai" and not api_base:
@@ -1506,6 +1534,7 @@ def login_oauth_provider(
token = login_github_copilot(print_fn=lambda _message: None)
if not (token and token.access):
raise WebUISettingsError("OAuth login failed", status=401)
invalidate_oauth_model_catalog(spec.name)
return settings_payload(config_path=config_path)
if spec.name == "xai_grok":
@@ -1591,6 +1620,7 @@ def complete_oauth_provider(
oauth_flows.remove(spec.name, flow_id, flow, cancel=False)
if not token.access:
raise WebUISettingsError("OAuth login failed", status=401)
invalidate_oauth_model_catalog(spec.name)
return settings_payload(config_path=config_path)
@@ -1629,6 +1659,7 @@ def logout_oauth_provider(
oauth_flows.clear(spec.name)
logout_xai_oauth()
invalidate_oauth_model_catalog(spec.name)
return settings_payload(config_path=config_path)
else:
raise WebUISettingsError("OAuth logout is not supported for this provider")
@@ -1636,6 +1667,7 @@ def logout_oauth_provider(
for path in (token_path, token_path.with_suffix(".lock")):
with suppress(FileNotFoundError):
path.unlink()
invalidate_oauth_model_catalog(spec.name)
return settings_payload(config_path=config_path)
+5 -1
View File
@@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.agent.context import TranscriptInput
from nanobot.agent.loop import AgentLoop, TurnContext, TurnKind
from nanobot.agent.tools.context import RequestContext
from nanobot.agent.tools.filesystem import ReadFileTool
@@ -148,7 +149,10 @@ async def test_pending_document_attachment_keeps_body_out_of_prompt(
runtime = loop.llm_runtime()
result = await loop._run_agent_loop(
[{"role": "user", "content": "hello"}],
TranscriptInput(
history=[{"role": "user", "content": "hello"}],
current_message=None,
),
runtime=runtime,
request_context=RequestContext(channel="cli", chat_id="c", runtime=runtime),
pending_queue=pending_queue,
+9 -9
View File
@@ -88,11 +88,11 @@ def _make_fake_compact(
state["count"] += 1
session = loop.sessions.get_or_create(key)
tail = list(session.messages[session.last_consolidated:])
tail = list(session.messages[session.last_archived:])
if not tail:
loop.sessions.save(session)
return ""
archive_end = session.last_consolidated + len(tail)
archive_end = session.last_archived + len(tail)
archive_msgs = tail
last_active = session.updated_at
@@ -109,7 +109,7 @@ def _make_fake_compact(
"last_active": last_active.isoformat(),
}
session.last_consolidated = archive_end
session.last_archived = archive_end
loop.sessions.save(session)
return s
@@ -399,12 +399,12 @@ class TestAutoCompact:
await loop.aclose()
@pytest.mark.asyncio
async def test_auto_compact_respects_last_consolidated(self, tmp_path):
"""_archive should only archive un-consolidated messages."""
async def test_auto_compact_respects_last_archived(self, tmp_path):
"""_archive should process only unarchived messages."""
loop = _make_loop(tmp_path, session_ttl_minutes=15)
session = loop.sessions.get_or_create("cli:test")
_add_turns(session, 14)
session.last_consolidated = 18
session.last_archived = 18
loop.sessions.save(session)
archived_messages = []
@@ -1302,9 +1302,9 @@ class TestSummaryPersistence:
assert "_last_summary" in reloaded.metadata
# Simulate /new command
session.clear()
loop.sessions.save(session)
loop.sessions.invalidate(session.key)
reloaded.clear()
loop.sessions.save(reloaded)
loop.sessions.invalidate(reloaded.key)
# After /new, metadata should no longer contain _last_summary
fresh = loop.sessions.get_or_create("cli:test")
+3 -3
View File
@@ -16,7 +16,7 @@ def _runtime(_session: Session | None = None):
def _make_session(
key: str = "cli:test",
messages: list | None = None,
last_consolidated: int = 0,
last_archived: int = 0,
updated_at: datetime | None = None,
metadata: dict | None = None,
) -> Session:
@@ -25,8 +25,8 @@ def _make_session(
key=key,
messages=messages or [],
metadata=metadata or {},
last_consolidated=last_consolidated,
)
session.last_archived = last_archived
if updated_at is not None:
session.updated_at = updated_at
return session
@@ -408,7 +408,7 @@ class TestCheckExpired:
last_active = datetime(2026, 1, 1, 10, 0, 0)
session = _make_session("cli:done", updated_at=last_active)
_add_turns(session, 2)
session.last_consolidated = len(session.messages)
session.last_archived = len(session.messages)
mock_sm.list_sessions.return_value = [
{"key": "cli:done", "updated_at": last_active.isoformat()},
]
-650
View File
@@ -1,650 +0,0 @@
"""Test session management with cache-friendly message handling."""
import asyncio
from collections.abc import Coroutine
from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.session.manager import Session, SessionManager
# Test constants
MEMORY_WINDOW = 50
KEEP_COUNT = MEMORY_WINDOW // 2 # 25
def create_session_with_messages(key: str, count: int, role: str = "user") -> Session:
"""Create a session and add the specified number of messages.
Args:
key: Session identifier
count: Number of messages to add
role: Message role (default: "user")
Returns:
Session with the specified messages
"""
session = Session(key=key)
for i in range(count):
session.add_message(role, f"msg{i}")
return session
def assert_messages_content(messages: list, start_index: int, end_index: int) -> None:
"""Assert that messages contain expected content from start to end index.
Args:
messages: List of message dictionaries
start_index: Expected first message index
end_index: Expected last message index
"""
assert len(messages) > 0
assert messages[0]["content"] == f"msg{start_index}"
assert messages[-1]["content"] == f"msg{end_index}"
def get_old_messages(session: Session, last_consolidated: int, keep_count: int) -> list:
"""Extract messages that would be consolidated using the standard slice logic.
Args:
session: The session containing messages
last_consolidated: Index of last consolidated message
keep_count: Number of recent messages to keep
Returns:
List of messages that would be consolidated
"""
return session.messages[last_consolidated:-keep_count]
class TestSessionLastConsolidated:
"""Test last_consolidated tracking to avoid duplicate processing."""
def test_initial_last_consolidated_zero(self) -> None:
"""Test that new session starts with last_consolidated=0."""
session = Session(key="test:initial")
assert session.last_consolidated == 0
def test_last_consolidated_persistence(self, tmp_path) -> None:
"""Test that last_consolidated persists across save/load."""
manager = SessionManager(Path(tmp_path))
session1 = create_session_with_messages("test:persist", 20)
session1.last_consolidated = 15
manager.save(session1)
session2 = manager.get_or_create("test:persist")
assert session2.last_consolidated == 15
assert len(session2.messages) == 20
def test_clear_resets_last_consolidated(self) -> None:
"""Test that clear() resets last_consolidated to 0."""
session = create_session_with_messages("test:clear", 10)
session.last_consolidated = 5
session.clear()
assert len(session.messages) == 0
assert session.last_consolidated == 0
class TestSessionImmutableHistory:
"""Test Session message immutability for cache efficiency."""
def test_initial_state(self) -> None:
"""Test that new session has empty messages list."""
session = Session(key="test:initial")
assert len(session.messages) == 0
def test_add_messages_appends_only(self) -> None:
"""Test that adding messages only appends, never modifies."""
session = Session(key="test:preserve")
session.add_message("user", "msg1")
session.add_message("assistant", "resp1")
session.add_message("user", "msg2")
assert len(session.messages) == 3
assert session.messages[0]["content"] == "msg1"
def test_get_history_returns_most_recent(self) -> None:
"""Test get_history returns the most recent messages."""
session = Session(key="test:history")
for i in range(10):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
history = session.get_history(max_messages=6)
assert len(history) == 6
assert history[0]["content"] == "msg7"
assert history[-1]["content"] == "resp9"
def test_get_history_with_all_messages(self) -> None:
"""Test get_history with max_messages larger than actual."""
session = create_session_with_messages("test:all", 5)
history = session.get_history(max_messages=100)
assert len(history) == 5
assert history[0]["content"] == "msg0"
def test_get_history_stable_for_same_session(self) -> None:
"""Test that get_history returns same content for same max_messages."""
session = create_session_with_messages("test:stable", 20)
history1 = session.get_history(max_messages=10)
history2 = session.get_history(max_messages=10)
assert history1 == history2
def test_messages_list_never_modified(self) -> None:
"""Test that messages list is never modified after creation."""
session = create_session_with_messages("test:immutable", 5)
original_len = len(session.messages)
session.get_history(max_messages=2)
assert len(session.messages) == original_len
for _ in range(10):
session.get_history(max_messages=3)
assert len(session.messages) == original_len
class TestSessionPersistence:
"""Test Session persistence and reload."""
@pytest.fixture
def temp_manager(self, tmp_path):
return SessionManager(Path(tmp_path))
def test_persistence_roundtrip(self, temp_manager):
"""Test that messages persist across save/load."""
session1 = create_session_with_messages("test:persistence", 20)
temp_manager.save(session1)
session2 = temp_manager.get_or_create("test:persistence")
assert len(session2.messages) == 20
assert session2.messages[0]["content"] == "msg0"
assert session2.messages[-1]["content"] == "msg19"
def test_get_history_after_reload(self, temp_manager):
"""Test that get_history works correctly after reload."""
session1 = create_session_with_messages("test:reload", 30)
temp_manager.save(session1)
session2 = temp_manager.get_or_create("test:reload")
history = session2.get_history(max_messages=10)
assert len(history) == 10
assert history[0]["content"] == "msg20"
assert history[-1]["content"] == "msg29"
def test_clear_resets_session(self, temp_manager):
"""Test that clear() properly resets session."""
session = create_session_with_messages("test:clear", 10)
assert len(session.messages) == 10
session.clear()
assert len(session.messages) == 0
class TestConsolidationTriggerConditions:
"""Test consolidation trigger conditions and logic."""
def test_consolidation_needed_when_messages_exceed_window(self):
"""Test consolidation logic: should trigger when messages exceed the window."""
session = create_session_with_messages("test:trigger", 60)
total_messages = len(session.messages)
messages_to_process = total_messages - session.last_consolidated
assert total_messages > MEMORY_WINDOW
assert messages_to_process > 0
expected_consolidate_count = total_messages - KEEP_COUNT
assert expected_consolidate_count == 35
def test_consolidation_skipped_when_within_keep_count(self):
"""Test consolidation skipped when total messages <= keep_count."""
session = create_session_with_messages("test:skip", 20)
total_messages = len(session.messages)
assert total_messages <= KEEP_COUNT
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
assert len(old_messages) == 0
def test_consolidation_skipped_when_no_new_messages(self):
"""Test consolidation skipped when messages_to_process <= 0."""
session = create_session_with_messages("test:already_consolidated", 40)
session.last_consolidated = len(session.messages) - KEEP_COUNT # 15
# Add a few more messages
for i in range(40, 42):
session.add_message("user", f"msg{i}")
total_messages = len(session.messages)
messages_to_process = total_messages - session.last_consolidated
assert messages_to_process > 0
# Simulate last_consolidated catching up
session.last_consolidated = total_messages - KEEP_COUNT
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
assert len(old_messages) == 0
class TestLastConsolidatedEdgeCases:
"""Test last_consolidated edge cases and data corruption scenarios."""
def test_last_consolidated_exceeds_message_count(self):
"""Test behavior when last_consolidated > len(messages) (data corruption)."""
session = create_session_with_messages("test:corruption", 10)
session.last_consolidated = 20
total_messages = len(session.messages)
messages_to_process = total_messages - session.last_consolidated
assert messages_to_process <= 0
old_messages = get_old_messages(session, session.last_consolidated, 5)
assert len(old_messages) == 0
def test_last_consolidated_negative_value(self):
"""Test behavior with negative last_consolidated (invalid state)."""
session = create_session_with_messages("test:negative", 10)
session.last_consolidated = -5
keep_count = 3
old_messages = get_old_messages(session, session.last_consolidated, keep_count)
# messages[-5:-3] with 10 messages gives indices 5,6
assert len(old_messages) == 2
assert old_messages[0]["content"] == "msg5"
assert old_messages[-1]["content"] == "msg6"
def test_messages_added_after_consolidation(self):
"""Test correct behavior when new messages arrive after consolidation."""
session = create_session_with_messages("test:new_messages", 40)
session.last_consolidated = len(session.messages) - KEEP_COUNT # 15
# Add new messages after consolidation
for i in range(40, 50):
session.add_message("user", f"msg{i}")
total_messages = len(session.messages)
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
expected_consolidate_count = total_messages - KEEP_COUNT - session.last_consolidated
assert len(old_messages) == expected_consolidate_count
assert_messages_content(old_messages, 15, 24)
def test_slice_behavior_when_indices_overlap(self):
"""Test slice behavior when last_consolidated >= total - keep_count."""
session = create_session_with_messages("test:overlap", 30)
session.last_consolidated = 12
old_messages = get_old_messages(session, session.last_consolidated, 20)
assert len(old_messages) == 0
class TestArchiveAllMode:
"""Test archive_all mode (used by /new command)."""
def test_archive_all_consolidates_everything(self):
"""Test archive_all=True consolidates all messages."""
session = create_session_with_messages("test:archive_all", 50)
archive_all = True
if archive_all:
old_messages = session.messages
assert len(old_messages) == 50
assert session.last_consolidated == 0
def test_archive_all_resets_last_consolidated(self):
"""Test that archive_all mode resets last_consolidated to 0."""
session = create_session_with_messages("test:reset", 40)
session.last_consolidated = 15
archive_all = True
if archive_all:
session.last_consolidated = 0
assert session.last_consolidated == 0
assert len(session.messages) == 40
def test_archive_all_vs_normal_consolidation(self):
"""Test difference between archive_all and normal consolidation."""
# Normal consolidation
session1 = create_session_with_messages("test:normal", 60)
session1.last_consolidated = len(session1.messages) - KEEP_COUNT
# archive_all mode
session2 = create_session_with_messages("test:all", 60)
session2.last_consolidated = 0
assert session1.last_consolidated == 35
assert len(session1.messages) == 60
assert session2.last_consolidated == 0
assert len(session2.messages) == 60
class TestCacheImmutability:
"""Test that consolidation doesn't modify session.messages (cache safety)."""
def test_consolidation_does_not_modify_messages_list(self):
"""Test that consolidation leaves messages list unchanged."""
session = create_session_with_messages("test:immutable", 50)
original_messages = session.messages.copy()
original_len = len(session.messages)
session.last_consolidated = original_len - KEEP_COUNT
assert len(session.messages) == original_len
assert session.messages == original_messages
def test_get_history_does_not_modify_messages(self):
"""Test that get_history doesn't modify messages list."""
session = create_session_with_messages("test:history_immutable", 40)
original_messages = [m.copy() for m in session.messages]
for _ in range(5):
history = session.get_history(max_messages=10)
assert len(history) == 10
assert len(session.messages) == 40
for i, msg in enumerate(session.messages):
assert msg["content"] == original_messages[i]["content"]
def test_consolidation_only_updates_last_consolidated(self):
"""Test that consolidation only updates last_consolidated field."""
session = create_session_with_messages("test:field_only", 60)
original_messages = session.messages.copy()
original_key = session.key
original_metadata = session.metadata.copy()
session.last_consolidated = len(session.messages) - KEEP_COUNT
assert session.messages == original_messages
assert session.key == original_key
assert session.metadata == original_metadata
assert session.last_consolidated == 35
class TestSliceLogic:
"""Test the slice logic: messages[last_consolidated:-keep_count]."""
def test_slice_extracts_correct_range(self):
"""Test that slice extracts the correct message range."""
session = create_session_with_messages("test:slice", 60)
old_messages = get_old_messages(session, 0, KEEP_COUNT)
assert len(old_messages) == 35
assert_messages_content(old_messages, 0, 34)
remaining = session.messages[-KEEP_COUNT:]
assert len(remaining) == 25
assert_messages_content(remaining, 35, 59)
def test_slice_with_partial_consolidation(self):
"""Test slice when some messages already consolidated."""
session = create_session_with_messages("test:partial", 70)
last_consolidated = 30
old_messages = get_old_messages(session, last_consolidated, KEEP_COUNT)
assert len(old_messages) == 15
assert_messages_content(old_messages, 30, 44)
def test_slice_with_various_keep_counts(self):
"""Test slice behavior with different keep_count values."""
session = create_session_with_messages("test:keep_counts", 50)
test_cases = [(10, 40), (20, 30), (30, 20), (40, 10)]
for keep_count, expected_count in test_cases:
old_messages = session.messages[0:-keep_count]
assert len(old_messages) == expected_count
def test_slice_when_keep_count_exceeds_messages(self):
"""Test slice when keep_count > len(messages)."""
session = create_session_with_messages("test:exceed", 10)
old_messages = session.messages[0:-20]
assert len(old_messages) == 0
class TestEmptyAndBoundarySessions:
"""Test empty sessions and boundary conditions."""
def test_empty_session_consolidation(self):
"""Test consolidation behavior with empty session."""
session = Session(key="test:empty")
assert len(session.messages) == 0
assert session.last_consolidated == 0
messages_to_process = len(session.messages) - session.last_consolidated
assert messages_to_process == 0
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
assert len(old_messages) == 0
def test_single_message_session(self):
"""Test consolidation with single message."""
session = Session(key="test:single")
session.add_message("user", "only message")
assert len(session.messages) == 1
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
assert len(old_messages) == 0
def test_exactly_keep_count_messages(self):
"""Test session with exactly keep_count messages."""
session = create_session_with_messages("test:exact", KEEP_COUNT)
assert len(session.messages) == KEEP_COUNT
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
assert len(old_messages) == 0
def test_just_over_keep_count(self):
"""Test session with one message over keep_count."""
session = create_session_with_messages("test:over", KEEP_COUNT + 1)
assert len(session.messages) == 26
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
assert len(old_messages) == 1
assert old_messages[0]["content"] == "msg0"
def test_very_large_session(self):
"""Test consolidation with very large message count."""
session = create_session_with_messages("test:large", 1000)
assert len(session.messages) == 1000
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
assert len(old_messages) == 975
assert_messages_content(old_messages, 0, 974)
remaining = session.messages[-KEEP_COUNT:]
assert len(remaining) == 25
assert_messages_content(remaining, 975, 999)
def test_session_with_gaps_in_consolidation(self):
"""Test session with potential gaps in consolidation history."""
session = create_session_with_messages("test:gaps", 50)
session.last_consolidated = 10
# Add more messages
for i in range(50, 60):
session.add_message("user", f"msg{i}")
old_messages = get_old_messages(session, session.last_consolidated, KEEP_COUNT)
expected_count = 60 - KEEP_COUNT - 10
assert len(old_messages) == expected_count
assert_messages_content(old_messages, 10, 34)
class TestNewCommandArchival:
"""Test /new archival behavior with the simplified consolidation flow."""
@staticmethod
def _make_loop(tmp_path: Path):
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import GenerationSettings, LLMResponse
bus = MessageBus()
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.estimate_prompt_tokens.return_value = (10_000, "test")
provider.generation = GenerationSettings(max_tokens=100)
loop = AgentLoop(
bus=bus,
provider=provider,
workspace=tmp_path,
model="test-model",
context_window_tokens=1,
)
loop.provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="ok", tool_calls=[]))
loop.tools.get_definitions = MagicMock(return_value=[])
return loop
@pytest.mark.asyncio
async def test_new_clears_session_immediately_even_if_archive_fails(self, tmp_path: Path) -> None:
"""/new clears session immediately; archive is fire-and-forget."""
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
session = loop.sessions.get_or_create("cli:test")
for i in range(5):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
loop.sessions.save(session)
call_count = 0
expected_runtime = loop.llm_runtime()
async def _failing_summarize(session, *, archive_end, runtime) -> None:
nonlocal call_count
assert runtime is expected_runtime
assert session.key == "cli:test"
assert archive_end == len(session.messages)
call_count += 1
loop.consolidator.archive_session = _failing_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime)
assert response is not None
assert "new session started" in response.content.lower()
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
await loop.aclose()
assert call_count == 1
@pytest.mark.asyncio
async def test_new_reuses_replay_prefix_and_archives_only_unconsolidated_messages(
self,
tmp_path: Path,
) -> None:
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
loop.set_runtime_context_window(128_000)
session = loop.sessions.get_or_create("cli:test")
for i in range(5):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
session.last_consolidated = len(session.messages) - 2
ordinary_history = session.get_history()
assert [message["content"] for message in ordinary_history] == [
"msg1",
"resp1",
"msg2",
"resp2",
"msg3",
"resp3",
"msg4",
"resp4",
]
loop.sessions.save(session)
expected_runtime = loop.llm_runtime()
scheduled: list[Coroutine[Any, Any, object]] = []
loop.schedule_background = scheduled.append # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime)
assert response is not None
assert "new session started" in response.content.lower()
assert len(scheduled) == 1
await scheduled[0]
await loop.aclose()
sent = loop.provider.chat_with_retry.call_args.kwargs["messages"]
assert sent[1:-1] == ordinary_history
assert "final 2 conversation messages" in sent[-1]["content"]
@pytest.mark.asyncio
async def test_new_clears_session_and_responds(self, tmp_path: Path) -> None:
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
session = loop.sessions.get_or_create("cli:test")
for i in range(3):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
loop.sessions.save(session)
expected_runtime = loop.llm_runtime()
async def _ok_summarize(session, *, archive_end, runtime) -> str:
assert runtime is expected_runtime
assert session.key == "cli:test"
assert archive_end == len(session.messages)
return "Summary."
loop.consolidator.archive_session = _ok_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime)
assert response is not None
assert "new session started" in response.content.lower()
assert loop.sessions.get_or_create("cli:test").messages == []
@pytest.mark.asyncio
async def test_aclose_drains_background_tasks(self, tmp_path: Path) -> None:
"""aclose waits for background tasks to complete."""
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
session = loop.sessions.get_or_create("cli:test")
for i in range(3):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
loop.sessions.save(session)
archived = asyncio.Event()
release_archive = asyncio.Event()
expected_runtime = loop.llm_runtime()
async def _slow_summarize(session, *, archive_end, runtime) -> str:
assert runtime is expected_runtime
assert session.key == "cli:test"
assert archive_end == len(session.messages)
await release_archive.wait()
archived.set()
return "Summary."
loop.consolidator.archive_session = _slow_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
await loop._process_message(new_msg, runtime=expected_runtime)
assert not archived.is_set()
release_archive.set()
await loop.aclose()
assert archived.is_set()
-112
View File
@@ -1,112 +0,0 @@
"""Tests for configurable consolidation_ratio."""
from unittest.mock import AsyncMock, MagicMock
import pytest
from pydantic import ValidationError
import nanobot.agent.memory as memory_module
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import GenerationSettings, LLMResponse
def _make_loop(
tmp_path,
*,
estimated_tokens: int = 0,
context_window_tokens: int = 200,
consolidation_ratio: float = 0.5,
) -> AgentLoop:
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings(max_tokens=0)
provider.estimate_prompt_tokens.return_value = (estimated_tokens, "test-counter")
_response = LLMResponse(content="ok", tool_calls=[])
provider.chat_with_retry = AsyncMock(return_value=_response)
provider.chat_stream_with_retry = AsyncMock(return_value=_response)
loop = AgentLoop(
bus=MessageBus(),
provider=provider,
workspace=tmp_path,
model="test-model",
context_window_tokens=context_window_tokens,
consolidation_ratio=consolidation_ratio,
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator._SAFETY_BUFFER = 0
return loop
def _session_with_turns(loop: AgentLoop, *, turns: int):
session = loop.sessions.get_or_create("cli:test")
session.messages = []
for i in range(turns):
session.messages.append({"role": "user", "content": f"u{i}", "timestamp": f"2026-01-01T00:00:{i:02d}"})
session.messages.append({"role": "assistant", "content": f"a{i}", "timestamp": f"2026-01-01T00:01:{i:02d}"})
loop.sessions.save(session)
return session
@pytest.mark.asyncio
@pytest.mark.parametrize(
("ratio", "context_window_tokens", "estimates", "expected_archives"),
[
(0.5, 200, [250, 90], 1),
(0.1, 1000, [1200, 800, 400, 50], 2),
(0.9, 200, [300, 175], 1),
],
)
async def test_consolidation_ratio_controls_target(
tmp_path,
monkeypatch,
ratio: float,
context_window_tokens: int,
estimates: list[int],
expected_archives: int,
) -> None:
loop = _make_loop(
tmp_path,
context_window_tokens=context_window_tokens,
consolidation_ratio=ratio,
)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = _session_with_turns(loop, turns=10)
remaining_estimates = list(estimates)
runtime = loop.llm_runtime()
def mock_estimate(_session, *, runtime):
return (remaining_estimates.pop(0), "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=runtime,
)
assert loop.consolidator.archive_session.await_count == expected_archives
def test_ratio_propagated_from_config_schema() -> None:
defaults = AgentDefaults()
assert defaults.consolidation_ratio == 0.5
defaults = AgentDefaults.model_validate({"consolidationRatio": 0.3})
assert defaults.consolidation_ratio == 0.3
dumped = defaults.model_dump(by_alias=True)
assert dumped["consolidationRatio"] == 0.3
def test_ratio_validation_rejects_out_of_range() -> None:
with pytest.raises(ValidationError):
AgentDefaults(consolidation_ratio=0.05)
with pytest.raises(ValidationError):
AgentDefaults(consolidation_ratio=1.0)
+285 -82
View File
@@ -1,4 +1,4 @@
"""Tests for the lightweight Consolidator — append-only to HISTORY.md."""
"""Tests for Memory checkpoint consolidation and history journaling."""
from dataclasses import replace
from unittest.mock import AsyncMock, MagicMock
@@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.agent.memory import (
_ARCHIVE_SUMMARY_MAX_CHARS,
_HISTORY_ENTRY_HARD_CAP,
Consolidator,
MemoryStore,
)
@@ -26,6 +26,8 @@ from nanobot.session.manager import Session
from nanobot.utils.llm_runtime import LLMRuntime
from nanobot.utils.prompt_templates import render_template
_ARCHIVE_PROMPT = render_template("agent/consolidator_archive.md", strip=True)
@pytest.fixture
def store(tmp_path):
@@ -98,8 +100,15 @@ def _build_test_messages(**kwargs):
]
async def _archive(consolidator, messages, runtime, *, session_key="test:session"):
return await consolidator.archive(
async def _archive(
consolidator,
messages,
runtime,
*,
session_key="test:session",
previous_summary=None,
):
return await consolidator.archiver.archive(
messages,
runtime=runtime,
session_key=session_key,
@@ -108,6 +117,7 @@ async def _archive(consolidator, messages, runtime, *, session_key="test:session
current_message="consolidate",
),
request_tools=[],
previous_summary=previous_summary,
)
@@ -201,7 +211,9 @@ class TestConsolidatorSummarize:
mock_provider.chat_with_retry.side_effect = Exception("API error")
messages = [{"role": "user", "content": "hello"}]
result = await _archive(consolidator, messages, runtime)
assert result is None # no summary on raw dump fallback
assert result is not None
assert "[RAW]" in result
assert "hello" in result
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1
assert "[RAW]" in entries[0]["content"]
@@ -226,22 +238,52 @@ class TestConsolidatorSummarize:
entries = store.read_unprocessed_history(since_cursor=0)
assert entries[0]["session_key"] == "slack:chat-2"
async def test_raw_fallback_represents_previous_checkpoint_and_new_chunk(
self,
consolidator,
mock_provider,
runtime,
):
runtime = replace(runtime, generation=GenerationSettings(max_tokens=96))
mock_provider.chat_with_retry.side_effect = RuntimeError("API error")
result = await _archive(
consolidator,
[{"role": "user", "content": "NEW_MARKER " + "new " * 200}],
runtime,
previous_summary="OLD_MARKER " + "old " * 200,
)
assert result is not None
assert "[Previous archived context]" in result
assert "OLD_MARKER" in result
assert "[Newly archived raw context]" in result
assert "NEW_MARKER" in result
assert "... (truncated)" in result
async def test_summarize_skips_empty_messages(self, consolidator, runtime):
result = await _archive(consolidator, [], runtime)
assert result is None
class TestConsolidatorPromptContract:
def test_archive_prompt_outputs_attribute_tags_without_missing_context_claims(self):
prompt = render_template("agent/consolidator_archive.md", strip=True, archive_count=4)
def test_archive_prompt_requests_a_cumulative_replacement_checkpoint(self):
prompt = _ARCHIVE_PROMPT
for section in ("## Merge rules", "## What to retain", "## Output"):
assert section in prompt
assert "replacement checkpoint" in prompt
assert "[Archived Context Summary]" in prompt
assert "current conversation state" in prompt
assert "SNIP" in prompt
assert "final 4 conversation messages" in prompt
for mark in ("[permanent]", "[durable]", "[ephemeral]", "[correction]", "[skip]"):
for mark in ("[permanent]", "[durable]", "[ephemeral]", "[correction]"):
assert mark in prompt
assert "check context below" not in prompt.lower()
assert "Do not output facts already present in the system prompt's Recent History" in prompt
assert "Do not mark something [skip] merely because it might already exist" in prompt
assert "working-state handoff" in prompt
assert "- [mark] fact" in prompt
assert "[skip]" not in prompt
assert "(nothing)" in prompt
assert "history.jsonl" not in prompt
class TestConsolidatorArchiveErrorHandling:
"""archive() must fall back when the LLM does not complete its overview.
@@ -270,7 +312,8 @@ class TestConsolidatorArchiveErrorHandling:
{"role": "assistant", "content": "Done, fixed the race condition."},
]
result = await _archive(consolidator, messages, runtime)
assert result is None
assert result is not None
assert "[RAW]" in result
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1
assert "[RAW]" in entries[0]["content"]
@@ -342,7 +385,7 @@ class TestConsolidatorTokenBudget:
):
"""No consolidation when tokens are within budget."""
session = MagicMock()
session.last_consolidated = 0
session.last_archived = 0
session.messages = [{"role": "user", "content": "hi"}]
session.key = "test:key"
consolidator.sessions._session_cache[session.key] = session
@@ -362,7 +405,7 @@ class TestConsolidatorTokenBudget:
with pytest.raises(RuntimeError, match="counter failed"):
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
async def test_estimate_uses_full_unconsolidated_tail(self, consolidator, runtime):
async def test_estimate_uses_full_unarchived_tail(self, consolidator, runtime):
"""Consolidation pressure must account for the full unarchived tail."""
session = Session(key="test:full-tail")
for i in range(160):
@@ -385,7 +428,7 @@ class TestConsolidatorTokenBudget:
session = Session(key="test:archived-replay")
for i in range(10):
session.add_message("user", f"msg-{i}")
session.last_consolidated = len(session.messages)
session.last_archived = len(session.messages)
captured: dict[str, list[dict]] = {}
@@ -420,8 +463,8 @@ class TestConsolidatorTokenBudget:
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
)
consolidator.pick_consolidation_boundary = MagicMock(return_value=(50, 800))
consolidator._build_messages = MagicMock(side_effect=_build_test_messages)
consolidator.pick_consolidation_boundary = MagicMock(return_value=50)
consolidator.archiver._build_messages = MagicMock(side_effect=_build_test_messages)
mock_provider.estimate_prompt_tokens.return_value = (100, "test-counter")
mock_provider.chat_with_retry.return_value = LLMResponse(
content="Token overflow summary.",
@@ -434,13 +477,13 @@ class TestConsolidatorTokenBudget:
assert [message["content"] for message in request["messages"][1:-1]] == [
f"m{i}" for i in range(50)
]
assert "final 50 conversation messages" in request["messages"][-1]["content"]
assert request["messages"][-1]["content"] == _ARCHIVE_PROMPT
assert request["tools"] == []
assert request["tool_choice"] == "none"
assert session.last_consolidated == 50
assert session.provider_state is None
assert "tool_choice" not in request
assert session.last_archived == 50
assert session.provider_state == _provider_state()
async def test_raw_archive_fallback_advances_last_consolidated(
async def test_raw_archive_fallback_advances_archive_watermark(
self, consolidator, runtime
):
"""When archive() falls back to raw-archive (LLM failed), the cursor
@@ -448,27 +491,26 @@ class TestConsolidatorTokenBudget:
on every subsequent maybe_consolidate_by_tokens() call, spamming
duplicate [RAW] entries into history.jsonl."""
consolidator._SAFETY_BUFFER = 0
session = MagicMock()
session.last_consolidated = 0
session.key = "test:key"
session = Session(key="test:key")
session.provider_state = _provider_state()
session.messages = [
{"role": "user" if i in {0, 50} else "assistant", "content": f"m{i}"}
for i in range(70)
]
session.metadata = {}
consolidator.sessions._session_cache[session.key] = session
consolidator.estimate_session_prompt_tokens = MagicMock(
side_effect=[(1200, "tiktoken"), (400, "tiktoken")]
)
# LLM consolidation fails after raw_archive fires.
consolidator.archive_session = AsyncMock(return_value=None)
consolidator.archive_session = AsyncMock(return_value="[RAW] checkpoint")
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive_session.assert_awaited_once()
# The chunk is considered "materialized" (as a raw-archive breadcrumb),
# so last_consolidated must have moved past it.
assert session.last_consolidated == 50
# so the archive watermark must have moved past it without touching
# the provider-owned continuation state.
assert session.last_archived == 50
assert session.provider_state == _provider_state()
async def test_raw_archive_fallback_breaks_round_loop(
self, consolidator, runtime
@@ -477,7 +519,7 @@ class TestConsolidatorTokenBudget:
same maybe_consolidate_by_tokens invocation bail after one fallback."""
consolidator._SAFETY_BUFFER = 0
session = MagicMock()
session.last_consolidated = 0
session.last_archived = 0
session.key = "test:key"
session.messages = [
{"role": "user" if i in {0, 20, 40, 60} else "assistant", "content": f"m{i}"}
@@ -489,11 +531,11 @@ class TestConsolidatorTokenBudget:
consolidator.estimate_session_prompt_tokens = MagicMock(
return_value=(1200, "tiktoken")
)
consolidator.archive_session = AsyncMock(return_value=None)
consolidator.archive_session = AsyncMock(return_value="[RAW] checkpoint")
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
# Exactly one fallback per call — not _MAX_CONSOLIDATION_ROUNDS.
# The fixed policy archives at most one prefix per call.
assert consolidator.archive_session.await_count == 1
async def test_boundary_respected_when_no_intermediate_user_turn(
@@ -502,7 +544,7 @@ class TestConsolidatorTokenBudget:
"""When boundary points past a long tool chain, the full chunk is archived."""
consolidator._SAFETY_BUFFER = 0
session = MagicMock()
session.last_consolidated = 0
session.last_archived = 0
session.key = "test:key"
session.messages = [
{
@@ -520,8 +562,8 @@ class TestConsolidatorTokenBudget:
await consolidator.maybe_consolidate_by_tokens(session, runtime=runtime)
consolidator.archive_session.assert_awaited_once()
# pick_consolidation_boundary finds the only boundary at idx=61
assert session.last_consolidated == 61
# The fixed recent tail expands backward to the user at idx=61.
assert session.last_archived == 61
class TestCompactIdleSession:
@@ -575,8 +617,8 @@ class TestCompactIdleSession:
reloaded = sessions.get_or_create("cli:test")
assert len(reloaded.messages) == 40
assert reloaded.messages[0]["content"] == "user msg 0"
assert reloaded.last_consolidated == 40
assert reloaded.provider_state is None
assert reloaded.last_archived == 40
assert reloaded.provider_state == _provider_state()
visible = reloaded.get_history(max_messages=40)
assert len(visible) == 8
assert visible[0]["content"] == "user msg 16"
@@ -608,30 +650,65 @@ class TestCompactIdleSession:
mock_provider.chat_with_retry.assert_awaited_once()
assert len(store.read_unprocessed_history(since_cursor=0)) == 1
reloaded = sessions.get_or_create("cli:short")
assert reloaded.last_consolidated == 2
assert reloaded.last_archived == 2
assert [message["content"] for message in reloaded.get_history()] == ["hello", "hi"]
@pytest.mark.asyncio
async def test_idle_compaction_with_no_new_messages_is_noop(
self, real_consolidator, mock_provider, store, runtime
):
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:archived-idle")
session.add_message("user", "already archived")
session.add_message("assistant", "old answer")
session.last_archived = 2
sessions.save(session)
sessions.invalidate("cli:archived-idle")
result = await real_consolidator.compact_idle_session(
"cli:archived-idle",
runtime=runtime,
)
assert result == ""
mock_provider.chat_with_retry.assert_not_awaited()
reloaded = sessions.get_or_create("cli:archived-idle")
assert reloaded.last_archived == 2
assert "_last_summary" not in reloaded.metadata
assert store.read_unprocessed_history(since_cursor=0) == []
@pytest.mark.asyncio
async def test_new_messages_advance_existing_archive_progress(
self, real_consolidator, mock_provider, runtime
):
mock_provider.chat_with_retry.return_value = MagicMock(
content="Summary.", finish_reason="stop"
)
mock_provider.chat_with_retry.side_effect = [
MagicMock(content="First replacement checkpoint.", finish_reason="stop"),
MagicMock(content="Second replacement checkpoint.", finish_reason="stop"),
]
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:incremental")
session.add_message("user", "first user")
session.add_message("assistant", "first assistant")
sessions.save(session)
await real_consolidator.compact_idle_session("cli:incremental", runtime=runtime)
first = await real_consolidator.compact_idle_session(
"cli:incremental",
runtime=runtime,
)
current = sessions.get_or_create("cli:incremental")
current.add_message("user", "second user")
current.add_message("assistant", "second assistant")
sessions.save(current)
await real_consolidator.compact_idle_session("cli:incremental", runtime=runtime)
second = await real_consolidator.compact_idle_session(
"cli:incremental",
runtime=runtime,
)
assert first == "First replacement checkpoint."
assert second == "Second replacement checkpoint."
assert mock_provider.chat_with_retry.await_count == 2
latest_build = real_consolidator.archiver._build_messages.call_args_list[-1].kwargs
assert latest_build["session_summary"]["text"] == "First replacement checkpoint."
latest_messages = mock_provider.chat_with_retry.await_args_list[-1].kwargs["messages"]
assert [message["content"] for message in latest_messages[1:5]] == [
"first user",
@@ -639,8 +716,91 @@ class TestCompactIdleSession:
"second user",
"second assistant",
]
assert "final 2 conversation messages" in latest_messages[-1]["content"]
assert sessions.get_or_create("cli:incremental").last_consolidated == 4
assert latest_messages[-1]["content"] == _ARCHIVE_PROMPT
sessions.invalidate("cli:incremental")
reloaded = sessions.get_or_create("cli:incremental")
assert reloaded.last_archived == 4
assert reloaded.metadata["_last_summary"]["text"] == second
@pytest.mark.asyncio
async def test_raw_fallback_preserves_previous_checkpoint_and_new_chunk(
self,
real_consolidator,
mock_provider,
store,
runtime,
):
mock_provider.chat_with_retry.side_effect = [
LLMResponse(content="Earlier durable checkpoint.", finish_reason="stop"),
RuntimeError("LLM unavailable"),
]
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:cumulative-fallback")
session.add_message("user", "first user")
session.add_message("assistant", "first answer")
sessions.save(session)
await real_consolidator.compact_idle_session(
"cli:cumulative-fallback",
runtime=runtime,
)
current = sessions.get_or_create("cli:cumulative-fallback")
current.add_message("user", "second user")
current.add_message("assistant", "newest working state")
sessions.save(current)
fallback = await real_consolidator.compact_idle_session(
"cli:cumulative-fallback",
runtime=runtime,
)
assert fallback is not None
assert "[Previous archived context]" in fallback
assert "Earlier durable checkpoint." in fallback
assert "[Newly archived raw context]" in fallback
assert "newest working state" in fallback
entries = store.read_unprocessed_history(0)
assert entries[0]["content"] == "Earlier durable checkpoint."
assert entries[1]["content"].startswith("[RAW] 2 messages")
sessions.invalidate("cli:cumulative-fallback")
reloaded = sessions.get_or_create("cli:cumulative-fallback")
assert reloaded.metadata["_last_summary"]["text"] == fallback
@pytest.mark.asyncio
async def test_nothing_keeps_previous_replacement_checkpoint(
self,
real_consolidator,
mock_provider,
runtime,
):
mock_provider.chat_with_retry.side_effect = [
LLMResponse(content="Existing checkpoint.", finish_reason="stop"),
LLMResponse(content="(nothing)", finish_reason="stop"),
]
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:nothing-after-summary")
session.add_message("user", "important first turn")
session.add_message("assistant", "important result")
sessions.save(session)
await real_consolidator.compact_idle_session(
"cli:nothing-after-summary",
runtime=runtime,
)
current = sessions.get_or_create("cli:nothing-after-summary")
current.add_message("user", "thanks")
current.add_message("assistant", "you're welcome")
sessions.save(current)
result = await real_consolidator.compact_idle_session(
"cli:nothing-after-summary",
runtime=runtime,
)
assert result == "(nothing)"
sessions.invalidate("cli:nothing-after-summary")
reloaded = sessions.get_or_create("cli:nothing-after-summary")
assert reloaded.last_archived == 4
assert reloaded.metadata["_last_summary"]["text"] == "Existing checkpoint."
@pytest.mark.asyncio
async def test_concurrent_append_remains_unarchived(
@@ -664,13 +824,13 @@ class TestCompactIdleSession:
reloaded = sessions.get_or_create("cli:concurrent")
assert len(reloaded.messages) == 4
assert reloaded.last_consolidated == 2
assert reloaded.last_archived == 2
@pytest.mark.asyncio
async def test_summarizes_retained_suffix_not_just_dropped_prefix(
self, real_consolidator, mock_provider, runtime
):
"""idleCompact must summarize over the full unconsolidated tail, including
"""idleCompact must summarize over the full unarchived tail, including
the recent suffix it retains. Otherwise a late user correction / final
result that lands in the kept suffix is excluded from the persisted
summary, leaving a stale wrong conclusion in history. Regression for #4264."""
@@ -705,6 +865,7 @@ class TestCompactIdleSession:
mock_provider.chat_with_retry.side_effect = RuntimeError("LLM unavailable")
sessions = real_consolidator.sessions
session = sessions.get_or_create("cli:rawdrop")
session.provider_state = _provider_state()
for i in range(18):
session.add_message("user", f"user msg {i}")
session.add_message("assistant", f"assistant msg {i}")
@@ -723,6 +884,7 @@ class TestCompactIdleSession:
reloaded = sessions.get_or_create("cli:rawdrop")
assert len(reloaded.messages) == 38
assert reloaded.messages[-1]["content"] == "RETAINED_SUFFIX_marker"
assert reloaded.provider_state == _provider_state()
@pytest.mark.asyncio
async def test_idle_compact_writes_session_key_to_history(
@@ -788,11 +950,16 @@ class TestCompactIdleSession:
result = await real_consolidator.compact_idle_session(
"cli:nothing", runtime=runtime, max_suffix=4
)
second = await real_consolidator.compact_idle_session(
"cli:nothing", runtime=runtime, max_suffix=4
)
assert result == "(nothing)"
assert second == ""
reloaded = sessions.get_or_create("cli:nothing")
assert "_last_summary" not in reloaded.metadata
assert real_consolidator.store.read_unprocessed_history(0) == []
mock_provider.chat_with_retry.assert_awaited_once()
@pytest.mark.asyncio
async def test_llm_failure_preserves_history_but_advances_replay_boundary(
@@ -809,7 +976,8 @@ class TestCompactIdleSession:
result = await real_consolidator.compact_idle_session(
"cli:fail", runtime=runtime, max_suffix=4
)
assert result is None
assert result is not None
assert "[RAW]" in result
# raw_archive should have been called (history.jsonl gets an entry)
entries = store.read_unprocessed_history(since_cursor=0)
@@ -818,7 +986,8 @@ class TestCompactIdleSession:
reloaded = sessions.get_or_create("cli:fail")
assert len(reloaded.messages) == 20
assert reloaded.messages[0]["content"] == "u0"
assert reloaded.last_consolidated == 20
assert reloaded.last_archived == 20
assert reloaded.metadata["_last_summary"]["text"] == result
assert [m["content"] for m in reloaded.get_history(max_messages=20)] == [
"u6",
"a6",
@@ -831,10 +1000,10 @@ class TestCompactIdleSession:
]
@pytest.mark.asyncio
async def test_respects_last_consolidated(
async def test_respects_last_archived(
self, real_consolidator, mock_provider, runtime
):
"""30 turns with last_consolidated=50 → only unconsolidated tail considered."""
"""30 turns with last_archived=50 → only the unarchived tail is considered."""
mock_provider.chat_with_retry.return_value = MagicMock(
content="Tail summary.", finish_reason="stop"
)
@@ -843,7 +1012,7 @@ class TestCompactIdleSession:
for i in range(30):
session.add_message("user", f"u{i}")
session.add_message("assistant", f"a{i}")
session.last_consolidated = 50 # Only 10 messages unconsolidated
session.last_archived = 50 # Only 10 messages remain unarchived
sessions.save(session)
result = await real_consolidator.compact_idle_session(
@@ -852,18 +1021,17 @@ class TestCompactIdleSession:
assert result == "Tail summary."
reloaded = sessions.get_or_create("cli:offset")
assert len(reloaded.messages) == 60
assert reloaded.last_consolidated == 60
assert reloaded.last_archived == 60
# Verify only the unconsolidated tail was processed:
# All 10 unconsolidated messages (50-59) are archived exactly once.
# Verify only the unarchived tail was processed:
# All 10 unarchived messages (50-59) are archived exactly once.
archived_call = mock_provider.chat_with_retry.call_args
sent_messages = archived_call.kwargs["messages"]
sent_content = [message.get("content") for message in sent_messages]
# The ordinary replay prefix contributes recent context, while the
# temporary instruction limits the new overview to the unarchived tail.
# The replacement overview covers all model-visible conversation context.
assert "u0" not in sent_content
assert "u26" in sent_content
assert "final 10 conversation messages" in sent_messages[-1]["content"]
assert sent_messages[-1]["content"] == _ARCHIVE_PROMPT
@pytest.mark.asyncio
async def test_full_archive_keeps_extended_legal_replay_suffix(
@@ -890,7 +1058,7 @@ class TestCompactIdleSession:
reloaded = sessions.get_or_create("cli:noncontiguous")
assert len(reloaded.messages) == 25
assert reloaded.last_consolidated == 25
assert reloaded.last_archived == 25
assert [m["content"] for m in reloaded.get_history(max_messages=25)] == [
"user-14",
"assistant-00",
@@ -905,7 +1073,7 @@ class TestCompactIdleSession:
"assistant-09",
]
# #4264: idle compaction now summarizes the full unconsolidated tail, so
# #4264: idle compaction now summarizes the full unarchived tail, so
# the dropped head (user-00) and retained suffix (user-14 through
# assistant-09) are all summarized.
archived_call = mock_provider.chat_with_retry.call_args
@@ -923,7 +1091,7 @@ class TestCompactIdleSession:
runtime,
):
tools = [{"type": "function", "function": {"name": "lookup"}}]
real_consolidator._get_tool_definitions.return_value = tools
real_consolidator.archiver._get_tool_definitions.return_value = tools
mock_provider.chat_with_retry.return_value = LLMResponse(
content="Overview from the temporary turn.",
finish_reason="stop",
@@ -952,9 +1120,9 @@ class TestCompactIdleSession:
"user",
]
assert sent_messages[2]["tool_calls"][0]["id"] == "call-1"
assert "final 4 conversation messages" in sent_messages[-1]["content"]
assert sent_messages[-1]["content"] == _ARCHIVE_PROMPT
assert call["tools"] == tools
assert call["tool_choice"] == "none"
assert "tool_choice" not in call
reloaded = sessions.get_or_create("cli:tool-history")
assert len(reloaded.messages) == 4
@@ -992,12 +1160,13 @@ class TestCompactIdleSession:
runtime=runtime,
)
assert result is None
assert result is not None
assert "[RAW]" in result
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1
assert entries[0]["content"].startswith("[RAW] ")
assert "important answer" in entries[0]["content"]
assert sessions.get_or_create("cli:unexpected-tool").last_consolidated == 2
assert sessions.get_or_create("cli:unexpected-tool").last_archived == 2
@pytest.mark.asyncio
async def test_empty_response_uses_raw_fallback(
@@ -1022,12 +1191,13 @@ class TestCompactIdleSession:
runtime=runtime,
)
assert result is None
assert result is not None
assert "[RAW]" in result
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1
assert entries[0]["content"].startswith("[RAW] ")
assert "important answer" in entries[0]["content"]
assert sessions.get_or_create("cli:empty-summary").last_consolidated == 2
assert sessions.get_or_create("cli:empty-summary").last_archived == 2
@pytest.mark.asyncio
async def test_oversized_prefix_raw_archives_without_flattened_llm_retry(
@@ -1048,15 +1218,16 @@ class TestCompactIdleSession:
runtime=runtime,
)
assert result is None
assert result is not None
assert "[RAW]" in result
mock_provider.chat_with_retry.assert_not_awaited()
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 1
assert entries[0]["content"].startswith("[RAW] ")
assert sessions.get_or_create("sdk:oversized").last_consolidated == 1
assert sessions.get_or_create("sdk:oversized").last_archived == 1
@pytest.mark.asyncio
async def test_incremental_scope_counts_only_model_visible_messages(
async def test_archive_context_contains_only_model_visible_messages(
self,
real_consolidator,
mock_provider,
@@ -1070,7 +1241,7 @@ class TestCompactIdleSession:
session = sessions.get_or_create("cli:commands")
session.add_message("user", "already archived user")
session.add_message("assistant", "already archived answer")
session.last_consolidated = 2
session.last_archived = 2
session.add_message("user", "/status", _command=True)
session.add_message("assistant", "status output", _command=True)
session.add_message("user", "new user")
@@ -1089,7 +1260,7 @@ class TestCompactIdleSession:
"new user",
"new answer",
]
assert "final 2 conversation messages" in sent[-1]["content"]
assert sent[-1]["content"] == _ARCHIVE_PROMPT
@pytest.mark.asyncio
async def test_reuses_real_prefix_for_unified_session_workspace(
@@ -1122,8 +1293,6 @@ class TestCompactIdleSession:
current_message="next project question",
channel="websocket",
workspace=project,
session_key=session.key,
unified_session=True,
)
await loop.consolidator.compact_idle_session(
@@ -1133,7 +1302,7 @@ class TestCompactIdleSession:
sent_messages = runtime.provider.chat_with_retry.call_args.kwargs["messages"]
assert sent_messages[:-1] == ordinary_messages[:-1]
assert "final 2 conversation messages" in sent_messages[-1]["content"]
assert sent_messages[-1]["content"] == _ARCHIVE_PROMPT
system = sent_messages[0]["content"]
assert "PROJECT_WORKSPACE_MARKER" in system
assert "GLOBAL_WORKSPACE_MARKER" not in system
@@ -1278,7 +1447,7 @@ class TestConsolidatorSessionRefresh:
session_after = sessions.get_or_create("cli:test")
assert len(session_after.messages) == 40
assert session_after.last_consolidated == 40
assert session_after.last_archived == 40
assert len(session_after.get_history(max_messages=40)) == 8
@@ -1303,6 +1472,21 @@ class TestRawArchiveTruncation:
assert len(entries) == 1
assert "hello" in entries[0]["content"]
def test_raw_archive_returns_the_sanitized_persisted_checkpoint(self, store):
messages = [
{
"role": "user",
"content": "<think>PRIVATE_REASONING</think>visible result",
}
]
checkpoint = store.raw_archive(messages, session_key="cli:test")
persisted = store.read_unprocessed_history(since_cursor=0)[0]["content"]
assert checkpoint == persisted
assert "PRIVATE_REASONING" not in checkpoint
assert "visible result" in checkpoint
def test_raw_archive_excludes_model_only_runtime_context(self, store):
content, marker = append_runtime_context(
"ship the feature",
@@ -1334,21 +1518,40 @@ class TestRawArchiveTruncation:
class TestArchivePersistence:
async def test_oversized_summary_is_capped_before_append(
async def test_archive_returns_the_sanitized_persisted_summary(
self, consolidator, mock_provider, store, runtime
):
mock_provider.chat_with_retry.return_value = MagicMock(
content="<think>PRIVATE_REASONING</think>safe summary",
finish_reason="stop",
has_tool_calls=False,
)
summary = await _archive(
consolidator,
[{"role": "user", "content": "hi"}],
runtime,
)
persisted = store.read_unprocessed_history(since_cursor=0)[0]["content"]
assert summary == persisted == "safe summary"
async def test_oversized_summary_uses_history_emergency_cap(
self, consolidator, mock_provider, store, runtime
):
"""A pathologically large LLM summary must not land full-length in
history.jsonl that would re-open the #3412 bloat vector from the
*success* path instead of the fallback path."""
mock_provider.chat_with_retry.return_value = MagicMock(
content="S" * (_ARCHIVE_SUMMARY_MAX_CHARS * 10),
content="S" * (_HISTORY_ENTRY_HARD_CAP * 2),
finish_reason="stop",
)
await _archive(
summary = await _archive(
consolidator,
[{"role": "user", "content": "hi"}],
runtime,
)
entry = store.read_unprocessed_history(since_cursor=0)[0]
assert len(entry["content"]) <= _ARCHIVE_SUMMARY_MAX_CHARS + 50
assert len(entry["content"]) <= _HISTORY_ENTRY_HARD_CAP + 50
assert summary == entry["content"]
+26 -9
View File
@@ -4,7 +4,7 @@ from pathlib import Path
import pytest
from nanobot.agent.context import ContextBuilder
from nanobot.agent.context import ContextBuilder, TranscriptInput
from nanobot.runtime_context import RuntimeContextBlock
# ---------------------------------------------------------------------------
@@ -133,10 +133,7 @@ class TestLoadBootstrapFiles:
(project / "SOUL.md").write_text("project soul collision", encoding="utf-8")
(project / "USER.md").write_text("project user collision", encoding="utf-8")
result = ContextBuilder(agent_home).build_system_prompt(
workspace=project,
include_memory_recent_history=False,
)
result = ContextBuilder(agent_home).build_system_prompt(workspace=project)
assert "selected project rules" in result
assert "global project rules" not in result
@@ -152,10 +149,7 @@ class TestLoadBootstrapFiles:
project.mkdir()
(agent_home / "AGENTS.md").write_text("default workspace rules", encoding="utf-8")
result = ContextBuilder(agent_home).build_system_prompt(
workspace=project,
include_memory_recent_history=False,
)
result = ContextBuilder(agent_home).build_system_prompt(workspace=project)
assert "default workspace rules" not in result
@@ -403,6 +397,15 @@ class TestBuildMessages:
assert "user-only runtime context" not in messages[-1]["content"]
assert "_meta" not in messages[-1]
def test_compatibility_builder_merges_system_role_without_history(self, tmp_path):
builder = _builder(tmp_path)
messages = builder.build_messages([], "system event", current_role="system")
assert len(messages) == 1
assert messages[0]["role"] == "system"
assert str(messages[0]["content"]).endswith("system event")
def test_explicit_skill_reference_loads_full_instructions_for_this_turn(self, tmp_path):
skill_dir = tmp_path / "skills" / "review"
skill_dir.mkdir(parents=True)
@@ -472,6 +475,20 @@ class TestBuildMessages:
assert "previous user message" in str(messages[1]["content"])
assert "new message" in str(messages[1]["content"])
def test_structured_transcript_preserves_fresh_turn_boundary(self, tmp_path):
builder = _builder(tmp_path)
transcript = TranscriptInput(
history=[{"role": "user", "content": "previous user message"}],
current_message="new message",
)
messages = builder.build_transcript(transcript)
assert [message["role"] for message in messages] == ["system", "user", "user"]
assert messages[-2]["content"] == "previous user message"
assert messages[-1]["content"] == "new message"
assert transcript.message_count == 3
def test_current_message_can_be_built_without_history_merge(self, tmp_path):
builder = _builder(tmp_path)
current = builder.build_current_message(
-168
View File
@@ -3,7 +3,6 @@
from __future__ import annotations
import datetime as datetime_module
import re
from datetime import datetime as real_datetime
from importlib.resources import files as pkg_files
from pathlib import Path
@@ -104,173 +103,6 @@ def test_provider_context_appended_after_user_content(tmp_path) -> None:
assert user_pos < context_pos, "user content must precede provider context"
def test_unprocessed_history_injected_into_system_prompt(tmp_path) -> None:
"""Entries in history.jsonl not yet consumed by Dream appear with timestamps."""
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
builder.memory.append_history("User asked about weather in Tokyo")
builder.memory.append_history("Agent fetched forecast via web_search")
prompt = builder.build_system_prompt()
assert "# Recent History" in prompt
assert "User asked about weather in Tokyo" in prompt
assert "Agent fetched forecast via web_search" in prompt
assert re.search(r"\[\d{4}-\d{2}-\d{2} \d{2}:\d{2}\]", prompt)
def test_recent_history_injection_is_session_scoped(tmp_path) -> None:
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
builder.memory.append_history("legacy entry without session")
builder.memory.append_history("telegram history", session_key="telegram:chat-1")
builder.memory.append_history("slack history", session_key="slack:chat-2")
prompt = builder.build_system_prompt(session_key="telegram:chat-1")
assert "# Recent History" in prompt
assert "telegram history" in prompt
assert "slack history" not in prompt
assert "legacy entry without session" not in prompt
def test_session_summary_replaces_interleaved_recent_history_entry(tmp_path) -> None:
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
session_key = "unified:default"
overview = "CURRENT_SESSION_OVERVIEW_MARKER"
builder.memory.append_history("another session event", session_key=session_key)
builder.memory.append_history(overview, session_key=session_key)
latest_cursor = builder.memory.append_history(
"later telegram event",
session_key="telegram:chat-1",
)
summary = {"text": overview, "last_active": "2026-08-19T10:00:00"}
prompt = builder.build_system_prompt(
session_key=session_key,
session_summary=summary,
unified_session=True,
)
assert "# Recent History" in prompt
assert "another session event" in prompt
assert "later telegram event" in prompt
assert "[Archived Context Summary]" in prompt
assert prompt.count(overview) == 1
builder.memory.set_last_dream_cursor(latest_cursor)
processed_prompt = builder.build_system_prompt(
session_key=session_key,
session_summary=summary,
unified_session=True,
)
assert "# Recent History" not in processed_prompt
assert processed_prompt.count(overview) == 1
def test_recent_history_injection_unified_excludes_cron_internals(tmp_path) -> None:
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
builder.memory.append_history("unified user history", session_key="unified:default")
builder.memory.append_history("channel user history", session_key="telegram:chat-1")
builder.memory.append_history("cron internal history", session_key="cron:job-1")
prompt = builder.build_system_prompt(
session_key="unified:default",
unified_session=True,
)
assert "unified user history" in prompt
assert "channel user history" in prompt
assert "cron internal history" not in prompt
def test_cron_recent_history_can_see_own_history_and_unified_context(tmp_path) -> None:
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
builder.memory.append_history("unified user history", session_key="unified:default")
builder.memory.append_history("own cron history", session_key="cron:job-1")
builder.memory.append_history("other cron history", session_key="cron:job-2")
prompt = builder.build_system_prompt(
session_key="cron:job-1",
unified_session=True,
)
assert "unified user history" in prompt
assert "own cron history" in prompt
assert "other cron history" not in prompt
def test_recent_history_capped_at_max(tmp_path) -> None:
"""Only the most recent _MAX_RECENT_HISTORY entries are injected."""
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
for i in range(builder._MAX_RECENT_HISTORY + 20):
builder.memory.append_history(f"entry-{i}")
prompt = builder.build_system_prompt()
assert "entry-0" not in prompt
assert "entry-19" not in prompt
assert f"entry-{builder._MAX_RECENT_HISTORY + 19}" in prompt
def test_recent_history_truncated_at_max_tokens(tmp_path) -> None:
"""Recent History section must be truncated to _MAX_HISTORY_TOKENS."""
import tiktoken
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
big_entry = "word " * (builder._MAX_HISTORY_TOKENS + 5_000)
builder.memory.append_history(big_entry)
prompt = builder.build_system_prompt()
history_section = prompt.split("# Recent History\n\n", 1)
assert len(history_section) == 2
enc = tiktoken.get_encoding("cl100k_base")
assert len(enc.encode(history_section[1])) <= builder._MAX_HISTORY_TOKENS
def test_no_recent_history_when_dream_has_processed_all(tmp_path) -> None:
"""If Dream has consumed everything, no Recent History section should appear."""
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
cursor = builder.memory.append_history("already processed entry")
builder.memory.set_last_dream_cursor(cursor)
prompt = builder.build_system_prompt()
assert "# Recent History" not in prompt
def test_partial_dream_processing_shows_only_remainder(tmp_path) -> None:
"""When Dream has processed some entries, only the unprocessed ones appear."""
workspace = _make_workspace(tmp_path)
builder = ContextBuilder(workspace)
builder.memory.append_history("old conversation about Python")
c2 = builder.memory.append_history("old conversation about Rust")
builder.memory.append_history("recent question about Docker")
builder.memory.append_history("recent question about K8s")
builder.memory.set_last_dream_cursor(c2)
prompt = builder.build_system_prompt()
assert "# Recent History" in prompt
assert "old conversation about Python" not in prompt
assert "old conversation about Rust" not in prompt
assert "recent question about Docker" in prompt
assert "recent question about K8s" in prompt
def test_execution_rules_in_system_prompt(tmp_path) -> None:
"""Execution rules should appear in the system prompt via the default templates."""
from nanobot.utils.helpers import sync_workspace_templates
+3 -3
View File
@@ -426,7 +426,7 @@ class TestEphemeralDirect:
bus=bus,
provider=provider,
workspace=tmp_path,
context_window_tokens=8000,
context_window_tokens=32_000,
)
return loop, store
@@ -606,7 +606,7 @@ class TestEphemeralDirect:
bus=MessageBus(),
provider=provider,
workspace=tmp_path,
context_window_tokens=8000,
context_window_tokens=32_000,
)
await loop.process_direct(
@@ -666,7 +666,7 @@ class TestEphemeralHooks:
bus=bus,
provider=provider,
workspace=tmp_path,
context_window_tokens=8000,
context_window_tokens=32_000,
hooks=[spy],
)
+9 -5
View File
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from nanobot.agent.context import TranscriptInput
from nanobot.agent.hook import (
AgentHook,
AgentHookContext,
@@ -459,7 +460,7 @@ async def test_agent_loop_extra_hook_receives_calls(tmp_path):
loop.tools.get_definitions = MagicMock(return_value=[])
result = await loop._run_agent_loop(
[{"role": "user", "content": "hi"}],
TranscriptInput(history=[{"role": "user", "content": "hi"}], current_message=None),
runtime=loop.llm_runtime(),
)
@@ -504,7 +505,7 @@ async def test_agent_loop_turn_hook_factories_receive_context(tmp_path):
runtime = loop.llm_runtime()
await loop._run_agent_loop(
[{"role": "user", "content": "hi"}],
TranscriptInput(history=[{"role": "user", "content": "hi"}], current_message=None),
runtime=runtime,
on_progress=on_progress,
request_context=RequestContext(
@@ -551,7 +552,7 @@ async def test_agent_loop_extra_hook_error_isolation(tmp_path):
loop.tools.get_definitions = MagicMock(return_value=[])
result = await loop._run_agent_loop(
[{"role": "user", "content": "hi"}],
TranscriptInput(history=[{"role": "user", "content": "hi"}], current_message=None),
runtime=loop.llm_runtime(),
)
@@ -577,7 +578,9 @@ async def test_agent_loop_extra_hooks_do_not_swallow_loop_hook_errors(tmp_path):
with pytest.raises(RuntimeError, match="progress failed"):
await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=bad_progress
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_progress=bad_progress,
)
@@ -596,7 +599,8 @@ async def test_agent_loop_no_hooks_backward_compat(tmp_path):
loop.max_iterations = 2
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime()
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
)
assert result.final_content == (
"I reached the maximum number of tool call iterations (2) "
+47
View File
@@ -0,0 +1,47 @@
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
def _provider() -> MagicMock:
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation = SimpleNamespace(
max_tokens=4096,
temperature=0.1,
reasoning_effort=None,
)
return provider
def test_request_concurrency_is_unlimited_by_default(
monkeypatch: pytest.MonkeyPatch,
loop_factory,
) -> None:
monkeypatch.delenv("NANOBOT_MAX_CONCURRENT_REQUESTS", raising=False)
loop = loop_factory(provider=_provider(), patch_deps=True)
assert loop._concurrency_gate is None
@pytest.mark.asyncio
async def test_positive_request_concurrency_keeps_explicit_cap(
monkeypatch: pytest.MonkeyPatch,
loop_factory,
) -> None:
monkeypatch.setenv("NANOBOT_MAX_CONCURRENT_REQUESTS", "2")
loop = loop_factory(provider=_provider(), patch_deps=True)
gate = loop._concurrency_gate
assert gate is not None
for _ in range(2):
await gate.acquire()
try:
assert gate.locked()
finally:
for _ in range(2):
gate.release()
+59 -114
View File
@@ -2,17 +2,22 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
import nanobot.agent.memory as memory_module
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import LLMResponse
def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -> AgentLoop:
def _make_loop(
tmp_path,
*,
estimated_tokens: int,
context_window_tokens: int,
max_tokens: int = 0,
) -> AgentLoop:
from nanobot.providers.base import GenerationSettings
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.generation = GenerationSettings(max_tokens=0)
provider.generation = GenerationSettings(max_tokens=max_tokens)
provider.estimate_prompt_tokens.return_value = (estimated_tokens, "test-counter")
_response = LLMResponse(content="ok", tool_calls=[])
provider.chat_with_retry = AsyncMock(return_value=_response)
@@ -24,6 +29,9 @@ def _make_loop(tmp_path, *, estimated_tokens: int, context_window_tokens: int) -
workspace=tmp_path,
model="test-model",
context_window_tokens=context_window_tokens,
# These tests isolate Memory consolidation; Runner request fitting is
# covered separately with realistic context windows.
context_block_limit=10_000,
)
loop.tools.get_definitions = MagicMock(return_value=[])
loop.consolidator._SAFETY_BUFFER = 0
@@ -41,17 +49,16 @@ async def test_prompt_below_threshold_does_not_consolidate(tmp_path) -> None:
@pytest.mark.asyncio
async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypatch) -> None:
async def test_prompt_above_threshold_triggers_consolidation(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
loop.sessions.save(session)
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _message: 500)
await loop.process_direct("hello", session_key="cli:test")
@@ -59,23 +66,46 @@ async def test_prompt_above_threshold_triggers_consolidation(tmp_path, monkeypat
@pytest.mark.asyncio
async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path, monkeypatch) -> None:
async def test_token_consolidation_refreshes_summary_for_current_request(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock( # type: ignore[method-assign]
return_value="FRESH_CHECKPOINT"
)
loop.consolidator.estimate_session_prompt_tokens = MagicMock( # type: ignore[method-assign]
return_value=(1000, "test")
)
loop.schedule_background = lambda coro: coro.close() # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
loop.sessions.save(session)
await loop.process_direct("hello", session_key="cli:test")
request_messages = loop.provider.chat_with_retry.await_args.kwargs["messages"]
system_prompt = request_messages[0]["content"]
assert "FRESH_CHECKPOINT" in system_prompt
assert all(message.get("content") != "u0" for message in request_messages)
assert loop.sessions.get_or_create("cli:test").last_archived == 12
@pytest.mark.asyncio
async def test_prompt_above_threshold_uses_fixed_recent_tail(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=1000, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
loop.sessions.save(session)
token_map = {"u1": 120, "a1": 120, "u2": 120, "a2": 120, "u3": 120}
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda message: token_map[message["content"]])
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
@@ -83,112 +113,29 @@ async def test_prompt_above_threshold_archives_until_next_user_boundary(tmp_path
archive_end = loop.consolidator.archive_session.await_args.kwargs["archive_end"]
archived_chunk = session.messages[:archive_end]
assert [message["content"] for message in archived_chunk] == ["u1", "a1", "u2", "a2"]
assert session.last_consolidated == 4
@pytest.mark.asyncio
async def test_consolidation_loops_until_target_met(tmp_path, monkeypatch) -> None:
"""Verify maybe_consolidate_by_tokens keeps looping until under threshold."""
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
{"role": "assistant", "content": "a3", "timestamp": "2026-01-01T00:00:05"},
{"role": "user", "content": "u4", "timestamp": "2026-01-01T00:00:06"},
assert [message["content"] for message in archived_chunk] == [
"u0", "a0", "u1", "a1", "u2", "a2", "u3", "a3", "u4", "a4", "u5", "a5",
]
loop.sessions.save(session)
call_count = [0]
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
if call_count[0] == 2:
return (300, "test")
return (80, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
assert loop.consolidator.archive_session.await_count == 2
assert session.last_consolidated == 6
assert session.last_archived == 12
@pytest.mark.asyncio
async def test_consolidation_continues_below_trigger_until_half_target(tmp_path, monkeypatch) -> None:
"""Once triggered, consolidation should continue until it drops below half threshold."""
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value=True) # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
{"role": "assistant", "content": "a2", "timestamp": "2026-01-01T00:00:03"},
{"role": "user", "content": "u3", "timestamp": "2026-01-01T00:00:04"},
{"role": "assistant", "content": "a3", "timestamp": "2026-01-01T00:00:05"},
{"role": "user", "content": "u4", "timestamp": "2026-01-01T00:00:06"},
]
loop.sessions.save(session)
call_count = [0]
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
if call_count[0] == 2:
return (150, "test")
return (80, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 100)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
runtime=loop.llm_runtime(),
)
assert loop.consolidator.archive_session.await_count == 2
assert session.last_consolidated == 6
@pytest.mark.asyncio
async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path, monkeypatch) -> None:
async def test_consolidation_persists_summary_for_next_prepare_session(tmp_path) -> None:
loop = _make_loop(tmp_path, estimated_tokens=0, context_window_tokens=200)
loop.consolidator.archive_session = AsyncMock(return_value="User discussed project status.") # type: ignore[method-assign]
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(5)
for role in ("user", "assistant")
]
loop.sessions.save(session)
call_count = [0]
def mock_estimate(_session, *, runtime):
call_count[0] += 1
if call_count[0] == 1:
return (500, "test")
return (80, "test")
return (500, "test")
loop.consolidator.estimate_session_prompt_tokens = mock_estimate # type: ignore[method-assign]
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 150)
await loop.consolidator.maybe_consolidate_by_tokens(
session,
@@ -235,7 +182,7 @@ async def test_preflight_consolidation_receives_pending_summary(tmp_path) -> Non
@pytest.mark.asyncio
async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) -> None:
async def test_preflight_consolidation_before_llm_call(tmp_path) -> None:
"""Verify preflight consolidation runs before the LLM call in process_direct."""
order: list[str] = []
@@ -258,13 +205,11 @@ async def test_preflight_consolidation_before_llm_call(tmp_path, monkeypatch) ->
session = loop.sessions.get_or_create("cli:test")
session.messages = [
{"role": "user", "content": "u1", "timestamp": "2026-01-01T00:00:00"},
{"role": "assistant", "content": "a1", "timestamp": "2026-01-01T00:00:01"},
{"role": "user", "content": "u2", "timestamp": "2026-01-01T00:00:02"},
{"role": role, "content": f"{role[0]}{turn}"}
for turn in range(10)
for role in ("user", "assistant")
]
loop.sessions.save(session)
monkeypatch.setattr(memory_module, "estimate_message_tokens", lambda _m: 500)
call_count = [0]
def mock_estimate(_session, *, runtime):
call_count[0] += 1
+16 -17
View File
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.agent.context import TranscriptInput
from nanobot.agent.hooks import create_file_edit_activity_hook
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.context import current_request_context
@@ -84,7 +85,9 @@ class TestToolEventProgress:
progress.append((content, tool_hint, tool_events))
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_progress=on_progress,
)
assert result.final_content == "Done"
@@ -155,7 +158,9 @@ class TestToolEventProgress:
file_events.extend(file_edit_events)
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_progress=on_progress,
)
assert result.final_content == "Done"
@@ -225,7 +230,9 @@ class TestToolEventProgress:
)
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_progress=on_progress,
)
assert result.final_content == "Done"
@@ -263,7 +270,9 @@ class TestToolEventProgress:
file_events.extend(file_edit_events)
await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_progress=on_progress,
)
assert file_events == []
@@ -373,7 +382,6 @@ class TestToolEventProgress:
"""The /goal command rewrites the prompt but must not bypass WebUI file-edit progress."""
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "test-model"
call_count = 0
@@ -460,7 +468,6 @@ class TestToolEventProgress:
"""Non-streaming channels should get one final reply, not token progress spam."""
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "openai-codex/gpt-5.5"
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="Hello", tool_calls=[]))
provider.chat_stream_with_retry = AsyncMock()
@@ -493,7 +500,6 @@ class TestToolEventProgress:
"""Streaming channels still receive provider deltas through stream events."""
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "openai-codex/gpt-5.5"
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
@@ -544,7 +550,6 @@ class TestToolEventProgress:
) -> None:
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "test-model"
responses = iter([
LLMResponse(content="first-", finish_reason="length"),
@@ -590,7 +595,6 @@ class TestToolEventProgress:
) -> None:
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "test-model"
call_count = 0
@@ -637,7 +641,6 @@ class TestToolEventProgress:
) -> None:
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "test-model"
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
@@ -728,7 +731,6 @@ class TestToolEventProgress:
"""A no-tools finalization must not be dropped after empty stream retries."""
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "openai-codex/gpt-5.5"
provider.chat_stream_with_retry = AsyncMock(side_effect=[
LLMResponse(content=None, tool_calls=[]),
@@ -776,7 +778,6 @@ class TestToolEventProgress:
) -> None:
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "openai-codex/gpt-5.5"
first_request_started = asyncio.Event()
release_first_request = asyncio.Event()
@@ -935,7 +936,6 @@ class TestToolEventProgress:
"""Recovered streaming output should use a new stream segment."""
bus = MessageBus()
provider = MagicMock()
provider.supports_progress_deltas = True
provider.get_default_model.return_value = "openai-codex/gpt-5.5"
async def chat_stream_with_retry(*, on_content_delta, on_stream_recover, **kwargs):
@@ -988,13 +988,12 @@ class TestToolEventProgress:
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_streamed_progress_is_not_repeated_before_tool_execution(
async def test_streamed_content_is_not_repeated_before_tool_execution(
self,
tmp_path: Path,
) -> None:
"""If content was already streamed as progress, tool setup should not repeat it."""
"""If content was already streamed, tool setup should not repeat it."""
loop = _make_loop(tmp_path)
loop.provider.supports_progress_deltas = True
tool_call = ToolCallRequest(id="call1", name="custom_tool", arguments={"path": "foo.txt"})
calls = iter([
LLMResponse(content="I will inspect it.", tool_calls=[tool_call]),
@@ -1029,7 +1028,7 @@ class TestToolEventProgress:
progress.append((content, tool_hint, tool_events))
result = await loop._run_agent_loop(
[],
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_progress=on_progress,
on_stream=on_stream,
+14 -7
View File
@@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from nanobot.agent.context import TranscriptInput
from nanobot.agent.goal_permission import goal_mutation_allowed, goal_mutation_permission
from nanobot.agent.tools.context import RequestContext
from nanobot.bus.outbound_events import StreamedResponseEvent
@@ -55,7 +56,7 @@ async def test_ephemeral_runner_enters_and_restores_turn_scopes(tmp_path):
loop.tools.get_definitions = MagicMock(return_value=[])
await loop._run_agent_loop(
[],
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
ephemeral=True,
turn_scopes=[goal_mutation_permission(True)],
@@ -340,7 +341,8 @@ async def test_loop_max_iterations_message_stays_stable(tmp_path):
loop.max_iterations = 2
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime()
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
)
assert result.final_content == (
@@ -362,7 +364,7 @@ async def test_loop_goal_turn_uses_standard_iteration_budget(tmp_path):
runtime = loop.llm_runtime()
result = await loop._run_agent_loop(
[],
TranscriptInput(history=[], current_message=None),
runtime=runtime,
request_context=RequestContext(
channel="cli",
@@ -401,7 +403,7 @@ async def test_loop_stream_filter_handles_think_only_prefix_without_crashing(tmp
endings.append(resuming)
result = await loop._run_agent_loop(
[],
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_stream=on_stream,
on_stream_end=on_stream_end,
@@ -428,7 +430,9 @@ async def test_loop_stream_filter_hides_partial_trailing_think_prefix(tmp_path):
deltas.append(delta)
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_stream=on_stream
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_stream=on_stream,
)
assert result.final_content == "Hello World"
@@ -451,7 +455,9 @@ async def test_loop_stream_filter_hides_complete_trailing_think_tag(tmp_path):
deltas.append(delta)
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_stream=on_stream
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_stream=on_stream,
)
assert result.final_content == "Hello World"
@@ -472,7 +478,8 @@ async def test_loop_retries_think_only_final_response(tmp_path):
loop.provider.chat_with_retry = chat_with_retry
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime()
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
)
assert result.final_content == "Recovered answer"
+46 -23
View File
@@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from loguru import logger
from nanobot.agent.context import ContextBuilder
from nanobot.agent.context import ContextBuilder, TranscriptInput
from nanobot.agent.loop import AgentLoop
from nanobot.agent.runner import AgentRunResult
from nanobot.agent.tools.context import RequestContext, request_context
@@ -79,6 +79,13 @@ def _agent_run_result(
)
def _assembled_messages(
builder: ContextBuilder,
transcript_input: TranscriptInput,
) -> list[dict]:
return builder.build_transcript(transcript_input, include_memory=False)
def _mk_loop() -> AgentLoop:
loop = AgentLoop.__new__(AgentLoop)
from nanobot.config.schema import AgentDefaults
@@ -930,10 +937,13 @@ async def test_runtime_checkpoint_keeps_provider_state_out_of_public_metadata(
session = loop.sessions.get_or_create("cli:private-checkpoint")
await loop._run_agent_loop(
[
{"role": "system", "content": "system"},
{"role": "user", "content": "question"},
],
TranscriptInput(
history=[
{"role": "system", "content": "system"},
{"role": "user", "content": "question"},
],
current_message=None,
),
runtime=loop.llm_runtime(),
session=session,
)
@@ -1008,7 +1018,7 @@ async def test_subagent_followup_state_is_durable_before_prompt_assembly(
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop.provider.can_resume_conversation_state.return_value = True
loop._build_initial_messages = MagicMock( # type: ignore[method-assign]
loop.context.build_system_prompt = MagicMock( # type: ignore[method-assign]
side_effect=RuntimeError("prompt boom"),
)
session = loop.sessions.get_or_create("cli:subagent-prompt-crash")
@@ -1041,8 +1051,8 @@ async def test_subagent_redelivery_does_not_duplicate_staged_provider_input(
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
loop.provider.can_resume_conversation_state.return_value = True
build_initial_messages = loop._build_initial_messages
loop._build_initial_messages = MagicMock( # type: ignore[method-assign]
build_system_prompt = loop.context.build_system_prompt
loop.context.build_system_prompt = MagicMock( # type: ignore[method-assign]
side_effect=RuntimeError("prompt boom"),
)
session = loop.sessions.get_or_create("cli:subagent-redelivery")
@@ -1066,7 +1076,7 @@ async def test_subagent_redelivery_does_not_duplicate_staged_provider_input(
message.get("content")
for message in persisted.provider_state.pending_messages
].count("subagent result") == 1
loop._build_initial_messages = build_initial_messages # type: ignore[method-assign]
loop.context.build_system_prompt = build_system_prompt # type: ignore[method-assign]
loop._run_agent_loop = AsyncMock( # type: ignore[method-assign]
side_effect=RuntimeError("provider boom"),
)
@@ -1319,7 +1329,8 @@ async def test_internal_continuation_queues_turn_without_fake_user_history(
calls: list[dict] = []
async def fake_run_agent_loop(initial_messages, *, metadata=None, **_kwargs):
async def fake_run_agent_loop(transcript_input, *, metadata=None, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
calls.append({"initial_messages": initial_messages, "metadata": metadata})
if len(calls) == 1:
return _agent_run_result(
@@ -1387,8 +1398,9 @@ async def test_internal_continuation_preserves_streaming_route_metadata(
calls = 0
async def fake_run_agent_loop(initial_messages, *, on_stream=None, on_stream_end=None, **_kwargs):
async def fake_run_agent_loop(transcript_input, *, on_stream=None, on_stream_end=None, **_kwargs):
nonlocal calls
initial_messages = _assembled_messages(loop.context, transcript_input)
calls += 1
if calls == 1:
return _agent_run_result(
@@ -1460,8 +1472,9 @@ async def test_websocket_internal_continuation_keeps_single_visible_run(
calls = 0
async def fake_run_agent_loop(initial_messages, **_kwargs):
async def fake_run_agent_loop(transcript_input, **_kwargs):
nonlocal calls
initial_messages = _assembled_messages(loop.context, transcript_input)
calls += 1
if calls == 1:
return _agent_run_result(
@@ -1623,7 +1636,7 @@ async def test_run_agent_loop_continuation_reads_latest_goal_metadata(
runtime = loop.llm_runtime()
await loop._run_agent_loop(
[],
TranscriptInput(history=[], current_message=None),
runtime=runtime,
session=session,
request_context=RequestContext(
@@ -1753,7 +1766,7 @@ async def test_stop_preserves_runtime_checkpoint_for_next_turn(tmp_path: Path) -
checkpoint_saved = asyncio.Event()
async def interrupted_run_agent_loop(_initial_messages, *, session=None, **_kwargs):
async def interrupted_run_agent_loop(_transcript_input, *, session=None, **_kwargs):
assert session is not None
loop._set_runtime_checkpoint(
session,
@@ -1813,7 +1826,8 @@ async def test_stop_preserves_runtime_checkpoint_for_next_turn(tmp_path: Path) -
assert interrupted.metadata.get(AgentLoop._PENDING_USER_TURN_KEY) is True
assert interrupted.metadata.get(AgentLoop._RUNTIME_CHECKPOINT_KEY) is not None
async def resumed_run_agent_loop(initial_messages, **_kwargs):
async def resumed_run_agent_loop(transcript_input, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
return _agent_run_result(
"next answer",
[*initial_messages, {"role": "assistant", "content": "next answer"}],
@@ -1864,7 +1878,8 @@ async def test_system_subagent_followup_is_persisted_before_prompt_assembly(tmp_
record_runtime = MagicMock(wraps=loop.runtime_event_publisher.record_turn_runtime)
loop.runtime_event_publisher.record_turn_runtime = record_runtime
async def fake_run_agent_loop(initial_messages, **kwargs):
async def fake_run_agent_loop(transcript_input, **kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
seen["initial_messages"] = initial_messages
seen["runtime"] = kwargs["runtime"]
seen["request_context"] = kwargs["request_context"]
@@ -1940,7 +1955,8 @@ async def test_turn_usage_is_persisted_with_the_saved_session(tmp_path: Path) ->
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
turn_usage = LLMUsage.reported(input_tokens=64, output_tokens=9)
async def fake_run_agent_loop(initial_messages, **_kwargs):
async def fake_run_agent_loop(transcript_input, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
return _agent_run_result(
"done",
[*initial_messages, {"role": "assistant", "content": "done"}],
@@ -1966,7 +1982,8 @@ async def test_system_subagent_followup_does_not_log_content(tmp_path: Path) ->
return_value=False
)
async def fake_run_agent_loop(initial_messages, **_kwargs):
async def fake_run_agent_loop(transcript_input, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
return _agent_run_result(
"done",
[*initial_messages, {"role": "assistant", "content": "done"}],
@@ -2022,7 +2039,8 @@ async def test_system_subagent_followup_uses_common_turn_lifecycle(tmp_path: Pat
setattr(loop, name, record)
async def fake_run_agent_loop(initial_messages, **_kwargs):
async def fake_run_agent_loop(transcript_input, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
return _agent_run_result(
"done",
[*initial_messages, {"role": "assistant", "content": "done"}],
@@ -2065,7 +2083,8 @@ async def test_multiple_subagent_followups_all_persist_as_standalone_history(tmp
loop = _make_full_loop(tmp_path)
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
async def fake_run_agent_loop(initial_messages, **_kwargs):
async def fake_run_agent_loop(transcript_input, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
return _agent_run_result(
"ack",
[*initial_messages, {"role": "assistant", "content": "ack"}],
@@ -2196,7 +2215,8 @@ async def test_system_subagent_followup_uses_thread_session_and_slack_metadata(t
seen: dict[str, object] = {}
async def fake_run_agent_loop(initial_messages, **kwargs):
async def fake_run_agent_loop(transcript_input, **kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
seen["initial_messages"] = initial_messages
seen["request_context"] = kwargs["request_context"]
return _agent_run_result(
@@ -2252,8 +2272,11 @@ async def test_turn_after_unanswered_user_keeps_tool_call_pairing(tmp_path: Path
session.add_message("user", "earlier question that never got an answer")
loop.sessions.save(session)
async def fake_run_agent_loop(initial_messages, **_kwargs):
assert [m["role"] for m in initial_messages] == ["system", "user"]
async def fake_run_agent_loop(transcript_input, **_kwargs):
initial_messages = _assembled_messages(loop.context, transcript_input)
assert [m["role"] for m in initial_messages] == ["system", "user", "user"]
assert initial_messages[-2]["content"] == "earlier question that never got an answer"
assert initial_messages[-1]["content"] == "and another thing"
return _agent_run_result(
"done",
[
+3 -2
View File
@@ -5,6 +5,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.agent.context import TranscriptInput
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.context import (
RequestContext,
@@ -133,7 +134,7 @@ async def test_loop_binds_request_context_for_tool_execution(tmp_path: Path) ->
metadata = {"slack": {"thread_ts": "111.222", "channel_type": "channel"}}
runtime = loop.llm_runtime()
await loop._run_agent_loop(
[],
TranscriptInput(history=[], current_message=None),
runtime=runtime,
request_context=RequestContext(
channel="slack",
@@ -234,7 +235,7 @@ async def test_agent_loop_restores_outer_request_context_after_runner_exception(
try:
with pytest.raises(RuntimeError, match="runner failed"):
await loop._run_agent_loop(
[],
TranscriptInput(history=[], current_message=None),
runtime=runtime,
request_context=RequestContext(
channel="slack",
-48
View File
@@ -113,54 +113,6 @@ class TestHistoryWithCursor:
entries = store.read_unprocessed_history(since_cursor=0)
assert len(entries) == 2
def test_prompt_history_filters_to_current_session(self, store):
store.append_history("legacy entry without session")
store.append_history("telegram entry", session_key="telegram:chat-1")
store.append_history("slack entry", session_key="slack:chat-2")
entries = store.read_recent_history_for_prompt(
since_cursor=0,
session_key="telegram:chat-1",
)
assert [e["content"] for e in entries] == ["telegram entry"]
assert [e["content"] for e in store.read_unprocessed_history(0)] == [
"legacy entry without session",
"telegram entry",
"slack entry",
]
def test_unified_prompt_history_excludes_internal_cron_sessions(self, store):
store.append_history("legacy entry without session")
store.append_history("unified entry", session_key="unified:default")
store.append_history("telegram entry", session_key="telegram:chat-1")
store.append_history("cron internal entry", session_key="cron:job-1")
entries = store.read_recent_history_for_prompt(
since_cursor=0,
session_key="unified:default",
unified_session=True,
)
assert [e["content"] for e in entries] == [
"legacy entry without session",
"unified entry",
"telegram entry",
]
def test_unified_cron_prompt_history_includes_own_cron_entry(self, store):
store.append_history("unified entry", session_key="unified:default")
store.append_history("other cron entry", session_key="cron:job-2")
store.append_history("own cron entry", session_key="cron:job-1")
entries = store.read_recent_history_for_prompt(
since_cursor=0,
session_key="cron:job-1",
unified_session=True,
)
assert [e["content"] for e in entries] == ["unified entry", "own cron entry"]
def test_read_unprocessed_skips_entries_without_cursor(self, store):
"""Regression: entries missing the cursor key should be silently skipped."""
store.history_file.write_text(
+185
View File
@@ -0,0 +1,185 @@
"""Test /new archival behavior."""
import asyncio
from collections.abc import Coroutine
from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.utils.prompt_templates import render_template
_ARCHIVE_PROMPT = render_template("agent/consolidator_archive.md", strip=True)
class TestNewCommandArchival:
"""Test /new archival behavior with the structured archive flow."""
@staticmethod
def _make_loop(tmp_path: Path):
from nanobot.agent.loop import AgentLoop
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import GenerationSettings, LLMResponse
bus = MessageBus()
provider = MagicMock()
provider.get_default_model.return_value = "test-model"
provider.estimate_prompt_tokens.return_value = (10_000, "test")
provider.generation = GenerationSettings(max_tokens=100)
loop = AgentLoop(
bus=bus,
provider=provider,
workspace=tmp_path,
model="test-model",
context_window_tokens=1,
)
loop.provider.chat_with_retry = AsyncMock(
return_value=LLMResponse(content="ok", tool_calls=[])
)
loop.tools.get_definitions = MagicMock(return_value=[])
return loop
@pytest.mark.asyncio
async def test_new_clears_session_immediately_even_if_archive_fails(
self,
tmp_path: Path,
) -> None:
"""/new clears session immediately; archive is fire-and-forget."""
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
session = loop.sessions.get_or_create("cli:test")
for i in range(5):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
loop.sessions.save(session)
call_count = 0
expected_runtime = loop.llm_runtime()
async def _failing_summarize(session, *, archive_end, runtime) -> None:
nonlocal call_count
assert runtime is expected_runtime
assert session.key == "cli:test"
assert archive_end == len(session.messages)
call_count += 1
loop.consolidator.archive_session = _failing_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime)
assert response is not None
assert "new session started" in response.content.lower()
session_after = loop.sessions.get_or_create("cli:test")
assert len(session_after.messages) == 0
await loop.aclose()
assert call_count == 1
@pytest.mark.asyncio
async def test_new_reuses_replay_prefix_and_archives_only_unarchived_messages(
self,
tmp_path: Path,
) -> None:
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
loop.set_runtime_context_window(128_000)
session = loop.sessions.get_or_create("cli:test")
for i in range(5):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
session.last_archived = len(session.messages) - 2
ordinary_history = session.get_history()
assert [message["content"] for message in ordinary_history] == [
"msg1",
"resp1",
"msg2",
"resp2",
"msg3",
"resp3",
"msg4",
"resp4",
]
loop.sessions.save(session)
expected_runtime = loop.llm_runtime()
scheduled: list[Coroutine[Any, Any, object]] = []
loop.schedule_background = scheduled.append # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime)
assert response is not None
assert "new session started" in response.content.lower()
assert len(scheduled) == 1
await scheduled[0]
await loop.aclose()
sent = loop.provider.chat_with_retry.call_args.kwargs["messages"]
assert sent[1:-1] == ordinary_history
assert sent[-1]["content"] == _ARCHIVE_PROMPT
@pytest.mark.asyncio
async def test_new_clears_session_and_responds(self, tmp_path: Path) -> None:
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
session = loop.sessions.get_or_create("cli:test")
for i in range(3):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
loop.sessions.save(session)
expected_runtime = loop.llm_runtime()
async def _ok_summarize(session, *, archive_end, runtime) -> str:
assert runtime is expected_runtime
assert session.key == "cli:test"
assert archive_end == len(session.messages)
return "Summary."
loop.consolidator.archive_session = _ok_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
response = await loop._process_message(new_msg, runtime=expected_runtime)
assert response is not None
assert "new session started" in response.content.lower()
assert loop.sessions.get_or_create("cli:test").messages == []
@pytest.mark.asyncio
async def test_aclose_drains_background_tasks(self, tmp_path: Path) -> None:
"""aclose waits for background tasks to complete."""
from nanobot.bus.events import InboundMessage
loop = self._make_loop(tmp_path)
session = loop.sessions.get_or_create("cli:test")
for i in range(3):
session.add_message("user", f"msg{i}")
session.add_message("assistant", f"resp{i}")
loop.sessions.save(session)
archived = asyncio.Event()
release_archive = asyncio.Event()
expected_runtime = loop.llm_runtime()
async def _slow_summarize(session, *, archive_end, runtime) -> str:
assert runtime is expected_runtime
assert session.key == "cli:test"
assert archive_end == len(session.messages)
await release_archive.wait()
archived.set()
return "Summary."
loop.consolidator.archive_session = _slow_summarize # type: ignore[method-assign]
new_msg = InboundMessage(channel="cli", sender_id="user", chat_id="test", content="/new")
await loop._process_message(new_msg, runtime=expected_runtime)
assert not archived.is_set()
release_archive.set()
await loop.aclose()
assert archived.is_set()
+62 -97
View File
@@ -10,6 +10,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.context import TranscriptInput
from nanobot.agent.context_governance import ContextWindowExceededError
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import (
LLMProvider,
@@ -34,6 +36,35 @@ def _make_usage_spec(provider, tools):
)
def test_initial_transcript_is_built_from_structured_turn_input() -> None:
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
transcript_input = TranscriptInput(
history=[{"role": "user", "content": "earlier"}],
current_message="fresh",
)
expected = [
{"role": "system", "content": "system"},
{"role": "user", "content": "earlier"},
{"role": "user", "content": "fresh"},
]
transcript_builder = MagicMock(return_value=expected)
spec = make_run_spec(
provider,
initial_messages=None,
transcript_input=transcript_input,
transcript_builder=transcript_builder,
tools=MagicMock(),
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
assert AgentRunner._initial_transcript(spec) == expected
transcript_builder.assert_called_once_with(transcript_input)
def test_usage_or_estimate_replaces_reported_zero_for_content(monkeypatch) -> None:
from nanobot.agent.runner import AgentRunner
@@ -56,6 +87,7 @@ def test_usage_or_estimate_replaces_reported_zero_for_content(monkeypatch) -> No
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
tool_definitions=tools.get_definitions(),
)
assert usage == LLMUsage.estimated(input_tokens=12, output_tokens=7).with_timing(
@@ -100,6 +132,7 @@ def test_usage_or_estimate_counts_tool_call_output_for_reported_zero(monkeypatch
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
tool_definitions=tools.get_definitions(),
)
assert usage == LLMUsage.estimated(input_tokens=13, output_tokens=9)
@@ -132,6 +165,7 @@ def test_usage_or_estimate_counts_error_without_estimating_tokens(
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
tool_definitions=tools.get_definitions(),
)
assert usage is not None
@@ -167,6 +201,7 @@ def test_usage_or_estimate_trusts_positive_reported_total(monkeypatch) -> None:
_make_usage_spec(provider, tools),
[{"role": "user", "content": "hello"}],
response,
tool_definitions=tools.get_definitions(),
)
assert usage is not None
@@ -336,14 +371,12 @@ async def test_runner_replays_provider_state_without_chat_projection_duplicates(
@pytest.mark.asyncio
async def test_runner_governs_tool_result_before_adding_it_to_provider_state():
async def test_runner_preserves_tool_result_before_rejecting_unfit_followup():
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = True
provider.supports_native_compaction.return_value = False
calls = 0
captured_context: ProviderCallContext | None = None
checkpoints: list[dict] = []
state = ProviderConversationState(
kind="openai_responses",
@@ -354,7 +387,7 @@ async def test_runner_governs_tool_result_before_adding_it_to_provider_state():
)
async def chat_with_retry(**kwargs):
nonlocal calls, captured_context
nonlocal calls
calls += 1
if calls == 1:
return LLMResponse(
@@ -368,7 +401,6 @@ async def test_runner_governs_tool_result_before_adding_it_to_provider_state():
],
provider_state=state,
)
captured_context = kwargs["provider_context"]
return LLMResponse(content="done")
provider.chat_with_retry = chat_with_retry
@@ -379,37 +411,36 @@ async def test_runner_governs_tool_result_before_adding_it_to_provider_state():
async def checkpoint(payload: dict) -> None:
checkpoints.append(payload)
await AgentRunner().run(make_run_spec(
provider,
initial_messages=[
{"role": "system", "content": "system"},
{"role": "user", "content": "read the file"},
],
tools=tools,
model="gpt-5.6",
context_window_tokens=3_000,
context_block_limit=200,
max_tokens=1_000,
max_iterations=3,
max_tool_result_chars=10_000,
checkpoint_callback=checkpoint,
))
with pytest.raises(ContextWindowExceededError):
await AgentRunner().run(make_run_spec(
provider,
initial_messages=[
{"role": "system", "content": "system"},
{"role": "user", "content": "read the file"},
],
tools=tools,
model="gpt-5.6",
context_window_tokens=3_000,
context_block_limit=200,
max_tokens=1_000,
max_iterations=3,
max_tool_result_chars=10_000,
checkpoint_callback=checkpoint,
))
assert captured_context is not None
assert captured_context.conversation_state is not None
pending = captured_context.conversation_state.pending_messages
assert len(pending) == 1
assert pending[0]["role"] == "tool"
assert "compacted to fit context" in pending[0]["content"]
assert pending[0]["content"] != "x" * 5_000
assert calls == 1
completed_checkpoint = next(
checkpoint
for checkpoint in checkpoints
if checkpoint["phase"] == "tools_completed"
)
checkpoint_pending = completed_checkpoint["provider_state"].pending_messages
assert "compacted to fit context" in checkpoint_pending[0]["content"]
assert checkpoint_pending[0]["content"] != "x" * 5_000
assert checkpoint_pending == [{
"role": "tool",
"tool_call_id": "call_1",
"name": "read_file",
"content": "x" * 5_000,
}]
@pytest.mark.asyncio
@@ -798,64 +829,6 @@ async def test_runner_times_out_never_ending_streaming_request():
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_closes_progress_reasoning_on_streaming_wall_timeout():
from nanobot.agent.hook import AgentHook
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
provider.supports_progress_deltas = True
events: list[tuple[str, str | None]] = []
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
try:
await on_content_delta("<think>working...</think>")
await asyncio.sleep(3600)
finally:
events.append(("provider_cancelled", None))
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
tools = MagicMock()
tools.get_definitions.return_value = []
class ProgressReasoningHook(AgentHook):
async def emit_reasoning(self, reasoning_content: str | None) -> None:
if reasoning_content:
events.append(("reasoning", reasoning_content))
async def emit_reasoning_end(self) -> None:
events.append(("reasoning_end", None))
real_wait_for = asyncio.wait_for
async def fake_wait_for(coro, *, timeout):
assert timeout == 300.0
return await real_wait_for(coro, timeout=0.01)
runner = AgentRunner()
with patch("nanobot.agent.runner.asyncio.wait_for", fake_wait_for):
result = await runner.run(make_run_spec(provider,
initial_messages=[{"role": "user", "content": "think forever"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=ProgressReasoningHook(),
progress_callback=AsyncMock(),
llm_timeout_s=1,
))
assert result.stop_reason == "error"
assert result.final_content == "Error calling LLM: timed out after 300s"
assert events == [
("reasoning", "working..."),
("provider_cancelled", None),
("reasoning_end", None),
]
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_replaces_empty_tool_result_with_marker():
from nanobot.agent.runner import AgentRunner
@@ -1285,13 +1258,8 @@ async def test_runner_accumulates_usage_and_preserves_cache_reads():
@pytest.mark.asyncio
async def test_runner_binds_on_retry_wait_to_retry_callback_not_progress():
"""Regression: provider retry heartbeats must route through
``retry_wait_callback``, not ``progress_callback``. Binding them to
the progress callback (as an earlier runtime refactor did) caused
internal retry diagnostics like "Model request failed, retry in 1s"
to leak to end-user channels as normal progress updates.
"""
async def test_runner_binds_on_retry_wait_callback():
"""Provider retry heartbeats use the explicitly supplied callback."""
from nanobot.agent.runner import AgentRunner
captured: dict = {}
@@ -1305,7 +1273,6 @@ async def test_runner_binds_on_retry_wait_to_retry_callback_not_progress():
tools = MagicMock()
tools.get_definitions.return_value = []
progress_cb = AsyncMock()
retry_wait_cb = AsyncMock()
runner = AgentRunner()
@@ -1318,12 +1285,10 @@ async def test_runner_binds_on_retry_wait_to_retry_callback_not_progress():
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
retry_wait_callback=retry_wait_cb,
))
assert captured["on_retry_wait"] is retry_wait_cb
assert captured["on_retry_wait"] is not progress_cb
# ---------------------------------------------------------------------------
+33 -18
View File
@@ -9,7 +9,9 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.tools import ToolResult
from nanobot.agent.tools.execution import execute_tool_calls
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
@@ -50,16 +52,36 @@ async def test_runner_returns_tool_exception_to_model_for_recovery():
{"name": "list_dir", "status": "error", "detail": "boom"}
]
tool_message = next(message for message in result.messages if message.get("role") == "tool")
retry_hint = "[Analyze the error above and try a different approach.]"
assert "Error: RuntimeError: boom" in tool_message["content"]
assert tool_message["content"].count(retry_hint) == 1
@pytest.mark.asyncio
async def test_tool_execution_does_not_duplicate_existing_retry_hint():
retry_hint = "\n\n[Analyze the error above and try a different approach.]"
tools = SimpleNamespace(
execute=AsyncMock(return_value=ToolResult.error("Error: boom" + retry_hint)),
)
results, events = await execute_tool_calls(
tools,
[ToolCallRequest(id="call_1", name="list_dir", arguments={})],
concurrent=False,
external_lookup_counts={},
workspace_violation_counts={},
hook=AgentHook(),
context=AgentHookContext(iteration=0, messages=[]),
)
assert results == ["Error: boom" + retry_hint]
assert results[0].count(retry_hint) == 1
assert events[0]["status"] == "error"
@pytest.mark.asyncio
@pytest.mark.parametrize("control_error", [KeyboardInterrupt, SystemExit])
async def test_runner_propagates_tool_control_flow_exceptions(control_error: type[BaseException]):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
async def test_tool_execution_propagates_control_flow_exceptions(control_error: type[BaseException]):
async def execute(_name, _args):
raise control_error("stop")
@@ -67,22 +89,15 @@ async def test_runner_propagates_tool_control_flow_exceptions(control_error: typ
get_definitions=lambda: [],
execute=execute,
)
runner = AgentRunner()
spec = make_run_spec(
provider,
initial_messages=[],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
with pytest.raises(control_error):
await runner._run_tool(
spec,
ToolCallRequest(id="call_1", name="list_dir", arguments={}),
await execute_tool_calls(
tools,
[ToolCallRequest(id="call_1", name="list_dir", arguments={})],
concurrent=False,
external_lookup_counts={},
workspace_violation_counts={},
hook=AgentHook(),
context=AgentHookContext(iteration=0, messages=[]),
)
+526 -264
View File
@@ -1,8 +1,7 @@
"""Tests for AgentRunner context governance: backfill, orphan cleanup, microcompact, snip_history."""
"""Tests for AgentRunner context governance: repair and request fitting."""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@@ -12,11 +11,14 @@ from nanobot.agent.context_governance import (
BACKFILL_CONTENT,
ContextGovernanceConfig,
ContextGovernor,
ContextWindowExceededError,
)
from nanobot.agent.runner import AgentRunSpec
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import (
LLMProvider,
LLMResponse,
LLMUsage,
ProviderConversationState,
ToolCallRequest,
)
@@ -28,8 +30,6 @@ def _governance_config(
provider,
tools,
spec: AgentRunSpec,
*,
inflight_start_index: int = 0,
) -> ContextGovernanceConfig:
return ContextGovernanceConfig(
provider=provider,
@@ -41,7 +41,6 @@ def _governance_config(
context_window_tokens=spec.runtime.context_window_tokens,
context_block_limit=spec.context_block_limit,
max_tokens=spec.runtime.generation.max_tokens,
inflight_start_index=inflight_start_index,
)
@@ -89,6 +88,508 @@ async def test_runner_propagates_context_governance_failure():
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_locally_fits_oversized_initial_transcript(monkeypatch):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="done"))
tools = MagicMock()
tools.get_definitions.return_value = []
old_content = "x" * 20_000
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda _provider, _model, messages, _tools: (
(600, "test-counter")
if any(message.get("content") == old_content for message in messages)
else (100, "test-counter")
),
)
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[
{"role": "system", "content": "system"},
{"role": "user", "content": "old question"},
{"role": "assistant", "content": old_content},
{"role": "user", "content": "continue"},
],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_tokens=100,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
assert provider.chat_with_retry.await_args.kwargs["messages"] == [
{"role": "system", "content": "system"},
{"role": "user", "content": "continue"},
]
assert any(message.get("content") == old_content for message in result.messages)
@pytest.mark.asyncio
async def test_runner_governs_messages_added_by_before_iteration_hook(monkeypatch):
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="unexpected"))
tools = MagicMock()
tools.get_definitions.return_value = []
oversized = "hook-added-oversized-message"
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda _provider, _model, messages, _tools: (
(2_000, "test-counter")
if any(message.get("content") == oversized for message in messages)
else (100, "test-counter")
),
)
class MutatingHook(AgentHook):
async def before_iteration(self, context: AgentHookContext) -> None:
context.messages.append({"role": "user", "content": oversized})
with pytest.raises(ContextWindowExceededError):
await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "hello"}],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=MutatingHook(),
))
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_drops_resumable_provider_state_when_request_is_fitted(monkeypatch):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = True
captured_contexts = []
old_content = "old-oversized-history"
candidate = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="local-model",
version=1,
payload={"items": [{"type": "message", "content": "fresh state"}]},
)
async def chat_with_retry(*, provider_context=None, **_kwargs):
captured_contexts.append(provider_context)
return LLMResponse(
content="done",
usage=LLMUsage.reported(input_tokens=100, output_tokens=10),
provider_state=candidate,
)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda _provider, _model, messages, _tools: (
(600, "test-counter")
if any(message.get("content") == old_content for message in messages)
else (100, "test-counter")
),
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_message_tokens",
lambda message: 450 if message.get("content") == old_content else 50,
)
saved_state = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="local-model",
version=1,
payload={"items": [{"type": "message", "content": "stale state"}]},
)
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[
{"role": "assistant", "content": old_content},
{"role": "user", "content": "continue"},
],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
provider_state=saved_state,
))
assert captured_contexts[0].conversation_state is None
assert result.provider_state is not None
assert result.provider_state.payload == candidate.payload
@pytest.mark.asyncio
async def test_runner_fits_each_malformed_retry_with_its_actual_tools(monkeypatch):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
calls: list[dict] = []
estimated_tools: list[object] = []
definitions = [{"type": "function", "function": {"name": "read_file"}}]
async def chat_with_retry(*, messages, tools=None, **_kwargs):
calls.append({"messages": [dict(message) for message in messages], "tools": tools})
if len(calls) < 3:
return LLMResponse(
content="bad tool request",
tool_calls=[ToolCallRequest(id=f"bad_{len(calls)}", name=None, arguments={})],
finish_reason="tool_calls",
usage=LLMUsage.reported(input_tokens=100, output_tokens=10),
)
return LLMResponse(
content="recovered",
usage=LLMUsage.reported(input_tokens=100, output_tokens=10),
)
def estimate(_provider, _model, messages, _tools):
estimated_tools.append(_tools)
user_count = sum(message.get("role") == "user" for message in messages)
return (600 if user_count > 1 else 100), "test-counter"
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = definitions
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
estimate,
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_message_tokens",
lambda _message: 300,
)
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "use a tool"}],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
assert [call["tools"] for call in calls] == [definitions, definitions, None]
assert definitions in estimated_tools
assert None in estimated_tools
assert [len(call["messages"]) for call in calls] == [1, 1, 1]
assert result.final_content == "recovered"
assert result.messages == [
{"role": "user", "content": "use a tool"},
{"role": "assistant", "content": "recovered"},
]
@pytest.mark.asyncio
async def test_runner_fits_empty_response_finalization_before_dispatch(monkeypatch):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
calls: list[dict] = []
async def chat_with_retry(*, messages, tools=None, **_kwargs):
calls.append({"messages": [dict(message) for message in messages], "tools": tools})
if len(calls) < 3:
return LLMResponse(
content=None,
usage=LLMUsage.reported(input_tokens=100, output_tokens=1),
)
return LLMResponse(
content="finalized",
usage=LLMUsage.reported(input_tokens=100, output_tokens=10),
)
def estimate(_provider, _model, messages, _tools):
contents = [str(message.get("content") or "") for message in messages]
has_original = "do task" in contents
has_finalization = any("conversation above" in content for content in contents)
return (600 if has_original and has_finalization else 100), "test-counter"
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
estimate,
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_message_tokens",
lambda _message: 300,
)
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "do task"}],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=3,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
assert len(calls) == 3
assert calls[-1]["tools"] is None
assert all(message.get("content") != "do task" for message in calls[-1]["messages"])
assert result.final_content == "finalized"
@pytest.mark.asyncio
async def test_runner_fits_max_iteration_finalization_before_dispatch(monkeypatch):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
calls: list[dict] = []
oversized_result = "oversized-current-tool-result"
async def chat_with_retry(*, messages, tools=None, **_kwargs):
calls.append({"messages": [dict(message) for message in messages], "tools": tools})
if len(calls) == 1:
return LLMResponse(
content="working",
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={})],
finish_reason="tool_calls",
usage=LLMUsage.reported(input_tokens=100, output_tokens=10),
)
return LLMResponse(
content="safe summary",
usage=LLMUsage.reported(input_tokens=100, output_tokens=10),
)
def estimate(_provider, _model, messages, _tools):
has_oversized = any(
message.get("content") == oversized_result for message in messages
)
return (600 if has_oversized else 100), "test-counter"
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
tools.execute = AsyncMock(return_value=oversized_result)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
estimate,
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_message_tokens",
lambda message: 600 if message.get("content") == oversized_result else 50,
)
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "inspect"}],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
assert len(calls) == 2
assert calls[-1]["tools"] is None
assert all(
message.get("content") != oversized_result
for message in calls[-1]["messages"]
)
assert any(message.get("content") == oversized_result for message in result.messages)
assert result.final_content == "safe summary"
@pytest.mark.parametrize(
("input_tokens", "expected_fitted"),
[(500, True), (100, False)],
)
def test_matching_reported_provider_usage_avoids_local_estimate(
monkeypatch,
input_tokens,
expected_fitted,
):
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
spec = make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "hello"}],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (_ for _ in ()).throw(
AssertionError("matching provider usage must be authoritative")
),
)
governor = ContextGovernor()
monkeypatch.setattr(governor, "fit_to_budget", lambda *_args, **_kwargs: [])
_messages, fitted = governor.fit_request(
_governance_config(provider, tools, spec),
spec.initial_messages,
LLMUsage.reported(input_tokens=input_tokens, output_tokens=10),
usage_matches_messages=True,
tool_definitions=tools.get_definitions(),
)
assert fitted is expected_fitted
def test_changed_messages_use_local_estimate_after_reported_usage(monkeypatch):
provider = MagicMock(spec=LLMProvider)
tools = MagicMock()
tools.get_definitions.return_value = []
spec = make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "new tool output"}],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
estimate = MagicMock(return_value=(600, "test-counter"))
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
estimate,
)
governor = ContextGovernor()
monkeypatch.setattr(governor, "fit_to_budget", lambda *_args, **_kwargs: [])
_messages, fitted = governor.fit_request(
_governance_config(provider, tools, spec),
spec.initial_messages,
LLMUsage.reported(input_tokens=900, output_tokens=10),
usage_matches_messages=False,
tool_definitions=tools.get_definitions(),
)
assert fitted is True
estimate.assert_called_once()
@pytest.mark.asyncio
async def test_runner_counts_resumed_provider_state_before_dispatch(monkeypatch):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
provider.can_resume_conversation_state.return_value = True
captured_contexts = []
async def chat_with_retry(*, provider_context=None, **_kwargs):
captured_contexts.append(provider_context)
return LLMResponse(
content="done",
usage=LLMUsage.reported(input_tokens=100, output_tokens=10),
)
provider.chat_with_retry = chat_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
current_message = {"role": "user", "content": "new delta"}
saved_state = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="local-model",
version=1,
payload={
"items": [{"type": "reasoning", "encrypted_content": "opaque"}],
"context_tokens": 450,
},
pending_messages=[current_message],
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (100, "test-counter"),
)
monkeypatch.setattr(
"nanobot.providers.conversation_state.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (100, "test-counter"),
)
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[current_message],
tools=tools,
model="local-model",
context_window_tokens=2_000,
context_block_limit=500,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
provider_state=saved_state,
))
assert captured_contexts[0].conversation_state is None
assert result.messages == [
current_message,
{"role": "assistant", "content": "done"},
]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("context_block_limit", "expected_budget"),
[(500, 500), (None, 0)],
)
async def test_runner_refuses_locally_fitted_request_that_still_cannot_fit(
monkeypatch,
context_block_limit,
expected_budget,
):
from nanobot.agent.runner import AgentRunner
provider = MagicMock(spec=LLMProvider)
provider.chat_with_retry = AsyncMock(return_value=LLMResponse(content="unexpected"))
tools = MagicMock()
tools.get_definitions.return_value = []
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (2_000, "test-counter"),
)
with pytest.raises(ContextWindowExceededError) as exc_info:
await AgentRunner().run(make_run_spec(
provider,
initial_messages=[
{"role": "system", "content": "oversized system"},
{"role": "user", "content": "oversized user"},
],
tools=tools,
model="local-model",
context_window_tokens=1_000,
context_block_limit=context_block_limit,
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
))
assert exc_info.value.estimated_tokens == 2_000
assert exc_info.value.input_budget == expected_budget
provider.chat_with_retry.assert_not_awaited()
def test_snip_history_drops_orphaned_tool_results_from_trimmed_slice(monkeypatch):
provider = MagicMock()
tools = MagicMock()
@@ -130,7 +631,11 @@ def test_snip_history_drops_orphaned_tool_results_from_trimmed_slice(monkeypatch
lambda msg: token_sizes.get(str(msg.get("content")), 40),
)
trimmed = ContextGovernor().snip_history(_governance_config(provider, tools, spec), messages)
trimmed = ContextGovernor().snip_history(
_governance_config(provider, tools, spec),
messages,
tool_definitions=tools.get_definitions(),
)
# After the fix, the user message is recovered so the sequence is valid
# for providers that require system → user (e.g. GLM error 1214).
@@ -182,7 +687,11 @@ def test_snip_history_reserves_budget_for_tool_definitions(monkeypatch):
lambda msg: token_sizes.get(str(msg.get("content")), 40),
)
trimmed = ContextGovernor().snip_history(_governance_config(provider, tools, spec), messages)
trimmed = ContextGovernor().snip_history(
_governance_config(provider, tools, spec),
messages,
tool_definitions=tools.get_definitions(),
)
contents = [message.get("content") for message in trimmed]
assert contents == ["system", "recent two"]
@@ -465,260 +974,6 @@ async def test_runner_backfill_only_mutates_model_context_not_returned_messages(
]
# ---------------------------------------------------------------------------
# Microcompact (stale tool result compaction)
# ---------------------------------------------------------------------------
def _microcompact_messages(*, total: int, tool_name: str, content: str) -> list[dict]:
messages: list[dict] = [{"role": "system", "content": "sys"}]
for i in range(total):
messages.append({
"role": "assistant",
"content": "",
"tool_calls": [{
"id": f"c{i}",
"type": "function",
"function": {"name": tool_name, "arguments": "{}"},
}],
})
messages.append({
"role": "tool",
"tool_call_id": f"c{i}",
"name": tool_name,
"content": content,
})
return messages
def test_microcompact_skips_when_prompt_under_hard_budget(monkeypatch):
"""Cache-friendly path: in-flight tool results stay stable while prompt fits."""
provider = MagicMock()
provider.generation = SimpleNamespace(max_tokens=0)
tools = MagicMock()
tools.get_definitions.return_value = []
total = 15
long_content = "x" * 600
messages = _microcompact_messages(total=total, tool_name="read_file", content=long_content)
spec = make_run_spec(provider,
initial_messages=messages,
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_tokens=0,
context_window_tokens=20_000,
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (1000, "test"),
)
result = ContextGovernor().compact_inflight_overflow(
_governance_config(provider, tools, spec),
messages,
set(),
)
assert result is messages
def test_microcompact_overflow_compacts_to_low_watermark(monkeypatch):
"""Overflow path: compact in-flight stale results with headroom for later calls."""
provider = MagicMock()
provider.generation = SimpleNamespace(max_tokens=0)
tools = MagicMock()
tools.get_definitions.return_value = []
total = 18
long_content = "x" * 600
messages = _microcompact_messages(total=total, tool_name="read_file", content=long_content)
spec = make_run_spec(provider,
initial_messages=messages,
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_tokens=0,
context_window_tokens=2224, # input budget 1200, low target 1020
)
def estimate(_provider, _model, msgs, _tools):
return sum(
100 if (content := msg.get("content")) == long_content
else 1 if isinstance(content, str) and "compacted to fit context" in content
else 0
for msg in msgs
if msg.get("role") == "tool"
), "test"
monkeypatch.setattr("nanobot.agent.context_governance.estimate_prompt_tokens_chain", estimate)
result = ContextGovernor().compact_inflight_overflow(
_governance_config(provider, tools, spec),
messages,
set(),
)
tool_msgs = [m for m in result if m.get("role") == "tool"]
compacted = [m for m in tool_msgs if "compacted to fit context" in str(m.get("content", ""))]
preserved = [m for m in tool_msgs if m.get("content") == long_content]
assert len(compacted) == 8
assert len(preserved) == total - 8
assert [m["tool_call_id"] for m in compacted] == [f"c{i}" for i in range(8)]
def test_microcompact_compacts_newest_when_it_alone_overflows(monkeypatch):
"""An unfit newest result tells the model to retry narrowly or report the limit."""
provider = MagicMock()
provider.generation = SimpleNamespace(max_tokens=0)
tools = MagicMock()
tools.get_definitions.return_value = []
long_content = "x" * 600
messages = _microcompact_messages(total=1, tool_name="read_file", content=long_content)
spec = make_run_spec(provider,
initial_messages=messages,
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_tokens=0,
context_window_tokens=2000,
context_block_limit=500,
)
def estimate(_provider, _model, msgs, _tools):
return sum(
1000 if msg.get("content") == long_content else 1
for msg in msgs
if msg.get("role") == "tool"
), "test"
monkeypatch.setattr("nanobot.agent.context_governance.estimate_prompt_tokens_chain", estimate)
compacted_tool_call_ids: set[str] = set()
result = ContextGovernor().compact_inflight_overflow(
_governance_config(provider, tools, spec),
messages,
compacted_tool_call_ids,
)
tool_msg = next(m for m in result if m.get("role") == "tool")
assert "compacted to fit context" in tool_msg["content"]
assert "Do not repeat the same call unchanged" in tool_msg["content"]
assert "Retry with a narrower path, query, range, or result limit" in tool_msg["content"]
assert "tell the user the task cannot fit" in tool_msg["content"]
assert compacted_tool_call_ids == {"c0"}
def test_context_governor_keeps_compaction_boundary_stable(monkeypatch):
provider = MagicMock()
provider.generation = SimpleNamespace(max_tokens=0)
tools = MagicMock()
tools.get_definitions.return_value = []
total = 18
long_content = "x" * 600
messages = _microcompact_messages(total=total, tool_name="read_file", content=long_content)
spec = make_run_spec(provider,
initial_messages=messages,
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_tokens=0,
context_window_tokens=2224,
)
def estimate(_provider, _model, msgs, _tools):
return sum(
100 if msg.get("content") == long_content else 1
for msg in msgs
if msg.get("role") == "tool"
), "test"
monkeypatch.setattr("nanobot.agent.context_governance.estimate_prompt_tokens_chain", estimate)
governor = ContextGovernor()
compacted_tool_call_ids: set[str] = set()
config = _governance_config(provider, tools, spec, inflight_start_index=0)
first = governor.compact_inflight_overflow(config, messages, compacted_tool_call_ids)
first_ids = set(compacted_tool_call_ids)
second = governor.compact_inflight_overflow(config, messages, compacted_tool_call_ids)
assert compacted_tool_call_ids == first_ids
assert [m.get("content") for m in second] == [m.get("content") for m in first]
def test_microcompact_preserves_short_results(monkeypatch):
"""Short tool results below the compaction threshold should not be replaced."""
provider = MagicMock()
provider.generation = SimpleNamespace(max_tokens=0)
tools = MagicMock()
tools.get_definitions.return_value = []
total = 15
messages = _microcompact_messages(total=total, tool_name="exec", content="short")
spec = make_run_spec(provider,
initial_messages=messages,
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_tokens=0,
context_window_tokens=2024,
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (2000, "test"),
)
result = ContextGovernor().compact_inflight_overflow(
_governance_config(provider, tools, spec),
messages,
set(),
)
assert result is messages # no copy needed — all stale results are short
def test_microcompact_skips_non_compactable_tools(monkeypatch):
"""Non-compactable tools (e.g. 'message') should never be replaced."""
provider = MagicMock()
provider.generation = SimpleNamespace(max_tokens=0)
tools = MagicMock()
tools.get_definitions.return_value = []
total = 15
long_content = "y" * 1000
messages = _microcompact_messages(total=total, tool_name="message", content=long_content)
spec = make_run_spec(provider,
initial_messages=messages,
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_tokens=0,
context_window_tokens=2024,
)
monkeypatch.setattr(
"nanobot.agent.context_governance.estimate_prompt_tokens_chain",
lambda *_args, **_kwargs: (2000, "test"),
)
result = ContextGovernor().compact_inflight_overflow(
_governance_config(provider, tools, spec),
messages,
set(),
)
assert result is messages # no compactable tools found
def test_governance_repairs_orphans_after_snip():
"""After snipping clips an assistant+tool_calls, orphan repair cleans up the tail."""
# Simulate snipping that keeps only the tail: drop the assistant with
@@ -818,7 +1073,11 @@ def test_snip_history_preserves_user_message_after_truncation(monkeypatch):
lambda msg: token_sizes.get(str(msg.get("content")), 100),
)
trimmed = ContextGovernor().snip_history(_governance_config(provider, tools, spec), messages)
trimmed = ContextGovernor().snip_history(
_governance_config(provider, tools, spec),
messages,
tool_definitions=tools.get_definitions(),
)
# The first non-system message MUST be user (not assistant).
non_system = [m for m in trimmed if m.get("role") != "system"]
@@ -863,7 +1122,11 @@ def test_snip_history_no_user_at_all_falls_back_gracefully(monkeypatch):
lambda msg: 100,
)
trimmed = ContextGovernor().snip_history(_governance_config(provider, tools, spec), messages)
trimmed = ContextGovernor().snip_history(
_governance_config(provider, tools, spec),
messages,
tool_definitions=tools.get_definitions(),
)
# Should not crash. The result should still be a valid list.
assert isinstance(trimmed, list)
@@ -871,7 +1134,6 @@ def test_snip_history_no_user_at_all_falls_back_gracefully(monkeypatch):
assert any(m.get("role") == "system" for m in trimmed)
# The _enforce_role_alternation safety net must be able to fix whatever
# _snip_history returns here — verify it produces a valid sequence.
from nanobot.providers.base import LLMProvider
fixed = LLMProvider._enforce_role_alternation(trimmed)
non_system = [m for m in fixed if m["role"] != "system"]
if non_system:
+8 -4
View File
@@ -10,6 +10,7 @@ import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.automation_turns import publish_next_deferred_turn
from nanobot.agent.context import TranscriptInput
from nanobot.agent.tools.context import RequestContext
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMResponse, ToolCallRequest
@@ -617,7 +618,7 @@ async def test_loop_injected_followup_preserves_image_media(tmp_path):
runtime = loop.llm_runtime()
result = await loop._run_agent_loop(
[{"role": "user", "content": "hello"}],
TranscriptInput(history=[{"role": "user", "content": "hello"}], current_message=None),
runtime=runtime,
request_context=RequestContext(channel="cli", chat_id="c", runtime=runtime),
pending_queue=pending_queue,
@@ -711,7 +712,10 @@ async def test_pending_injection_resolves_its_own_runtime_context(tmp_path):
runtime = loop.llm_runtime()
result = await loop._run_agent_loop(
[{"role": "user", "content": "initial message from user A"}],
TranscriptInput(
history=[{"role": "user", "content": "initial message from user A"}],
current_message=None,
),
runtime=runtime,
session=session,
request_context=RequestContext(
@@ -812,7 +816,7 @@ async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_p
runtime = loop.llm_runtime()
result = await loop._run_agent_loop(
[{"role": "user", "content": "hello"}],
TranscriptInput(history=[{"role": "user", "content": "hello"}], current_message=None),
runtime=runtime,
request_context=RequestContext(channel="cli", chat_id="c", runtime=runtime),
pending_queue=pending_queue,
@@ -1476,7 +1480,7 @@ async def test_pending_queue_preserves_overflow_for_next_injection_cycle(tmp_pat
runtime = loop.llm_runtime()
result = await loop._run_agent_loop(
[{"role": "user", "content": "hello"}],
TranscriptInput(history=[{"role": "user", "content": "hello"}], current_message=None),
runtime=runtime,
request_context=RequestContext(channel="cli", chat_id="c", runtime=runtime),
pending_queue=pending_queue,
+20 -68
View File
@@ -1,4 +1,4 @@
"""Tests for provider progress delta routing in the shared runner."""
"""Tests for runner progress hooks and provider event routing."""
import asyncio
from unittest.mock import AsyncMock, MagicMock
@@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.hook import CompositeHook
from nanobot.agent.hooks import FileEditActivityHook
from nanobot.agent.progress_hook import AgentProgressHook
from nanobot.agent.runner import AgentRunner
@@ -17,45 +16,9 @@ from nanobot.providers.base import LLMResponse, ToolCallRequest
_MAX_TOOL_RESULT_CHARS = AgentDefaults().max_tool_result_chars
@pytest.mark.asyncio
async def test_runner_streams_provider_progress_deltas_by_default():
"""Direct runner users keep the existing opt-in provider progress behavior."""
provider = MagicMock()
provider.supports_progress_deltas = True
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
await on_content_delta("he")
await on_content_delta("llo")
return LLMResponse(content="hello", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
tools = MagicMock()
tools.get_definitions.return_value = []
progress_cb = AsyncMock()
runner = AgentRunner()
result = await runner.run(make_run_spec(provider,
initial_messages=[
{"role": "system", "content": "system"},
{"role": "user", "content": "hi"},
],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
))
assert result.final_content == "hello"
assert [call.args[0] for call in progress_cb.await_args_list] == ["he", "llo"]
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_routes_hosted_tool_events_to_structured_progress():
provider = MagicMock()
provider.supports_progress_deltas = True
async def chat_stream_with_retry(*, on_content_delta, on_tool_call_delta, **kwargs):
await on_tool_call_delta({
@@ -88,13 +51,17 @@ async def test_runner_routes_hosted_tool_events_to_structured_progress():
tools.get_definitions.return_value = []
progress_events: list[dict] = []
progress_text: list[str] = []
streamed_text: list[str] = []
async def progress_cb(content, *, tool_events=None, **kwargs):
progress_text.append(content)
if tool_events:
progress_events.extend(tool_events)
hook = CompositeHook([AgentProgressHook(on_progress=progress_cb)])
async def stream_cb(content: str) -> None:
streamed_text.append(content)
hook = AgentProgressHook(on_progress=progress_cb, on_stream=stream_cb)
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "search X"}],
@@ -102,7 +69,6 @@ async def test_runner_routes_hosted_tool_events_to_structured_progress():
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
hook=hook,
))
@@ -133,14 +99,14 @@ async def test_runner_routes_hosted_tool_events_to_structured_progress():
"embeds": [],
},
]
assert progress_text == ['search X "nanobot oauth"', "", "done"]
assert progress_text == ['search X "nanobot oauth"', ""]
assert streamed_text == ["done"]
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_fails_pending_hosted_tool_when_model_request_fails():
provider = MagicMock()
provider.supports_progress_deltas = True
async def chat_stream_with_retry(*, on_tool_call_delta, **kwargs):
await on_tool_call_delta({
@@ -166,7 +132,10 @@ async def test_runner_fails_pending_hosted_tool_when_model_request_fails():
if tool_events:
progress_events.extend(tool_events)
hook = CompositeHook([AgentProgressHook(on_progress=progress_cb)])
async def stream_cb(_content: str) -> None:
pass
hook = AgentProgressHook(on_progress=progress_cb, on_stream=stream_cb)
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "search X"}],
@@ -174,7 +143,6 @@ async def test_runner_fails_pending_hosted_tool_when_model_request_fails():
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
hook=hook,
))
@@ -200,7 +168,6 @@ async def test_runner_fails_pending_hosted_tool_when_model_request_fails():
@pytest.mark.asyncio
async def test_runner_emits_write_file_diff_from_tool_execution_snapshots(tmp_path):
provider = MagicMock()
provider.supports_progress_deltas = True
call_count = 0
progress_events: list[dict] = []
(tmp_path / "big.txt").write_text("old\n", encoding="utf-8")
@@ -218,7 +185,7 @@ async def test_runner_emits_write_file_diff_from_tool_execution_snapshots(tmp_pa
def prepare_call(self, name, params):
return tool, params, None
async def chat_stream_with_retry(**kwargs):
async def chat_with_retry(**kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
@@ -235,8 +202,7 @@ async def test_runner_emits_write_file_diff_from_tool_execution_snapshots(tmp_pa
)
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
provider.chat_with_retry = chat_with_retry
tools = Tools()
runner = AgentRunner()
@@ -246,7 +212,6 @@ async def test_runner_emits_write_file_diff_from_tool_execution_snapshots(tmp_pa
model="test-model",
max_iterations=2,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
workspace=tmp_path,
hook=FileEditActivityHook(on_progress=progress_cb, workspace=tmp_path),
))
@@ -263,13 +228,11 @@ async def test_runner_emits_write_file_diff_from_tool_execution_snapshots(tmp_pa
and event["diff"]["format"] == "unified"
for event in progress_events
)
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_emits_edit_file_diff_from_tool_execution_snapshots(tmp_path):
provider = MagicMock()
provider.supports_progress_deltas = True
call_count = 0
progress_events: list[dict] = []
target = tmp_path / "notes.txt"
@@ -288,7 +251,7 @@ async def test_runner_emits_edit_file_diff_from_tool_execution_snapshots(tmp_pat
def prepare_call(self, name, params):
return tool, params, None
async def chat_stream_with_retry(**kwargs):
async def chat_with_retry(**kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
@@ -309,8 +272,7 @@ async def test_runner_emits_edit_file_diff_from_tool_execution_snapshots(tmp_pat
)
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
provider.chat_with_retry = chat_with_retry
tools = Tools()
runner = AgentRunner()
@@ -320,7 +282,6 @@ async def test_runner_emits_edit_file_diff_from_tool_execution_snapshots(tmp_pat
model="test-model",
max_iterations=2,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
workspace=tmp_path,
hook=FileEditActivityHook(on_progress=progress_cb, workspace=tmp_path),
))
@@ -335,13 +296,11 @@ async def test_runner_emits_edit_file_diff_from_tool_execution_snapshots(tmp_pat
and event["diff"]["format"] == "unified"
for event in progress_events
)
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_marks_file_edit_activity_failed_when_tool_errors(tmp_path):
provider = MagicMock()
provider.supports_progress_deltas = True
call_count = 0
progress_events: list[dict] = []
@@ -358,7 +317,7 @@ async def test_runner_marks_file_edit_activity_failed_when_tool_errors(tmp_path)
def prepare_call(self, name, params):
return tool, params, None
async def chat_stream_with_retry(**kwargs):
async def chat_with_retry(**kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
@@ -375,8 +334,7 @@ async def test_runner_marks_file_edit_activity_failed_when_tool_errors(tmp_path)
)
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
provider.chat_with_retry = chat_with_retry
tools = Tools()
runner = AgentRunner()
@@ -386,7 +344,6 @@ async def test_runner_marks_file_edit_activity_failed_when_tool_errors(tmp_path)
model="test-model",
max_iterations=2,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
workspace=tmp_path,
hook=FileEditActivityHook(on_progress=progress_cb, workspace=tmp_path),
))
@@ -395,13 +352,11 @@ async def test_runner_marks_file_edit_activity_failed_when_tool_errors(tmp_path)
assert progress_events[-1]["path"] == "aborted.txt"
assert progress_events[-1]["phase"] == "error"
assert progress_events[-1]["status"] == "error"
provider.chat_with_retry.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_marks_file_edit_activity_failed_when_cancelled(tmp_path):
provider = MagicMock()
provider.supports_progress_deltas = True
progress_events: list[dict] = []
executing = asyncio.Event()
target = tmp_path / "cancelled.txt"
@@ -426,7 +381,7 @@ async def test_runner_marks_file_edit_activity_failed_when_cancelled(tmp_path):
def prepare_call(self, name, params):
return tool, params, None
async def chat_stream_with_retry(**kwargs):
async def chat_with_retry(**kwargs):
return LLMResponse(
content=None,
tool_calls=[
@@ -439,8 +394,7 @@ async def test_runner_marks_file_edit_activity_failed_when_cancelled(tmp_path):
usage=None,
)
provider.chat_stream_with_retry = chat_stream_with_retry
provider.chat_with_retry = AsyncMock()
provider.chat_with_retry = chat_with_retry
tools = Tools()
runner = AgentRunner()
@@ -450,7 +404,6 @@ async def test_runner_marks_file_edit_activity_failed_when_cancelled(tmp_path):
model="test-model",
max_iterations=2,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
progress_callback=progress_cb,
workspace=tmp_path,
hook=FileEditActivityHook(on_progress=progress_cb, workspace=tmp_path),
)))
@@ -464,4 +417,3 @@ async def test_runner_marks_file_edit_activity_failed_when_cancelled(tmp_path):
assert progress_events[-1]["path"] == "cancelled.txt"
assert progress_events[-1]["status"] == "error"
assert progress_events[-1]["error"] == "Task interrupted before this tool finished."
provider.chat_with_retry.assert_not_awaited()
+300 -23
View File
@@ -9,12 +9,15 @@ channels, gated by ``context.streamed_reasoning`` rather than
from __future__ import annotations
import asyncio
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.progress_hook import AgentProgressHook
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMResponse, LLMUsage, ToolCallRequest
@@ -35,6 +38,63 @@ class _RecordingHook(AgentHook):
self.end_calls += 1
class _StreamRecordingHook(_RecordingHook):
def __init__(self) -> None:
super().__init__()
self.streamed: list[str] = []
def wants_streaming(self) -> bool:
return True
async def on_stream(self, _ctx: AgentHookContext, delta: str) -> None:
self.streamed.append(delta)
class _LifecycleRecordingHook(AgentHook):
def __init__(self) -> None:
super().__init__()
self.events: list[str] = []
def wants_streaming(self) -> bool:
return True
async def emit_reasoning(self, reasoning_content: str | None) -> None:
if reasoning_content:
self.events.append(f"reasoning:{reasoning_content}")
async def emit_reasoning_end(self) -> None:
self.events.append("reasoning_end")
async def on_stream(self, _ctx: AgentHookContext, delta: str) -> None:
self.events.append(f"content:{delta}")
async def on_stream_end(self, _ctx: AgentHookContext, *, resuming: bool) -> None:
self.events.append(f"stream_end:{resuming}")
async def before_execute_tools(self, context: AgentHookContext) -> None:
names = ",".join(call.name for call in context.tool_calls)
self.events.append(f"local_tools:{names}")
async def on_provider_tool_event(
self,
_context: AgentHookContext,
event: dict[str, Any],
) -> None:
self.events.append(f"hosted_tool:{event.get('phase')}")
class _BlockingReasoningEndHook(_LifecycleRecordingHook):
def __init__(self) -> None:
super().__init__()
self.reasoning_end_started = asyncio.Event()
self.release_reasoning_end = asyncio.Event()
async def emit_reasoning_end(self) -> None:
self.reasoning_end_started.set()
await self.release_reasoning_end.wait()
await super().emit_reasoning_end()
@pytest.mark.asyncio
async def test_runner_preserves_reasoning_fields_in_assistant_history():
"""Reasoning fields ride along on the persisted assistant message so
@@ -201,7 +261,6 @@ async def test_runner_emits_reasoning_content_even_when_answer_was_streamed():
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
provider.supports_progress_deltas = True
async def chat_stream_with_retry(*, on_content_delta=None, **kwargs):
if on_content_delta:
@@ -218,12 +277,7 @@ async def test_runner_emits_reasoning_content_even_when_answer_was_streamed():
tools = MagicMock()
tools.get_definitions.return_value = []
progress_calls: list[str] = []
async def _progress(content: str, **_kwargs):
progress_calls.append(content)
hook = _RecordingHook()
hook = _StreamRecordingHook()
runner = AgentRunner()
result = await runner.run(make_run_spec(provider,
initial_messages=[{"role": "user", "content": "question"}],
@@ -232,11 +286,10 @@ async def test_runner_emits_reasoning_content_even_when_answer_was_streamed():
max_iterations=3,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=hook,
progress_callback=_progress,
))
assert result.final_content == "The answer."
assert progress_calls, "answer should have streamed via progress callback"
assert hook.streamed == ["The ", "answer."]
assert hook.emitted == ["step-by-step deduction"]
@@ -247,7 +300,6 @@ async def test_runner_does_not_double_emit_when_inline_think_already_streamed():
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
provider.supports_progress_deltas = True
async def chat_stream_with_retry(*, on_content_delta=None, **kwargs):
if on_content_delta:
@@ -263,10 +315,16 @@ async def test_runner_does_not_double_emit_when_inline_think_already_streamed():
tools = MagicMock()
tools.get_definitions.return_value = []
async def _progress(content: str, **_kwargs):
reasoning_events: list[str] = []
async def _progress(content: str, *, reasoning: bool = False, **_kwargs):
if reasoning:
reasoning_events.append(content)
async def _stream(_content: str) -> None:
pass
hook = _RecordingHook()
hook = AgentProgressHook(on_progress=_progress, on_stream=_stream)
runner = AgentRunner()
result = await runner.run(make_run_spec(provider,
initial_messages=[{"role": "user", "content": "question"}],
@@ -275,12 +333,10 @@ async def test_runner_does_not_double_emit_when_inline_think_already_streamed():
max_iterations=3,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=hook,
progress_callback=_progress,
))
assert result.final_content == "The answer."
assert hook.emitted == ["working..."]
assert hook.end_calls >= 1, "reasoning stream must be closed once the answer starts"
assert reasoning_events == ["working..."]
@pytest.mark.asyncio
@@ -320,14 +376,6 @@ async def test_runner_closes_reasoning_stream_after_one_shot_response():
assert hook.end_calls == 1
class _StreamRecordingHook(_RecordingHook):
def wants_streaming(self) -> bool:
return True
async def on_stream(self, _ctx: AgentHookContext, delta: str) -> None:
pass
@pytest.mark.asyncio
async def test_runner_streams_native_thinking_deltas_without_post_hoc_dup():
"""Anthropic-style ``on_thinking_delta`` should fan out to ``emit_reasoning``;
@@ -370,6 +418,235 @@ async def test_runner_streams_native_thinking_deltas_without_post_hoc_dup():
assert hook.emitted == ["part1", "part2"]
@pytest.mark.asyncio
async def test_runner_closes_native_reasoning_before_streaming_answer():
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
async def chat_stream_with_retry(
*, on_content_delta=None, on_thinking_delta=None, **kwargs
):
if on_thinking_delta:
await on_thinking_delta("inspect")
if on_content_delta:
await on_content_delta("done")
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
hook = _LifecycleRecordingHook()
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "q"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=hook,
))
assert result.final_content == "done"
assert hook.events == [
"reasoning:inspect",
"reasoning_end",
"content:done",
"stream_end:False",
]
@pytest.mark.asyncio
async def test_runner_closes_native_reasoning_before_local_tool_execution():
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
responses = iter([
LLMResponse(
content="",
finish_reason="tool_calls",
tool_calls=[ToolCallRequest(id="call-1", name="list_dir", arguments={"path": "."})],
usage=None,
),
LLMResponse(content="done", tool_calls=[], usage=None),
])
async def chat_stream_with_retry(
*, on_content_delta=None, on_thinking_delta=None, **kwargs
):
response = next(responses)
if response.tool_calls:
if on_thinking_delta:
await on_thinking_delta("inspect")
elif on_content_delta:
await on_content_delta("done")
return response
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
tools.execute = AsyncMock(return_value="tool result")
hook = _LifecycleRecordingHook()
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "inspect"}],
tools=tools,
model="test-model",
max_iterations=2,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=hook,
))
assert result.final_content == "done"
assert hook.events == [
"reasoning:inspect",
"reasoning_end",
"stream_end:True",
"local_tools:list_dir",
"content:done",
"stream_end:False",
]
@pytest.mark.asyncio
async def test_runner_closes_native_reasoning_before_hosted_tool_event():
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
async def chat_stream_with_retry(
*, on_content_delta=None, on_thinking_delta=None, on_tool_call_delta=None, **kwargs
):
if on_thinking_delta:
await on_thinking_delta("search")
if on_tool_call_delta:
await on_tool_call_delta({
"kind": "hosted_tool",
"phase": "start",
"call_id": "search-1",
"name": "web_search",
"arguments": {"query": "nanobot"},
})
await on_tool_call_delta({
"kind": "hosted_tool",
"phase": "end",
"call_id": "search-1",
"name": "web_search",
"arguments": {"query": "nanobot"},
"result": {"count": 1},
})
if on_content_delta:
await on_content_delta("done")
return LLMResponse(content="done", tool_calls=[], usage=None)
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
hook = _LifecycleRecordingHook()
result = await AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "search"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=hook,
))
assert result.final_content == "done"
assert hook.events == [
"reasoning:search",
"reasoning_end",
"hosted_tool:start",
"hosted_tool:end",
"content:done",
"stream_end:False",
]
@pytest.mark.asyncio
async def test_runner_closes_native_reasoning_when_stream_is_cancelled():
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
reasoning_started = asyncio.Event()
release_provider = asyncio.Event()
async def chat_stream_with_retry(
*, on_thinking_delta=None, **kwargs
):
if on_thinking_delta:
await on_thinking_delta("inspect")
reasoning_started.set()
await release_provider.wait()
raise AssertionError("the cancelled provider call should not complete")
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
hook = _LifecycleRecordingHook()
task = asyncio.create_task(AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "inspect"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=hook,
)))
await reasoning_started.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert hook.events == ["reasoning:inspect", "reasoning_end"]
@pytest.mark.asyncio
async def test_runner_settles_native_reasoning_end_before_propagating_cancellation():
from nanobot.agent.runner import AgentRunner
provider = MagicMock()
async def chat_stream_with_retry(
*, on_content_delta=None, on_thinking_delta=None, **kwargs
):
if on_thinking_delta:
await on_thinking_delta("inspect")
if on_content_delta:
await on_content_delta("done")
raise AssertionError("the cancelled provider call should not complete")
provider.chat_stream_with_retry = chat_stream_with_retry
tools = MagicMock()
tools.get_definitions.return_value = []
hook = _BlockingReasoningEndHook()
task = asyncio.create_task(AgentRunner().run(make_run_spec(
provider,
initial_messages=[{"role": "user", "content": "inspect"}],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
hook=hook,
)))
await hook.reasoning_end_started.wait()
task.cancel()
await asyncio.sleep(0)
hook.release_reasoning_end.set()
with pytest.raises(asyncio.CancelledError):
await task
assert hook.events == ["reasoning:inspect", "reasoning_end"]
@pytest.mark.asyncio
async def test_runner_strips_thinking_tags_from_native_thinking_deltas():
from nanobot.agent.runner import AgentRunner
+6 -5
View File
@@ -9,6 +9,7 @@ import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.runner import AgentRunner
from nanobot.agent.tools import ToolResult
from nanobot.agent.tools.execution import is_ssrf_violation
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import LLMResponse, ToolCallRequest
@@ -66,20 +67,20 @@ async def test_runner_does_not_abort_on_workspace_violation_anymore():
def test_is_ssrf_violation_recognizes_private_url_blocks():
"""SSRF rejections are classified separately from workspace boundaries."""
ssrf_msg = "Error: Command blocked by safety guard (internal/private URL detected)"
assert AgentRunner._is_ssrf_violation(ssrf_msg) is True
assert AgentRunner._is_ssrf_violation(
assert is_ssrf_violation(ssrf_msg) is True
assert is_ssrf_violation(
"URL validation failed: Blocked: host resolves to private/internal address 192.168.1.2"
) is True
# Workspace-bound markers are NOT classified as SSRF.
assert AgentRunner._is_ssrf_violation(
assert is_ssrf_violation(
"Error: Command blocked by safety guard (path outside working dir)"
) is False
assert AgentRunner._is_ssrf_violation(
assert is_ssrf_violation(
"Path /tmp/x is outside allowed directory /ws"
) is False
# Deny / allowlist filter messages stay non-fatal too.
assert AgentRunner._is_ssrf_violation(
assert is_ssrf_violation(
"Error: Command blocked by deny pattern filter"
) is False
+70 -41
View File
@@ -3,14 +3,17 @@
from __future__ import annotations
import asyncio
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent.runner_helpers import make_run_spec
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.runner import AgentRunner
from nanobot.agent.tools.base import Tool, ToolResult
from nanobot.agent.tools.context import ToolContext
from nanobot.agent.tools.execution import execute_tool_calls
from nanobot.agent.tools.loader import ToolLoader
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.config.schema import AgentDefaults
@@ -150,31 +153,69 @@ def _tool_message(result, tool_call_id: str) -> dict:
@pytest.mark.asyncio
async def test_runner_propagates_tool_preparation_failure():
async def test_tool_execution_propagates_preparation_failure():
tools = MagicMock()
tools.prepare_call.side_effect = RuntimeError("tool preparation failed")
tools.execute = AsyncMock()
with pytest.raises(RuntimeError, match="tool preparation failed"):
await AgentRunner()._run_tool(
make_run_spec(
MagicMock(),
initial_messages=[],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
),
ToolCallRequest(id="call-1", name="demo", arguments={}),
{},
{},
await execute_tool_calls(
tools,
[ToolCallRequest(id="call-1", name="demo", arguments={})],
concurrent=False,
external_lookup_counts={},
workspace_violation_counts={},
hook=AgentHook(),
context=AgentHookContext(iteration=0, messages=[]),
)
tools.execute.assert_not_awaited()
@pytest.mark.asyncio
async def test_runner_batches_read_only_tools_before_exclusive_work():
async def test_tool_execution_propagates_cancellation_without_error_hook():
tools = MagicMock()
tools.prepare_call.return_value = (None, {}, None)
tools.execute = AsyncMock(side_effect=asyncio.CancelledError)
events: list[str] = []
class RecordingHook(AgentHook):
async def before_execute_tool(
self,
context: AgentHookContext,
tool_call: ToolCallRequest,
tool: Any,
params: Any,
) -> None:
events.append("before")
async def on_execute_tool_error(
self,
context: AgentHookContext,
tool_call: ToolCallRequest,
tool: Any,
params: Any,
error: Any,
) -> None:
events.append("error")
with pytest.raises(asyncio.CancelledError):
await execute_tool_calls(
tools,
[ToolCallRequest(id="call-1", name="demo", arguments={})],
concurrent=False,
external_lookup_counts={},
workspace_violation_counts={},
hook=RecordingHook(),
context=AgentHookContext(iteration=0, messages=[]),
)
assert events == ["before"]
@pytest.mark.asyncio
async def test_tool_execution_batches_read_only_tools_before_exclusive_work():
tools = ToolRegistry()
shared_events: list[str] = []
read_a = _DelayTool("read_a", delay=0.05, read_only=True, shared_events=shared_events)
@@ -184,24 +225,18 @@ async def test_runner_batches_read_only_tools_before_exclusive_work():
tools.register(read_b)
tools.register(write_a)
provider = MagicMock()
runner = AgentRunner()
await runner._execute_tools(
make_run_spec(provider,
initial_messages=[],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
concurrent_tools=True,
),
await execute_tool_calls(
tools,
[
ToolCallRequest(id="ro1", name="read_a", arguments={}),
ToolCallRequest(id="ro2", name="read_b", arguments={}),
ToolCallRequest(id="rw1", name="write_a", arguments={}),
],
{},
{},
concurrent=True,
external_lookup_counts={},
workspace_violation_counts={},
hook=AgentHook(),
context=AgentHookContext(iteration=0, messages=[]),
)
assert shared_events[0:2] == ["start:read_a", "start:read_b"]
@@ -212,7 +247,7 @@ async def test_runner_batches_read_only_tools_before_exclusive_work():
@pytest.mark.asyncio
async def test_runner_does_not_batch_exclusive_read_only_tools():
async def test_tool_execution_does_not_batch_exclusive_read_only_tools():
tools = ToolRegistry()
shared_events: list[str] = []
read_a = _DelayTool("read_a", delay=0.03, read_only=True, shared_events=shared_events)
@@ -228,24 +263,18 @@ async def test_runner_does_not_batch_exclusive_read_only_tools():
tools.register(ddg_like)
tools.register(read_b)
provider = MagicMock()
runner = AgentRunner()
await runner._execute_tools(
make_run_spec(provider,
initial_messages=[],
tools=tools,
model="test-model",
max_iterations=1,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
concurrent_tools=True,
),
await execute_tool_calls(
tools,
[
ToolCallRequest(id="ro1", name="read_a", arguments={}),
ToolCallRequest(id="ddg1", name="ddg_like", arguments={}),
ToolCallRequest(id="ro2", name="read_b", arguments={}),
],
{},
{},
concurrent=True,
external_lookup_counts={},
workspace_violation_counts={},
hook=AgentHook(),
context=AgentHookContext(iteration=0, messages=[]),
)
assert shared_events[0] == "start:read_a"
+24 -25
View File
@@ -148,28 +148,28 @@ def test_retain_recent_legal_suffix_keeps_recent_messages():
assert session.messages[-1]["content"] == "msg9"
def test_retain_recent_legal_suffix_adjusts_last_consolidated():
def test_retain_recent_legal_suffix_adjusts_last_archived():
session = Session(key="test:trim-cons")
for i in range(10):
session.messages.append({"role": "user", "content": f"msg{i}"})
session.last_consolidated = 7
session.last_archived = 7
session.retain_recent_legal_suffix(4)
assert len(session.messages) == 4
assert session.last_consolidated == 1
assert session.last_archived == 1
def test_retain_recent_legal_suffix_zero_clears_session():
session = Session(key="test:trim-zero")
for i in range(10):
session.messages.append({"role": "user", "content": f"msg{i}"})
session.last_consolidated = 5
session.last_archived = 5
session.retain_recent_legal_suffix(0)
assert session.messages == []
assert session.last_consolidated == 0
assert session.last_archived == 0
def test_retain_recent_legal_suffix_keeps_legal_tool_boundary():
@@ -188,15 +188,15 @@ def test_retain_recent_legal_suffix_keeps_legal_tool_boundary():
assert history[0]["content"] == "keep"
# --- last_consolidated > 0 ---
# --- last_archived > 0 ---
def test_orphan_trim_with_last_consolidated():
"""Orphan trimming works correctly when session is partially consolidated."""
def test_orphan_trim_with_last_archived():
"""Orphan trimming works correctly when a session is partially archived."""
session = Session(key="test:consolidated")
for i in range(10):
session.messages.append({"role": "user", "content": f"old {i}"})
session.messages.extend(_tool_turn("cons", i))
session.last_consolidated = 30
session.last_archived = 30
session.messages.append({"role": "user", "content": "recent"})
for i in range(15):
@@ -213,7 +213,7 @@ def test_get_history_replays_recent_messages_after_full_archive():
for i in range(10):
session.messages.append({"role": "user", "content": f"u{i}"})
session.messages.append({"role": "assistant", "content": f"a{i}"})
session.last_consolidated = len(session.messages)
session.last_archived = len(session.messages)
history = session.get_history(max_messages=100)
@@ -229,8 +229,8 @@ def test_get_history_replays_recent_messages_after_full_archive():
]
def test_get_history_extends_compacted_replay_to_preceding_user():
session = Session(key="test:compacted-tool-turn")
def test_get_history_extends_archived_replay_to_preceding_user():
session = Session(key="test:archived-tool-turn")
session.messages.extend(
[
{"role": "user", "content": "old"},
@@ -242,7 +242,7 @@ def test_get_history_extends_compacted_replay_to_preceding_user():
{"role": "assistant", "content": "done"},
]
)
session.last_consolidated = len(session.messages)
session.last_archived = len(session.messages)
history = session.get_history(max_messages=100)
@@ -251,8 +251,8 @@ def test_get_history_extends_compacted_replay_to_preceding_user():
_assert_no_orphans(history)
def test_compacted_tool_turn_can_extend_past_message_cap():
session = Session(key="test:long-compacted-tool-turn")
def test_archived_tool_turn_can_extend_past_message_cap():
session = Session(key="test:long-archived-tool-turn")
session.messages.extend(
[
{"role": "user", "content": "old"},
@@ -263,7 +263,7 @@ def test_compacted_tool_turn_can_extend_past_message_cap():
for i in range(50):
session.messages.extend(_tool_turn("keep", i))
session.messages.append({"role": "assistant", "content": "done"})
session.last_consolidated = len(session.messages)
session.last_archived = len(session.messages)
history = session.get_history(max_messages=120)
@@ -635,7 +635,7 @@ def test_fork_session_allows_index_equal_to_user_count(tmp_path):
assert [m["content"] for m in forked.messages] == ["round1", "answer1"]
def test_fork_session_drops_summary_when_fork_point_is_inside_consolidated_prefix(tmp_path):
def test_fork_session_drops_summary_when_fork_point_is_inside_archived_prefix(tmp_path):
manager = SessionManager(tmp_path)
source = manager.get_or_create("websocket:source")
source.messages = [
@@ -644,7 +644,7 @@ def test_fork_session_drops_summary_when_fork_point_is_inside_consolidated_prefi
{"role": "user", "content": "round2 fork me"},
{"role": "assistant", "content": "answer2"},
]
source.last_consolidated = 4
source.last_archived = 4
source.metadata["_last_summary"] = {"text": "round2 fork me and answer2"}
manager.save(source)
@@ -656,7 +656,7 @@ def test_fork_session_drops_summary_when_fork_point_is_inside_consolidated_prefi
assert forked is not None
assert [m["content"] for m in forked.messages] == ["round1", "answer1"]
assert forked.last_consolidated == 0
assert forked.last_archived == 0
assert "_last_summary" not in forked.metadata
@@ -880,7 +880,7 @@ def test_retain_recent_legal_suffix_returns_all_on_zero():
session = Session(key="test:zero-return")
for i in range(5):
session.messages.append({"role": "user", "content": f"msg{i}"})
session.last_consolidated = 3
session.last_archived = 3
result = session.retain_recent_legal_suffix(0)
@@ -889,22 +889,21 @@ def test_retain_recent_legal_suffix_returns_all_on_zero():
assert session.messages == []
def test_retain_recent_legal_suffix_last_consolidated_correct_in_else_branch():
"""last_consolidated after retain_recent_legal_suffix should reflect how
many retained messages were inside the old consolidated prefix."""
def test_retain_recent_legal_suffix_last_archived_correct_in_else_branch():
"""last_archived should count retained messages from the old archived prefix."""
session = Session(key="test:else-lc-correct")
# 20 messages: u0..u9, a0..a9
for i in range(10):
session.messages.append({"role": "user", "content": f"u{i}"})
for i in range(10):
session.messages.append({"role": "assistant", "content": f"a{i}"})
session.last_consolidated = 12 # u0..u9, a0, a1 consolidated
session.last_archived = 12 # u0..u9, a0, a1 archived
result = session.retain_recent_legal_suffix(4)
# Retained messages start from latest user (u9) + max_messages forward
# so retained = [u9, a0..a9][:4] → but these are from original indices 9..12
# Of those, indices 9,10,11 are < 12 (before_lc), so new_lc = 3
assert session.last_consolidated == 3
assert session.last_archived == 3
# already_cons should count dropped messages with original index < 12
assert result.already_consolidated_count == 9
+2 -2
View File
@@ -114,7 +114,7 @@ async def test_removed_session_model_preset_falls_back_and_clears_metadata(tmp_p
provider=base,
workspace=tmp_path,
model="base-model",
context_window_tokens=8_000,
context_window_tokens=16_000,
)
loop.schedule_background = lambda coro: coro.close() # type: ignore[method-assign]
session_key = "sdk:removed-preset"
@@ -196,7 +196,7 @@ async def test_sdk_custom_model_preset_metadata_does_not_select_runtime(
provider=base,
workspace=tmp_path,
model="base-model",
context_window_tokens=8_000,
context_window_tokens=16_000,
)
loop.schedule_background = lambda coro: coro.close() # type: ignore[method-assign]
bot = Nanobot(loop)
+1 -1
View File
@@ -179,7 +179,7 @@ def test_compact_probe_keeps_delivery_in_visible_suffix():
{"role": "assistant", "content": "a2"},
{"role": "assistant", "content": "a3"},
]
probe = Session(key="test:probe", messages=tail, last_consolidated=0)
probe = Session(key="test:probe", messages=tail)
probe.retain_recent_legal_suffix(3, extend_to_user=True)
+29
View File
@@ -520,6 +520,35 @@ class TestCancelBySession:
count = await sm.cancel_by_session("nonexistent")
assert count == 0
@pytest.mark.asyncio
async def test_cancels_active_and_queued_tasks(self, tmp_path):
sm = _manager(tmp_path, max_concurrent_subagents=1)
active_entered = asyncio.Event()
queued_entered = asyncio.Event()
async def _blocked_run(spec):
task = spec.initial_messages[-1]["content"]
if task == "active":
active_entered.set()
else:
queued_entered.set()
await asyncio.Event().wait()
sm.runner.run = _blocked_run
runtime = _runtime()
await sm.spawn("active", runtime=runtime, session_key="s1")
await asyncio.wait_for(active_entered.wait(), timeout=1.0)
await sm.spawn("queued", runtime=runtime, session_key="s1")
await asyncio.sleep(0)
assert not queued_entered.is_set()
assert await sm.cancel_by_session("s1") == 2
await asyncio.sleep(0)
assert not queued_entered.is_set()
assert sm._running_tasks == {}
assert sm._session_tasks == {}
@pytest.mark.asyncio
async def test_already_done_not_counted(self, tmp_path):
sm = _manager(tmp_path)
+1 -1
View File
@@ -254,7 +254,7 @@ class TestDispatch:
assert isinstance(second.event, StreamEndEvent)
@pytest.mark.asyncio
async def test_processing_lock_serializes(self):
async def test_same_session_dispatches_serialize(self):
from nanobot.bus.events import InboundMessage, OutboundMessage
loop, bus = _make_loop()
+1 -1
View File
@@ -291,7 +291,7 @@ class TestCmdNewUnifiedSession:
archived = loop.consolidator.archive_session.call_args.args[0]
assert archived.key == "unified:default"
assert archived.messages == expected_snapshot
assert archived.last_consolidated == 0
assert archived.last_archived == 0
loop.consolidator.archive_session.assert_called_once_with(
archived,
archive_end=len(expected_snapshot),
+137 -30
View File
@@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from nanobot.agent.context import TranscriptInput
from nanobot.agent.tools.context import RequestContext
from nanobot.config.schema import AgentDefaults
from nanobot.providers.base import GenerationSettings
@@ -211,8 +212,8 @@ async def test_spawn_forwards_temperature_to_run_spec(tmp_path):
@pytest.mark.asyncio
async def test_spawn_tool_rejects_when_at_concurrency_limit(tmp_path):
"""SpawnTool should return an error string when the concurrency limit is reached."""
async def test_background_spawn_waits_for_concurrency_capacity(tmp_path):
"""Background tasks should be accepted and start when capacity becomes available."""
from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.spawn import SpawnTool
from nanobot.bus.queue import MessageBus
@@ -224,14 +225,23 @@ async def test_spawn_tool_rejects_when_at_concurrency_limit(tmp_path):
workspace=tmp_path,
bus=bus,
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_concurrent_subagents=1,
)
mgr._announce_result = AsyncMock()
# Block the first subagent so it stays "running"
release = asyncio.Event()
first_entered = asyncio.Event()
second_entered = asyncio.Event()
release_first = asyncio.Event()
release_second = asyncio.Event()
async def fake_run(spec):
await release.wait()
task = spec.initial_messages[-1]["content"]
if task == "first task":
first_entered.set()
await release_first.wait()
else:
second_entered.set()
await release_second.wait()
return SimpleNamespace(
stop_reason="done",
final_content="done",
@@ -250,19 +260,24 @@ async def test_spawn_tool_rejects_when_at_concurrency_limit(tmp_path):
session_key="test:c1",
runtime=_runtime(provider),
)):
# First spawn succeeds
result = await tool.execute(task="first task")
assert "started" in result
first_result = await tool.execute(task="first task")
assert "started" in first_result
await asyncio.wait_for(first_entered.wait(), timeout=1.0)
# Second spawn should be rejected (default limit is 1)
result = await tool.execute(task="second task")
assert "Cannot spawn subagent" in result
assert "concurrency limit reached" in result
second_result = await tool.execute(task="second task")
assert "started" in second_result
tasks = list(mgr._running_tasks.values())
await asyncio.sleep(0)
assert not second_entered.is_set()
phases = {status.task_description: status.phase for status in mgr._task_statuses.values()}
assert phases == {"first task": "initializing", "second task": "queued"}
# Release the first subagent
release.set()
# Allow cleanup
await asyncio.gather(*mgr._running_tasks.values(), return_exceptions=True)
release_first.set()
await asyncio.wait_for(second_entered.wait(), timeout=1.0)
release_second.set()
await asyncio.gather(*tasks, return_exceptions=True)
await asyncio.sleep(0)
assert mgr._running_tasks == {}
@pytest.mark.asyncio
@@ -300,7 +315,7 @@ async def test_spawn_tool_waits_for_inline_result():
@pytest.mark.asyncio
async def test_inline_spawn_counts_toward_concurrency_limit(tmp_path):
async def test_inline_spawn_waits_for_concurrency_capacity(tmp_path):
from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.context import RequestContext, request_context
from nanobot.agent.tools.spawn import SpawnTool
@@ -312,12 +327,19 @@ async def test_inline_spawn_counts_toward_concurrency_limit(tmp_path):
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_concurrent_subagents=1,
)
release = asyncio.Event()
entered = asyncio.Event()
first_entered = asyncio.Event()
second_entered = asyncio.Event()
release_first = asyncio.Event()
release_second = asyncio.Event()
async def fake_run(spec):
entered.set()
await release.wait()
task = spec.initial_messages[-1]["content"]
if task == "first":
first_entered.set()
await release_first.wait()
else:
second_entered.set()
await release_second.wait()
return SimpleNamespace(
stop_reason="done",
final_content="done",
@@ -334,19 +356,100 @@ async def test_inline_spawn_counts_toward_concurrency_limit(tmp_path):
runtime=_runtime(MagicMock()),
)):
first = asyncio.create_task(tool.execute(task="first", wait=True))
await asyncio.wait_for(entered.wait(), timeout=1.0)
await asyncio.wait_for(first_entered.wait(), timeout=1.0)
second = await tool.execute(task="second", wait=True)
second = asyncio.create_task(tool.execute(task="second", wait=True))
await asyncio.sleep(0)
assert "concurrency limit reached" in second
assert manager.get_running_count() == 1
release.set()
assert not second.done()
assert not second_entered.is_set()
assert manager.get_running_count() == 2
release_first.set()
assert await first == "done"
await asyncio.wait_for(second_entered.wait(), timeout=1.0)
release_second.set()
assert await second == "done"
assert manager.get_running_count() == 0
assert manager._session_tasks == {}
@pytest.mark.asyncio
async def test_runner_executes_inline_spawn_batch_concurrently(tmp_path):
"""Adjacent blocking consultations should share one concurrent tool batch."""
from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.agent.subagent import SubagentManager
from nanobot.agent.tools.context import RequestContext, request_context
from nanobot.agent.tools.execution import execute_tool_calls
from nanobot.agent.tools.registry import ToolRegistry
from nanobot.agent.tools.spawn import SpawnTool
from nanobot.bus.queue import MessageBus
from nanobot.providers.base import ToolCallRequest
manager = SubagentManager(
workspace=tmp_path,
bus=MessageBus(),
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
max_concurrent_subagents=2,
)
both_entered = asyncio.Event()
release = asyncio.Event()
entered: list[str] = []
async def fake_run(spec):
entered.append(spec.initial_messages[-1]["content"])
if len(entered) == 2:
both_entered.set()
await release.wait()
return SimpleNamespace(
stop_reason="done",
final_content=spec.initial_messages[-1]["content"],
error=None,
tool_events=[],
)
manager.runner.run = AsyncMock(side_effect=fake_run)
tools = ToolRegistry()
tools.register(SpawnTool(manager))
runtime = _runtime(MagicMock())
calls = [
ToolCallRequest(
id="spawn-1",
name="spawn",
arguments={"task": "first", "wait": True},
),
ToolCallRequest(
id="spawn-2",
name="spawn",
arguments={"task": "second", "wait": True},
),
]
with request_context(RequestContext(
channel="test",
chat_id="c1",
session_key="test:c1",
runtime=runtime,
)):
execution = asyncio.create_task(execute_tool_calls(
tools,
calls,
concurrent=True,
external_lookup_counts={},
workspace_violation_counts={},
hook=AgentHook(),
context=AgentHookContext(iteration=0, messages=[], session_key="test:c1"),
))
await asyncio.wait_for(both_entered.wait(), timeout=1.0)
release.set()
results, events = await execution
assert set(entered) == {"first", "second"}
assert results == ["first", "second"]
assert [event["status"] for event in events] == ["ok", "ok"]
assert manager._running_tasks == {}
@pytest.mark.asyncio
async def test_cancel_by_session_cancels_inline_subagent(tmp_path):
from nanobot.agent.subagent import SubagentManager
@@ -391,6 +494,7 @@ def test_subagent_default_max_concurrent_matches_agent_defaults(tmp_path):
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
)
assert AgentDefaults().max_concurrent_subagents == 4
assert mgr.max_concurrent_subagents == AgentDefaults().max_concurrent_subagents
@@ -465,7 +569,10 @@ async def test_agent_loop_syncs_updated_max_iterations_before_run(tmp_path):
loop.runner.run = AsyncMock(side_effect=fake_run)
loop.max_iterations = 55
await loop._run_agent_loop([], runtime=loop.llm_runtime())
await loop._run_agent_loop(
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
)
loop.runner.run.assert_awaited_once()
@@ -506,7 +613,7 @@ async def test_drain_pending_no_block_when_no_subagents(tmp_path):
runtime = loop.llm_runtime()
await loop._run_agent_loop(
[{"role": "user", "content": "test"}],
TranscriptInput(history=[{"role": "user", "content": "test"}], current_message=None),
runtime=runtime,
session=None,
request_context=RequestContext(channel="test", chat_id="c1", runtime=runtime),
@@ -565,7 +672,7 @@ async def test_terminal_drain_timeout(tmp_path):
runtime = loop.llm_runtime()
await loop._run_agent_loop(
[{"role": "user", "content": "test"}],
TranscriptInput(history=[{"role": "user", "content": "test"}], current_message=None),
runtime=runtime,
session=session,
request_context=RequestContext(
@@ -639,7 +746,7 @@ async def test_terminal_drain_reuses_one_timeout_budget(tmp_path):
loop.subagents._running_tasks["sub-deadline-1"] = hang_task
await loop._run_agent_loop(
[{"role": "user", "content": "test"}],
TranscriptInput(history=[{"role": "user", "content": "test"}], current_message=None),
runtime=loop.llm_runtime(),
session=session,
pending_queue=pending_queue,
+126 -6
View File
@@ -799,7 +799,7 @@ def test_provider_login_can_set_xai_grok_as_main_provider(tmp_path):
saved = Config.model_validate(json.loads(config_path.read_text(encoding="utf-8")))
assert saved.agents.defaults.provider == "xai_grok"
assert saved.agents.defaults.model == "xai-grok/grok-4.5"
assert saved.agents.defaults.model == "xai-grok/grok-4.6"
assert saved.agents.defaults.context_window_tokens == 500_000
assert saved.agents.defaults.model_preset is None
assert make_provider(saved).__class__.__name__ == "XAIGrokProvider"
@@ -2534,7 +2534,11 @@ def test_webui_yes_still_refuses_invalid_custom_model_setup(
def test_open_webui_browser_redacts_bootstrap_secret(monkeypatch, capsys) -> None:
opened: list[str] = []
url = "http://127.0.0.1:8765/#/?bootstrapSecret=super-secret"
monkeypatch.setattr("webbrowser.open", lambda value: opened.append(value))
monkeypatch.setattr(
cli_webui_support,
"_launch_browser",
lambda value: opened.append(value) or True,
)
cli_webui_support._open_webui_browser(url, wait=False)
@@ -2544,6 +2548,42 @@ def test_open_webui_browser_redacts_bootstrap_secret(monkeypatch, capsys) -> Non
assert "super-secret" not in output
def test_open_webui_browser_reports_launch_failure(monkeypatch, capsys) -> None:
monkeypatch.setattr(cli_webui_support, "_launch_browser", lambda _value: False)
cli_webui_support._open_webui_browser("http://127.0.0.1:8765/", wait=False)
assert "Could not open browser; visit http://127.0.0.1:8765/" in _strip_ansi(
capsys.readouterr().out
)
def test_launch_browser_uses_macos_foreground_opener(monkeypatch) -> None:
seen: list[list[str]] = []
monkeypatch.setattr(cli_webui_support.sys, "platform", "darwin")
monkeypatch.setattr(
cli_webui_support.subprocess,
"run",
lambda command, **_kwargs: seen.append(command) or SimpleNamespace(returncode=0),
)
assert cli_webui_support._launch_browser("http://127.0.0.1:8765/") is True
assert seen == [["open", "http://127.0.0.1:8765/"]]
def test_launch_browser_uses_default_browser_off_macos(monkeypatch) -> None:
opened: list[tuple[str, int, bool]] = []
monkeypatch.setattr(cli_webui_support.sys, "platform", "linux")
monkeypatch.setattr(
cli_webui_support.webbrowser,
"open",
lambda url, *, new, autoraise: opened.append((url, new, autoraise)) or True,
)
assert cli_webui_support._launch_browser("http://127.0.0.1:8765/") is True
assert opened == [("http://127.0.0.1:8765/", 2, True)]
def test_webui_foreground_attaches_to_existing_managed_gateway(monkeypatch, tmp_path: Path) -> None:
config_file = tmp_path / "config.json"
config_file.write_text("{}")
@@ -2614,12 +2654,14 @@ def test_webui_foreground_attaches_to_existing_managed_gateway(monkeypatch, tmp_
assert seen["lease_release_wait_for_stop"] is False
def test_attach_to_background_gateway_detaches_on_ctrl_c(capsys) -> None:
def test_attach_to_background_gateway_detaches_on_ctrl_c(capsys, tmp_path: Path) -> None:
stopped = False
log_path = tmp_path / "gateway.log"
log_path.touch()
class _FakeRuntime:
def status(self):
return SimpleNamespace(running=True)
return SimpleNamespace(running=True, log_path=log_path)
def stop(self):
nonlocal stopped
@@ -2639,10 +2681,88 @@ def test_attach_to_background_gateway_detaches_on_ctrl_c(capsys) -> None:
assert "WebUI launcher detached" in rendered
def test_attach_to_background_gateway_checks_owned_sidecar() -> None:
def test_attach_to_background_gateway_follows_only_new_logs(capsys, tmp_path: Path) -> None:
log_path = tmp_path / "gateway.log"
log_path.write_text("historical log\n", encoding="utf-8")
polls = 0
class _FakeRuntime:
def status(self):
return SimpleNamespace(running=True)
return SimpleNamespace(running=True, log_path=log_path)
def _append_then_interrupt(_seconds: float) -> None:
nonlocal polls
if polls == 0:
with log_path.open("a", encoding="utf-8") as handle:
handle.write("[websocket] live log\n")
polls += 1
return
raise KeyboardInterrupt
cli_webui_support._attach_to_background_gateway(
_FakeRuntime(),
sleep=_append_then_interrupt,
)
output = capsys.readouterr().out
assert "[websocket] live log" in output
assert "historical log" not in output
def test_read_new_gateway_logs_recovers_after_truncation(tmp_path: Path) -> None:
log_path = tmp_path / "gateway.log"
log_path.write_text("a much longer historical log line\n", encoding="utf-8")
cursor = cli_webui_support._start_gateway_log_cursor(log_path)
log_path.write_text("fresh log\n", encoding="utf-8")
lines = cli_webui_support._read_new_gateway_logs(log_path, cursor)
assert lines == ["fresh log"]
assert cursor.offset == log_path.stat().st_size
def test_read_new_gateway_logs_detects_fast_rewrite_past_offset(tmp_path: Path) -> None:
log_path = tmp_path / "gateway.log"
log_path.write_text("historical log\n", encoding="utf-8")
cursor = cli_webui_support._start_gateway_log_cursor(log_path)
log_path.write_text("first fresh log\nsecond fresh log\n", encoding="utf-8")
lines = cli_webui_support._read_new_gateway_logs(log_path, cursor)
assert lines == ["first fresh log", "second fresh log"]
def test_read_new_gateway_logs_waits_for_complete_utf8_line(tmp_path: Path) -> None:
log_path = tmp_path / "gateway.log"
log_path.touch()
cursor = cli_webui_support._start_gateway_log_cursor(log_path)
encoded = "模型 ready\n".encode()
log_path.write_bytes(encoded[:2])
assert cli_webui_support._read_new_gateway_logs(log_path, cursor) == []
with log_path.open("ab") as handle:
handle.write(encoded[2:])
assert cli_webui_support._read_new_gateway_logs(log_path, cursor) == ["模型 ready"]
def test_read_new_gateway_logs_tolerates_missing_file(tmp_path: Path) -> None:
log_path = tmp_path / "missing.log"
cursor = cli_webui_support._start_gateway_log_cursor(log_path)
lines = cli_webui_support._read_new_gateway_logs(log_path, cursor)
assert lines == []
assert cursor.offset == 0
def test_attach_to_background_gateway_checks_owned_sidecar(tmp_path: Path) -> None:
log_path = tmp_path / "gateway.log"
log_path.touch()
class _FakeRuntime:
def status(self):
return SimpleNamespace(running=True, log_path=log_path)
def sidecar_exited() -> None:
raise WebUIDevError("WebUI development server exited unexpectedly (code 23)")
+82 -1
View File
@@ -1,4 +1,85 @@
from nanobot.cli.entry import _native_tui_candidate
import os
import subprocess
import sys
from pathlib import Path
from nanobot.cli import entry
from nanobot.cli.entry import _agent_invocation_args, _native_tui_candidate
def test_root_command_routes_to_agent_without_copying_agent_options() -> None:
assert _agent_invocation_args([]) == []
assert _agent_invocation_args(["agent", "--theme", "dark"]) == ["--theme", "dark"]
assert _agent_invocation_args(["--workspace", "./project"]) == [
"--workspace",
"./project",
]
assert _agent_invocation_args(["-mhello"]) == ["-mhello"]
def test_root_metadata_and_subcommands_keep_the_root_cli() -> None:
for args in (
["--help"],
["--version"],
["--install-completion"],
["gateway"],
["webui"],
):
assert _agent_invocation_args(args) is None
def test_root_shell_completion_keeps_root_subcommands() -> None:
env = os.environ.copy()
env.update(
{
"_NANOBOT_COMPLETE": "complete_bash",
"COMP_WORDS": "nanobot ",
"COMP_CWORD": "1",
}
)
script = (
"import sys; "
"from nanobot.cli.entry import main; "
"sys.argv = ['nanobot']; "
"main()"
)
result = subprocess.run(
[sys.executable, "-c", script],
cwd=Path(__file__).parents[2],
env=env,
capture_output=True,
text=True,
check=False,
)
assert result.returncode == 0, result.stderr
assert {"agent", "gateway", "webui"} <= set(result.stdout.splitlines())
assert "not supported" not in result.stderr
def test_root_alias_dispatches_the_shared_agent_command(monkeypatch) -> None:
calls: dict[str, object] = {}
monkeypatch.setattr(entry.sys, "argv", ["nanobot", "-m", "hello"])
monkeypatch.setattr(
entry,
"set_cli_process_identity",
lambda args: calls.__setitem__("identity", args),
)
monkeypatch.setattr(entry, "_configure_windows_console", lambda: None)
monkeypatch.setattr(
entry,
"_run_agent",
lambda args, *, prog_name: calls.update(args=args, prog_name=prog_name),
)
entry.main()
assert calls == {
"identity": ["agent", "-m", "hello"],
"args": ["-m", "hello"],
"prog_name": "nanobot",
}
def test_native_agent_invocations_use_the_lightweight_entrypoint() -> None:
+18
View File
@@ -58,6 +58,24 @@ def test_legacy_console_entrypoint_still_sets_subcommand_identity(
assert commands == [["webui"]]
def test_legacy_console_entrypoint_routes_bare_command_to_agent(
monkeypatch: pytest.MonkeyPatch,
) -> None:
identities: list[list[str]] = []
launches: list[tuple[list[str], str]] = []
monkeypatch.setattr("nanobot.cli.commands.set_cli_process_identity", identities.append)
monkeypatch.setattr(
"nanobot.cli.entry._run_agent",
lambda args, *, prog_name: launches.append((args, prog_name)),
)
result = CliRunner().invoke(app, [])
assert result.exit_code == 0
assert identities == [["agent"]]
assert launches == [([], "nanobot")]
def test_named_executable_creates_stable_role_symlink(
tmp_path: Path,
) -> None:
+9 -2
View File
@@ -11,6 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from nanobot.agent.context import TranscriptInput
from nanobot.bus.events import InboundMessage
from nanobot.providers.base import LLMResponse, LLMUsage
@@ -311,10 +312,16 @@ class TestRestartCommand:
LLMResponse(content="second", usage=None),
])
first = await loop._run_agent_loop([], runtime=loop.llm_runtime())
first = await loop._run_agent_loop(
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
)
assert first.usage == LLMUsage.reported(input_tokens=9, output_tokens=4)
second = await loop._run_agent_loop([], runtime=loop.llm_runtime())
second = await loop._run_agent_loop(
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
)
assert second.usage == LLMUsage.estimated(input_tokens=123, output_tokens=7)
@pytest.mark.asyncio
+1
View File
@@ -142,6 +142,7 @@ def test_launcher_passes_the_canonical_model_preset_to_the_tui(
assert captured["NANOBOT_TUI_BOOTSTRAP_URL"] == (
"http://127.0.0.1:8765/webui/bootstrap"
)
assert captured["NANOBOT_TUI_HEALTH_URL"] == "http://127.0.0.1:18790/health"
assert captured["NANOBOT_TUI_BOOTSTRAP_SECRET"] == "bootstrap-secret"
assert "NANOBOT_TUI_WS_URL" not in captured
assert "NANOBOT_TUI_API_TOKEN" not in captured
+1 -1
View File
@@ -110,7 +110,7 @@ class TestMidTurnCommandDispatchedDirectly:
loop = MagicMock()
loop.sessions = MagicMock()
loop.sessions.get_or_create = MagicMock(return_value=MagicMock(
messages=[], last_consolidated=0, clear=MagicMock(),
messages=[], last_archived=0, clear=MagicMock(),
))
loop.sessions.save = MagicMock()
loop.sessions.invalidate = MagicMock()
+40 -1
View File
@@ -7,6 +7,7 @@ import pytest
from nanobot.cron.service import CronJobSkippedError, CronService
from nanobot.cron.types import CronJob, CronPayload, CronSchedule
from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META
async def _wait_until(predicate, *, timeout: float = 1.0, interval: float = 0.01) -> None:
@@ -292,7 +293,12 @@ def test_load_store_migrates_legacy_delivery_context(tmp_path) -> None:
"deliver": True,
"channel": "telegram",
"to": "user-1",
"channelMeta": {"message_thread_id": 42},
"channelMeta": {
"message_thread_id": 42,
RUNTIME_CONTEXT_INPUT_META: [
{"source": "webui_quote", "content": "stale quote"}
],
},
"sessionKey": "telegram:user-1:topic:42",
},
"state": {},
@@ -411,6 +417,39 @@ def test_add_job_preserves_origin_delivery_context(tmp_path) -> None:
assert reloaded.payload.origin_metadata == metadata
@pytest.mark.asyncio
async def test_start_heals_runtime_context_from_pending_external_add(tmp_path) -> None:
"""Flattened runtime blocks from older action files must not be replayed."""
store_path = tmp_path / "cron" / "jobs.json"
external = CronService(store_path)
job = external.add_job(
name="quoted reminder",
schedule=CronSchedule(kind="every", every_ms=60_000),
message="remember this",
origin_metadata={"webui": True},
**_bound_chat("quoted"),
)
action_path = tmp_path / "cron" / "action.jsonl"
action = json.loads(action_path.read_text(encoding="utf-8"))
action["params"]["payload"]["origin_metadata"][RUNTIME_CONTEXT_INPUT_META] = [
{"source": "webui_quote", "content": "quoted reply"}
]
action_path.write_text(json.dumps(action), encoding="utf-8")
owner = CronService(store_path)
await owner.start()
try:
loaded = owner.get_job(job.id)
assert loaded is not None
assert loaded.payload.origin_metadata == {"webui": True}
raw = json.loads(store_path.read_text(encoding="utf-8"))
assert raw["jobs"][0]["payload"]["originMetadata"] == {"webui": True}
finally:
owner.stop()
@pytest.mark.asyncio
async def test_channel_meta_and_session_key_survive_store_reload(tmp_path) -> None:
store_path = tmp_path / "cron" / "jobs.json"
@@ -146,6 +146,51 @@ def test_controller_uses_governed_messages_for_provider_state_delta() -> None:
assert governed_checkpoint.pending_messages[-1]["content"] == "compacted result"
def test_controller_estimates_active_state_plus_pending_delta(monkeypatch) -> None:
provider = _provider()
current_message = {"role": "user", "content": "new delta"}
state = ProviderConversationState(
kind="openai_responses",
provider="openai:test",
model="gpt-5.6",
version=1,
payload={
"items": [{"type": "reasoning", "encrypted_content": "opaque"}],
"context_tokens": 450,
},
pending_messages=[current_message],
)
controller = ProviderConversationStateController(
provider=provider,
model="gpt-5.6",
messages=[current_message],
state=state,
)
seen = {}
def estimate(_provider, _model, messages, tools):
seen["messages"] = messages
seen["tools"] = tools
return 100, "test-counter"
monkeypatch.setattr(
"nanobot.providers.conversation_state.estimate_prompt_tokens_chain",
estimate,
)
tokens = controller.estimate_request_context_tokens(
[current_message],
model_messages=[current_message],
tool_definitions=[{"type": "web_search"}],
)
assert tokens == 550
assert seen == {
"messages": [current_message],
"tools": [{"type": "web_search"}],
}
def test_transient_response_preserves_only_durable_request_messages() -> None:
provider = _provider()
current_message = {"role": "user", "content": "continue"}
@@ -112,14 +112,49 @@ class TestEnforceRoleAlternation:
assert result[1]["content"] is None
assert result[2]["role"] == "tool"
def test_non_string_content_uses_latest(self):
def test_consecutive_user_messages_preserve_text_before_multimodal_content(self):
image = {
"type": "image_url",
"image_url": {"url": "data:image/png;base64,aW1hZ2U="},
}
msgs = [
{"role": "user", "content": [{"type": "text", "text": "A"}]},
{"role": "user", "content": "B"},
{"role": "user", "content": "Earlier unanswered question"},
{
"role": "user",
"content": [image, {"type": "text", "text": "The error is here"}],
},
]
result = LLMProvider._enforce_role_alternation(msgs)
assert len(result) == 1
assert result[0]["content"] == "B"
assert result == [{
"role": "user",
"content": [
{"type": "text", "text": "Earlier unanswered question"},
image,
{"type": "text", "text": "The error is here"},
],
}]
def test_consecutive_user_messages_preserve_multimodal_content_before_text(self):
image = {
"type": "image_url",
"image_url": {"url": "data:image/png;base64,aW1hZ2U="},
}
msgs = [
{
"role": "user",
"content": [image, {"type": "text", "text": "First question"}],
},
{"role": "user", "content": "Follow-up detail"},
]
result = LLMProvider._enforce_role_alternation(msgs)
assert result == [{
"role": "user",
"content": [
image,
{"type": "text", "text": "First question"},
{"type": "text", "text": "Follow-up detail"},
],
}]
def test_original_messages_not_mutated(self):
msgs = [
+504
View File
@@ -0,0 +1,504 @@
from __future__ import annotations
import base64
import json
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from types import SimpleNamespace
import httpx
import pytest
from nanobot.providers.oauth_model_catalog import (
OAuthModelCatalog,
get_oauth_model_catalog,
invalidate_oauth_model_catalog,
)
from nanobot.providers.openai_codex_provider import (
DEFAULT_OPENAI_CODEX_MODELS_URL,
OPENAI_CODEX_CATALOG_CLIENT_VERSION,
)
from nanobot.providers.registry import ProviderModelSpec
from nanobot.providers.xai_grok_provider import DEFAULT_XAI_GROK_MODELS_URL
from nanobot.providers.xai_oauth import XAIToken
@pytest.fixture(autouse=True)
def _clear_oauth_catalogs() -> None:
for provider in ("openai_codex", "xai_grok", "github_copilot"):
invalidate_oauth_model_catalog(provider)
yield
for provider in ("openai_codex", "xai_grok", "github_copilot"):
invalidate_oauth_model_catalog(provider)
def _fallback_model() -> ProviderModelSpec:
return ProviderModelSpec(id="provider/fallback", label="Fallback")
def test_xai_catalog_fetches_remote_models_and_reuses_capability_metadata(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
original_client = httpx.Client
captured: dict[str, object] = {}
payload = (
base64.urlsafe_b64encode(
json.dumps({"sub": "user-42", "email": "user@example.com"}).encode()
)
.decode()
.rstrip("=")
)
token = XAIToken(
access=f"header.{payload}.signature",
refresh="refresh-token",
expires=int(time.time() * 1000) + 3_600_000,
account_id="user@example.com",
)
def handler(request: httpx.Request) -> httpx.Response:
captured["request"] = request
return httpx.Response(
200,
json={
"data": [
{
"id": "grok-4.6",
"name": "Grok 4.6",
"description": "Latest frontier model",
"owned_by": "xAI",
"context_window": 500_000,
"supports_backend_search": True,
"reasoning_efforts": [
{"value": "xhigh"},
{"value": "high"},
{"value": "low"},
],
},
{
"id": "grok-next",
"_meta": {
"name": "Grok Next",
"context_window": 750_000,
"reasoning_efforts": ["high", "low"],
},
},
]
},
request=request,
)
def fake_client(**kwargs: object) -> httpx.Client:
captured["kwargs"] = kwargs
return original_client(
transport=httpx.MockTransport(handler),
timeout=kwargs["timeout"],
follow_redirects=kwargs["follow_redirects"],
)
monkeypatch.setattr(
"nanobot.providers.xai_grok_provider.get_xai_oauth_storage_path",
lambda: tmp_path / "auth" / "xai.json",
)
monkeypatch.setattr(
"nanobot.providers.xai_grok_provider.get_xai_oauth_login_status",
lambda: token,
)
monkeypatch.setattr(
"nanobot.providers.xai_grok_provider.get_xai_oauth_token",
lambda **_kwargs: token,
)
monkeypatch.setattr("nanobot.providers.xai_grok_provider.httpx.Client", fake_client)
catalog = get_oauth_model_catalog("xai_grok")
assert catalog.source == "remote"
assert [model.id for model in catalog.models] == [
"xai-grok/grok-4.6",
"xai-grok/grok-next",
]
grok = catalog.find("grok-4.6")
assert grok is not None
assert grok.description == "Latest frontier model"
assert grok.context_window == 500_000
assert grok.reasoning_efforts == ("xhigh", "high", "low")
assert grok.supports_backend_search is True
next_model = catalog.find("xai-grok/grok-next")
assert next_model is not None
assert next_model.label == "Grok Next"
assert next_model.context_window == 750_000
assert next_model.reasoning_efforts == ("high", "low")
request = captured["request"]
assert isinstance(request, httpx.Request)
assert str(request.url) == DEFAULT_XAI_GROK_MODELS_URL
assert request.headers["Authorization"] == f"Bearer {token.access}"
assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
assert request.headers["x-userid"] == "user-42"
assert request.headers["x-email"] == "user@example.com"
assert captured["kwargs"] == {"timeout": 10.0, "follow_redirects": False}
assert get_oauth_model_catalog("xai_grok").source == "cache"
def test_openai_codex_catalog_uses_account_catalog_and_filters_hidden_models(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
original_client = httpx.Client
captured: dict[str, object] = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["request"] = request
return httpx.Response(
200,
json={
"models": [
{
"slug": "gpt-new",
"display_name": "GPT New",
"description": "New model",
"context_window": 300_000,
"priority": 2,
"visibility": "list",
"supported_reasoning_levels": [
{"effort": "low"},
{"effort": "high"},
],
},
{
"slug": "gpt-first",
"display_name": "GPT First",
"priority": 1,
},
{
"slug": "internal-model",
"display_name": "Internal",
"visibility": "hide",
"priority": 0,
},
]
},
request=request,
)
def fake_client(**kwargs: object) -> httpx.Client:
captured["kwargs"] = kwargs
return original_client(
transport=httpx.MockTransport(handler),
timeout=kwargs["timeout"],
follow_redirects=kwargs["follow_redirects"],
)
class Storage:
def load(self) -> SimpleNamespace:
return SimpleNamespace(access="secret", account_id="account-42")
def get_token_path(self) -> Path:
return tmp_path / "auth" / "openai-codex.json"
monkeypatch.setattr(
"nanobot.providers.openai_codex_provider.FileTokenStorage",
lambda **_kwargs: Storage(),
)
monkeypatch.setattr(
"nanobot.providers.openai_codex_provider.get_codex_token",
lambda **_kwargs: SimpleNamespace(access="secret", account_id="account-42"),
)
monkeypatch.setattr("nanobot.providers.openai_codex_provider.httpx.Client", fake_client)
catalog = get_oauth_model_catalog("openai_codex")
assert catalog.source == "remote"
assert [model.id for model in catalog.models] == [
"openai-codex/gpt-first",
"openai-codex/gpt-new",
]
assert catalog.models[1].context_window == 300_000
assert catalog.models[1].reasoning_efforts == ("low", "high")
request = captured["request"]
assert isinstance(request, httpx.Request)
assert request.url.copy_with(query=None) == httpx.URL(DEFAULT_OPENAI_CODEX_MODELS_URL)
assert request.url.params["client_version"] == OPENAI_CODEX_CATALOG_CLIENT_VERSION
assert request.headers["Authorization"] == "Bearer secret"
assert request.headers["chatgpt-account-id"] == "account-42"
def test_github_copilot_catalog_only_lists_compatible_chat_models(
monkeypatch: pytest.MonkeyPatch,
tmp_path: Path,
) -> None:
original_client = httpx.Client
captured: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
captured.append(request)
if request.url.path.endswith("/copilot_internal/v2/token"):
return httpx.Response(
200,
json={
"token": "copilot-secret",
"endpoints": {"api": "https://api.individual.githubcopilot.com"},
},
request=request,
)
return httpx.Response(
200,
json={
"data": [
{
"id": "claude-sonnet",
"name": "Claude Sonnet",
"model_picker_enabled": True,
"policy": {"state": "enabled"},
"supported_endpoints": ["/chat/completions"],
"capabilities": {
"supports": {"reasoning_effort": ["low", "high"]},
"limits": {"max_context_window_tokens": 200_000},
},
},
{
"id": "gpt-5.4-mini",
"name": "GPT-5.4 Mini",
"model_picker_enabled": True,
"supported_endpoints": ["/responses"],
},
{
"id": "unknown-responses-only",
"name": "Unknown Responses only",
"model_picker_enabled": True,
"supported_endpoints": ["/responses"],
},
{
"id": "disabled",
"model_picker_enabled": True,
"policy": {"state": "disabled"},
"supported_endpoints": ["/chat/completions"],
},
]
},
request=request,
)
def fake_client(**kwargs: object) -> httpx.Client:
return original_client(
transport=httpx.MockTransport(handler),
timeout=kwargs["timeout"],
follow_redirects=kwargs["follow_redirects"],
)
class Storage:
def load(self) -> SimpleNamespace:
return SimpleNamespace(access="github-secret", account_id="octocat")
def get_token_path(self) -> Path:
return tmp_path / "auth" / "github-copilot.json"
monkeypatch.setattr(
"nanobot.providers.github_copilot_provider.get_storage",
lambda: Storage(),
)
monkeypatch.setattr("nanobot.providers.github_copilot_provider.httpx.Client", fake_client)
catalog = get_oauth_model_catalog("github_copilot")
assert catalog.source == "remote"
assert [model.id for model in catalog.models] == [
"github-copilot/claude-sonnet",
"github-copilot/gpt-5.4-mini",
]
assert catalog.models[0].context_window == 200_000
assert catalog.models[0].reasoning_efforts == ("low", "high")
assert len(captured) == 2
assert captured[0].headers["Authorization"] == "token github-secret"
assert captured[1].headers["Authorization"] == "Bearer copilot-secret"
assert str(captured[1].url) == "https://api.individual.githubcopilot.com/models"
assert get_oauth_model_catalog("github_copilot").source == "cache"
assert get_oauth_model_catalog(
"github_copilot",
proxy="http://proxy.example:8080",
).source == "remote"
assert len(captured) == 4
def test_catalog_single_flights_concurrent_refreshes() -> None:
calls = 0
calls_lock = threading.Lock()
barrier = threading.Barrier(8)
def fetch(_proxy: str | None) -> tuple[ProviderModelSpec, ...]:
nonlocal calls
with calls_lock:
calls += 1
time.sleep(0.05)
return (ProviderModelSpec(id="provider/remote", label="Remote"),)
catalog = OAuthModelCatalog(fallback_models=(_fallback_model(),), fetch=fetch)
def get_catalog(_index: int):
barrier.wait()
return catalog.get(cache_key="shared")
with ThreadPoolExecutor(max_workers=8) as pool:
results = list(pool.map(get_catalog, range(8)))
assert calls == 1
assert {result.models[0].id for result in results} == {"provider/remote"}
assert [result.source for result in results].count("remote") == 1
assert [result.source for result in results].count("cache") == 7
def test_catalog_invalidation_discards_an_inflight_account_refresh() -> None:
started = threading.Event()
release = threading.Event()
identity = ["old-account"]
def fetch(_proxy: str | None) -> tuple[ProviderModelSpec, ...]:
current = identity[0]
if current == "old-account":
started.set()
assert release.wait(timeout=2)
return (ProviderModelSpec(id=f"provider/{current}", label=current),)
catalog = OAuthModelCatalog(fallback_models=(_fallback_model(),), fetch=fetch)
with ThreadPoolExecutor(max_workers=2) as pool:
old_future = pool.submit(catalog.get, cache_key="old-key")
assert started.wait(timeout=2)
identity[0] = "new-account"
catalog.invalidate()
new_future = pool.submit(catalog.get, cache_key="new-key")
new_result = new_future.result(timeout=2)
release.set()
old_result = old_future.result(timeout=2)
assert old_result.source == "fallback"
assert new_result.models[0].id == "provider/new-account"
identity[0] = "old-account"
assert catalog.get(cache_key="old-key").models[0].id == "provider/old-account"
def test_catalog_bounds_failure_only_keys() -> None:
calls = 0
def fetch(_proxy: str | None) -> tuple[ProviderModelSpec, ...]:
nonlocal calls
calls += 1
raise httpx.ConnectError("offline")
catalog = OAuthModelCatalog(
fallback_models=(_fallback_model(),),
fetch=fetch,
max_entries=2,
)
for key in ("one", "two", "three"):
assert catalog.get(cache_key=key).source == "fallback"
assert calls == 3
assert catalog.get(cache_key="one").source == "fallback"
assert calls == 4
def test_catalog_returns_stale_then_negative_caches_refresh_failure() -> None:
now = [0.0]
calls = 0
def fetch(_proxy: str | None) -> tuple[ProviderModelSpec, ...]:
nonlocal calls
calls += 1
if calls > 1:
raise httpx.ConnectError("offline")
return (ProviderModelSpec(id="provider/remote", label="Remote"),)
catalog = OAuthModelCatalog(
fallback_models=(_fallback_model(),),
fetch=fetch,
fresh_ttl_s=10,
stale_ttl_s=100,
failure_ttl_s=30,
monotonic=lambda: now[0],
wall_clock=lambda: 123.0,
)
assert catalog.get(cache_key="one").source == "remote"
now[0] = 11
stale = catalog.get(cache_key="one")
assert stale.source == "stale"
assert stale.models[0].id == "provider/remote"
assert catalog.get(cache_key="one").source == "stale"
assert calls == 2
now[0] = 101
fallback = catalog.get(cache_key="one")
assert fallback.source == "fallback"
assert fallback.models[0].id == "provider/fallback"
assert calls == 3
@pytest.mark.parametrize(
"failure",
[
httpx.ConnectError("offline"),
ValueError("invalid JSON"),
httpx.HTTPStatusError(
"unauthorized",
request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
response=httpx.Response(401),
),
httpx.HTTPStatusError(
"rate limited",
request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
response=httpx.Response(429),
),
httpx.HTTPStatusError(
"upstream failure",
request=httpx.Request("GET", DEFAULT_XAI_GROK_MODELS_URL),
response=httpx.Response(503),
),
],
)
def test_catalog_falls_back_for_remote_failures(failure: Exception) -> None:
calls = 0
def fetch(_proxy: str | None) -> tuple[ProviderModelSpec, ...]:
nonlocal calls
calls += 1
raise failure
catalog = OAuthModelCatalog(
fallback_models=(_fallback_model(),),
fetch=fetch,
failure_ttl_s=30,
)
first = catalog.get(cache_key="one")
second = catalog.get(cache_key="one")
assert first.source == "fallback"
assert second.source == "fallback"
assert first.models == (_fallback_model(),)
assert calls == 1
def test_catalog_treats_empty_remote_list_as_failure_and_can_be_invalidated() -> None:
calls = 0
def fetch(_proxy: str | None) -> tuple[ProviderModelSpec, ...]:
nonlocal calls
calls += 1
return () if calls == 1 else (ProviderModelSpec(id="provider/new", label="New"),)
catalog = OAuthModelCatalog(
fallback_models=(_fallback_model(),),
fetch=fetch,
failure_ttl_s=30,
)
assert catalog.get(cache_key="one").source == "fallback"
catalog.invalidate()
refreshed = catalog.get(cache_key="one")
assert refreshed.source == "remote"
assert refreshed.models[0].id == "provider/new"
assert calls == 2
-6
View File
@@ -63,9 +63,3 @@ def test_explicit_provider_import_still_works(monkeypatch) -> None:
finally:
monkeypatch.undo()
setattr(sys.modules["nanobot"], "providers", original_package)
def test_openai_codex_supports_progress_deltas() -> None:
from nanobot.providers.openai_codex_provider import OpenAICodexProvider
assert OpenAICodexProvider.supports_progress_deltas is True
+228 -119
View File
@@ -1,6 +1,5 @@
from __future__ import annotations
import base64
import json
import time
from types import SimpleNamespace
@@ -12,21 +11,19 @@ import pytest
from nanobot.config.schema import Config
from nanobot.providers.base import LLMUsage
from nanobot.providers.factory import make_provider
from nanobot.providers.registry import find_by_name
from nanobot.providers.oauth_model_catalog import OAuthModelCatalogSnapshot
from nanobot.providers.registry import ProviderModelSpec, find_by_name
from nanobot.providers.xai_grok_provider import (
DEFAULT_XAI_GROK_MODEL,
DEFAULT_XAI_GROK_MODELS_URL,
XAIGrokProvider,
_bounded_error_body,
_build_headers,
_build_model_headers,
_build_reasoning_options,
_build_xai_http_error,
_fetch_xai_model_capabilities,
_parse_xai_model_capabilities,
_request_xai,
_xai_error_response,
_XAIHTTPError,
_XAIIncompleteHostedToolError,
)
@@ -51,22 +48,41 @@ def _mock_model_capabilities(
*,
supports_backend_search: bool,
) -> None:
async def fake_fetch(*_args, **_kwargs):
return {"grok-4.5": supports_backend_search}
def fake_catalog(*_args, **_kwargs):
return OAuthModelCatalogSnapshot(
models=(
ProviderModelSpec(
id="xai-grok/grok-4.5",
label="Grok 4.5",
supports_backend_search=supports_backend_search,
),
ProviderModelSpec(
id="xai-grok/grok-4.6",
label="Grok 4.6",
supports_backend_search=supports_backend_search,
),
),
source="remote",
fetched_at=1,
)
monkeypatch.setattr(
"nanobot.providers.xai_grok_provider._fetch_xai_model_capabilities",
fake_fetch,
"nanobot.providers.xai_grok_provider.get_xai_grok_model_catalog",
fake_catalog,
)
def test_xai_grok_registry_exposes_curated_x_search_model() -> None:
def test_xai_grok_registry_exposes_curated_x_search_models() -> None:
spec = find_by_name("xai_grok")
assert spec is not None
assert spec.is_oauth is True
assert spec.backend == "xai_grok"
assert spec.builtin_models[0].id == DEFAULT_XAI_GROK_MODEL
assert [model.id for model in spec.builtin_models] == [
"xai-grok/grok-4.6",
"xai-grok/grok-4.5",
]
assert spec.builtin_models[0].context_window == 500000
assert "when supported" in spec.builtin_models[0].description
@@ -117,7 +133,7 @@ async def test_provider_injects_hosted_x_search_and_required_proxy_headers(monke
assert response.content == "answer [[1]](https://x.com/example/status/1)"
url, headers, body = calls[0]
assert url == "https://cli-chat-proxy.grok.com/v1/responses"
assert body["model"] == "grok-4.5"
assert body["model"] == "grok-4.6"
assert body["tools"] == [
{
"type": "function",
@@ -132,12 +148,13 @@ async def test_provider_injects_hosted_x_search_and_required_proxy_headers(monke
assert body["stream_tool_calls"] is True
assert body["reasoning"] == {"summary": "concise", "effort": "high"}
assert body["store"] is False
assert body["max_turns"] == 5
assert headers["Authorization"] == "Bearer subscription-token"
assert headers["X-XAI-Token-Auth"] == "xai-grok-cli"
assert headers["x-authenticateresponse"] == "authenticate-response"
assert headers["x-grok-client-identifier"] == "nanobot"
assert headers["x-grok-client-mode"] == "headless"
assert headers["x-grok-model-override"] == "grok-4.5"
assert headers["x-grok-model-override"] == "grok-4.6"
@pytest.mark.asyncio
@@ -147,7 +164,7 @@ async def test_explicit_parameterized_x_search_is_preserved_without_catalog_look
_mock_token(monkeypatch)
bodies: list[dict[str, Any]] = []
async def unexpected_catalog_lookup(*_args, **_kwargs):
def unexpected_catalog_lookup(*_args, **_kwargs):
raise AssertionError("explicit raw tools must not depend on model catalog metadata")
async def fake_request(_url, _headers, body, **_kwargs):
@@ -155,7 +172,7 @@ async def test_explicit_parameterized_x_search_is_preserved_without_catalog_look
return "ok", [], "stop", {}, None
monkeypatch.setattr(
"nanobot.providers.xai_grok_provider._fetch_xai_model_capabilities",
"nanobot.providers.xai_grok_provider.get_xai_grok_model_catalog",
unexpected_catalog_lookup,
)
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
@@ -164,10 +181,12 @@ async def test_explicit_parameterized_x_search_is_preserved_without_catalog_look
"allowed_x_handles": ["nanobot_ai"],
"from_date": "2026-01-01",
}
provider = XAIGrokProvider(extra_body={
"parallel_tool_calls": False,
"tools": [hosted_tool, {"type": "code_interpreter", "container": "auto"}],
})
provider = XAIGrokProvider(
extra_body={
"parallel_tool_calls": False,
"tools": [hosted_tool, {"type": "code_interpreter", "container": "auto"}],
}
)
response = await provider.chat(
[{"role": "user", "content": "search"}],
@@ -210,7 +229,7 @@ async def test_explicit_empty_tools_disables_catalog_lookup_and_hosted_tool(monk
_mock_token(monkeypatch)
bodies: list[dict[str, Any]] = []
async def unexpected_catalog_lookup(*_args, **_kwargs):
def unexpected_catalog_lookup(*_args, **_kwargs):
raise AssertionError("explicitly disabled X Search must not fetch model capabilities")
async def fake_request(_url, _headers, body, **_kwargs):
@@ -218,7 +237,7 @@ async def test_explicit_empty_tools_disables_catalog_lookup_and_hosted_tool(monk
return "ok", [], "stop", {}, None
monkeypatch.setattr(
"nanobot.providers.xai_grok_provider._fetch_xai_model_capabilities",
"nanobot.providers.xai_grok_provider.get_xai_grok_model_catalog",
unexpected_catalog_lookup,
)
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
@@ -226,23 +245,28 @@ async def test_explicit_empty_tools_disables_catalog_lookup_and_hosted_tool(monk
response = await provider.chat(
[{"role": "user", "content": "hello"}],
tools=[{
"type": "function",
"function": {
"name": "read_file",
"description": "Read a file",
"parameters": {"type": "object"},
},
}],
tools=[
{
"type": "function",
"function": {
"name": "read_file",
"description": "Read a file",
"parameters": {"type": "object"},
},
}
],
)
assert response.content == "ok"
assert bodies[0]["tools"] == [{
"type": "function",
"name": "read_file",
"description": "Read a file",
"parameters": {"type": "object"},
}]
assert bodies[0]["tools"] == [
{
"type": "function",
"name": "read_file",
"description": "Read a file",
"parameters": {"type": "object"},
}
]
assert "max_turns" not in bodies[0]
@pytest.mark.asyncio
@@ -281,35 +305,8 @@ async def test_provider_keeps_local_x_search_when_model_does_not_support_hosted_
"parameters": {"type": "object"},
}
]
@pytest.mark.asyncio
async def test_provider_fails_closed_and_caches_model_catalog_failure(monkeypatch) -> None:
_mock_token(monkeypatch)
fetch_calls = 0
bodies: list[dict[str, Any]] = []
async def failing_fetch(*_args, **_kwargs):
nonlocal fetch_calls
fetch_calls += 1
raise httpx.ConnectError("catalog unavailable")
async def fake_request(_url, _headers, body, **_kwargs):
bodies.append(body)
return "ok", [], "stop", {}, None
monkeypatch.setattr(
"nanobot.providers.xai_grok_provider._fetch_xai_model_capabilities",
failing_fetch,
)
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
provider = XAIGrokProvider()
await provider.chat([{"role": "user", "content": "first"}])
await provider.chat([{"role": "user", "content": "second"}])
assert fetch_calls == 1
assert all({"type": "x_search"} not in body["tools"] for body in bodies)
assert "max_turns" not in bodies[0]
assert bodies[0]["instructions"] == ""
@pytest.mark.asyncio
@@ -395,7 +392,10 @@ async def test_factory_builds_xai_provider_and_applies_explicit_body_overrides(m
"providers": {
"xaiGrok": {
"proxy": "http://127.0.0.1:7890",
"extraBody": {"parallel_tool_calls": False},
"extraBody": {
"parallel_tool_calls": False,
"max_turns": 2,
},
}
},
}
@@ -408,6 +408,7 @@ async def test_factory_builds_xai_provider_and_applies_explicit_body_overrides(m
assert provider.proxy == "http://127.0.0.1:7890"
assert response.content == "ok"
assert bodies[0]["parallel_tool_calls"] is False
assert bodies[0]["max_turns"] == 2
assert {"type": "x_search"} in bodies[0]["tools"]
@@ -527,75 +528,183 @@ async def test_raw_response_request_streams_hosted_x_search_lifecycle(monkeypatc
assert "large hosted result" not in json.dumps(tool_events)
def test_model_capabilities_follow_upstream_aliases_and_default_to_disabled() -> None:
capabilities = _parse_xai_model_capabilities(
{
"data": [
{"id": "grok-4.5", "supportsBackendSearch": False},
{
"model": "grok-search",
"supports_backend_search": True,
},
{
"modelId": "grok-meta",
"_meta": {"supportsBackendSearch": True},
},
{"id": "grok-unknown"},
]
}
)
assert capabilities == {
"grok-4.5": False,
"grok-search": True,
"grok-meta": True,
"grok-unknown": False,
}
@pytest.mark.asyncio
async def test_model_capability_request_uses_subscription_headers(monkeypatch) -> None:
async def test_raw_response_request_streams_official_x_search_lifecycle(monkeypatch) -> None:
original_client = httpx.AsyncClient
captured: dict[str, Any] = {}
events = [
{
"type": "response.output_item.added",
"item": {
"type": "x_search_call",
"id": "x-search-1",
"status": "in_progress",
"action": {"query": "nanobot oauth"},
},
},
{
"type": "response.output_item.done",
"item": {
"type": "x_search_call",
"id": "x-search-1",
"status": "completed",
"action": {"query": "nanobot oauth"},
},
},
{
"type": "response.completed",
"response": {"status": "completed", "usage": {}},
},
]
content = "".join(f"data: {json.dumps(event)}\n\n" for event in events)
def handler(request: httpx.Request) -> httpx.Response:
captured["request"] = request
return httpx.Response(
200,
json={"data": [{"id": "grok-search", "supportsBackendSearch": True}]},
request=request,
)
return httpx.Response(200, content=content, request=request)
def fake_client(**kwargs) -> httpx.AsyncClient:
captured["kwargs"] = kwargs
return original_client(
transport=httpx.MockTransport(handler),
timeout=kwargs["timeout"],
follow_redirects=kwargs["follow_redirects"],
)
monkeypatch.setattr("nanobot.providers.xai_grok_provider.httpx.AsyncClient", fake_client)
payload = base64.urlsafe_b64encode(
json.dumps({"sub": "user-42", "email": "user@example.com"}).encode()
).decode().rstrip("=")
access_token = f"header.{payload}.signature"
headers = _build_model_headers(_token(access_token))
tool_events: list[dict[str, Any]] = []
capabilities = await _fetch_xai_model_capabilities(
DEFAULT_XAI_GROK_MODELS_URL,
headers,
await _request_xai(
"https://cli-chat-proxy.grok.com/v1/responses",
_build_headers("secret", "grok-4.6"),
{"model": "grok-4.6", "tools": [{"type": "x_search"}]},
on_tool_call_delta=lambda event: _append(tool_events, event),
)
request = captured["request"]
assert isinstance(request, httpx.Request)
assert request.method == "GET"
assert str(request.url) == DEFAULT_XAI_GROK_MODELS_URL
assert request.headers["Authorization"] == f"Bearer {access_token}"
assert request.headers["X-XAI-Token-Auth"] == "xai-grok-cli"
assert request.headers["x-userid"] == "user-42"
assert request.headers["x-email"] == "user@example.com"
assert captured["kwargs"] == {"timeout": 10.0, "follow_redirects": False}
assert capabilities == {"grok-search": True}
assert [(event["phase"], event["name"]) for event in tool_events] == [
("start", "x_search"),
("end", "x_search"),
]
assert tool_events[-1]["result"] == {"status": "completed"}
@pytest.mark.asyncio
async def test_raw_response_rejects_unfinished_hosted_tool_and_closes_progress(
monkeypatch,
) -> None:
original_client = httpx.AsyncClient
events = [
{
"type": "response.custom_tool_call_input.done",
"item_id": "x-search-1",
"input": '{"query":"nanobot oauth"}',
},
{"type": "response.output_text.delta", "delta": "I will keep searching."},
{
"type": "response.completed",
"response": {
"status": "completed",
"usage": {"input_tokens": 8, "output_tokens": 4, "total_tokens": 12},
},
},
]
content = "".join(f"data: {json.dumps(event)}\n\n" for event in events)
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, content=content, request=request)
def fake_client(**kwargs) -> httpx.AsyncClient:
return original_client(
transport=httpx.MockTransport(handler),
timeout=kwargs["timeout"],
)
monkeypatch.setattr("nanobot.providers.xai_grok_provider.httpx.AsyncClient", fake_client)
tool_events: list[dict[str, Any]] = []
with pytest.raises(_XAIIncompleteHostedToolError) as caught:
await _request_xai(
"https://cli-chat-proxy.grok.com/v1/responses",
_build_headers("secret", "grok-4.6"),
{"model": "grok-4.6", "tools": [{"type": "x_search"}]},
on_tool_call_delta=lambda event: _append(tool_events, event),
)
assert caught.value.usage == LLMUsage.reported(input_tokens=8, output_tokens=4)
assert [event["phase"] for event in tool_events] == ["start", "error"]
assert "before this hosted tool completed" in tool_events[-1]["error"]
@pytest.mark.asyncio
async def test_provider_recovers_unfinished_hosted_tool_once_and_preserves_usage(
monkeypatch,
) -> None:
_mock_token(monkeypatch)
_mock_model_capabilities(monkeypatch, supports_backend_search=True)
attempts = 0
request_ids: list[str] = []
streamed: list[str] = []
recovered: list[bool] = []
first_usage = LLMUsage.reported(input_tokens=10, output_tokens=2)
second_usage = LLMUsage.reported(input_tokens=11, output_tokens=4)
async def fake_request(_url, headers, body, **kwargs):
nonlocal attempts
attempts += 1
request_ids.append(headers["x-grok-req-id"])
assert body["max_turns"] == 5
if attempts == 1:
await kwargs["on_content_delta"]("I will keep searching.")
raise _XAIIncompleteHostedToolError(
[{"name": "x_search", "call_id": "search-1"}],
usage=first_usage,
)
await kwargs["on_content_delta"]("Final researched answer.")
return "Final researched answer.", [], "stop", second_usage, None
async def on_recover() -> None:
recovered.append(True)
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
provider = XAIGrokProvider()
response = await provider.chat_stream_with_retry(
[{"role": "user", "content": "Search X"}],
on_content_delta=lambda delta: _append(streamed, delta),
on_stream_recover=on_recover,
)
assert attempts == 2
assert len(set(request_ids)) == 2
assert recovered == [True]
assert streamed == ["I will keep searching.", "Final researched answer."]
assert response.content == "Final researched answer."
assert response.usage == first_usage + second_usage
@pytest.mark.asyncio
async def test_provider_preserves_usage_when_hosted_tool_recovery_also_fails(
monkeypatch,
) -> None:
_mock_token(monkeypatch)
_mock_model_capabilities(monkeypatch, supports_backend_search=True)
attempts = 0
usage = LLMUsage.reported(input_tokens=10, output_tokens=2)
async def fake_request(*_args, **_kwargs):
nonlocal attempts
attempts += 1
raise _XAIIncompleteHostedToolError(
[{"name": "x_search", "call_id": f"search-{attempts}"}],
usage=usage,
)
monkeypatch.setattr("nanobot.providers.xai_grok_provider._request_xai", fake_request)
provider = XAIGrokProvider()
response = await provider.chat_stream_with_retry(
[{"role": "user", "content": "Search X"}],
on_stream_recover=lambda: _append([], True),
)
assert attempts == 2
assert response.finish_reason == "error"
assert response.usage == usage + usage
@pytest.mark.asyncio
@@ -57,9 +57,41 @@ def test_loaded_corrupt_offset_keeps_messages(tmp_path: Path):
def test_valid_offset_is_preserved():
session = _session(10, 4)
assert session.last_consolidated == 4
assert session.last_archived == 4
assert len(session.get_history()) == 8
def test_last_archived_field_migrates_with_legacy_alias(tmp_path: Path):
manager = SessionManager(tmp_path)
path = manager._get_session_path("chan:chat")
path.parent.mkdir(parents=True, exist_ok=True)
messages = [
{"role": "user", "content": "first"},
{"role": "assistant", "content": "second"},
]
path.write_text(
"\n".join([
json.dumps({
"_type": "metadata",
"key": "chan:chat",
"metadata": {},
"last_archived": 1,
}),
*(json.dumps(message) for message in messages),
]) + "\n",
encoding="utf-8",
)
session = manager.get_or_create("chan:chat")
assert session.last_archived == 1
assert session.last_consolidated == 1
manager.save(session)
metadata = json.loads(path.read_text(encoding="utf-8").splitlines()[0])
assert metadata["last_archived"] == 1
assert metadata["last_consolidated"] == 1
def test_loaded_null_metadata_becomes_empty_dict(tmp_path: Path):
"""Session jsonl metadata:null must load as {} so agent .pop/.get work."""
manager = SessionManager(tmp_path)
+3 -6
View File
@@ -141,14 +141,13 @@ def test_internal_continuation_requires_budget_boundary_and_queue():
)
def test_save_skip_matches_prefix_when_current_message_merged():
def test_save_skip_matches_prefix_when_current_message_was_persisted():
skip = _save_skip_for_turn(
message_metadata=None,
initial_message_count=2, # [system, merged user]
history_count=1,
initial_message_count=3, # [system, history user, current user]
input_persisted_early=True,
)
assert skip == 2
assert skip == 3
def test_save_skip_unchanged_for_standalone_current_message():
@@ -156,12 +155,10 @@ def test_save_skip_unchanged_for_standalone_current_message():
assert _save_skip_for_turn(
message_metadata=None,
initial_message_count=3,
history_count=1,
input_persisted_early=True,
) == 3
assert _save_skip_for_turn(
message_metadata=None,
initial_message_count=3,
history_count=1,
input_persisted_early=False,
) == 2
+37
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import json
from unittest.mock import MagicMock
import pytest
@@ -11,6 +12,7 @@ from nanobot.agent.tools.message import MessageTool
from nanobot.agent.tools.spawn import SpawnTool
from nanobot.cron.service import CronService
from nanobot.providers.base import GenerationSettings, LLMProvider
from nanobot.runtime_context import RUNTIME_CONTEXT_INPUT_META, RuntimeContextBlock
from nanobot.session.keys import UNIFIED_SESSION_KEY
from nanobot.utils.llm_runtime import LLMRuntime
@@ -299,6 +301,41 @@ async def test_webui_cron_tool_uses_origin_session_when_unified_enabled(tmp_path
assert jobs[0].payload.origin_metadata == {"webui": True}
@pytest.mark.asyncio
async def test_cron_tool_snapshots_only_persistable_request_metadata(tmp_path) -> None:
"""Live runtime context must not poison a persisted WebUI cron job."""
store_path = tmp_path / "jobs.json"
service = CronService(store_path)
tool = CronTool(service)
await service.start()
try:
with request_context(
RequestContext(
channel="websocket",
chat_id="chat-123",
metadata={
"webui": True,
RUNTIME_CONTEXT_INPUT_META: [
RuntimeContextBlock(source="webui_quote", content="quoted reply")
],
"opaque": object(),
},
session_key=UNIFIED_SESSION_KEY,
)
):
result = await tool.execute(action="add", message="standup", every_seconds=300)
assert result.startswith("Created job")
jobs = service.list_jobs()
assert len(jobs) == 1
assert jobs[0].payload.origin_metadata == {"webui": True}
raw = json.loads(store_path.read_text(encoding="utf-8"))
assert raw["jobs"][0]["payload"]["originMetadata"] == {"webui": True}
finally:
service.stop()
@pytest.mark.asyncio
async def test_cron_tool_preserves_thread_scoped_session_key(tmp_path) -> None:
"""Channel-provided thread session keys should remain the cron owner."""
+40 -19
View File
@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from nanobot.agent.context import TranscriptInput
from nanobot.agent.loop import AgentLoop
from nanobot.agent.tools.message import MessageTool
from nanobot.bus.events import InboundMessage, OutboundMessage
@@ -24,7 +25,12 @@ class TestMessageToolSuppressLogic:
"""Final reply suppressed only when message tool sends to the same target."""
@pytest.mark.asyncio
async def test_suppress_when_sent_to_same_target(self, tmp_path: Path) -> None:
@pytest.mark.parametrize("ephemeral", [False, True])
async def test_suppress_when_sent_to_same_target(
self,
tmp_path: Path,
ephemeral: bool,
) -> None:
loop = _make_loop(tmp_path)
tool_call = ToolCallRequest(
id="call1", name="message",
@@ -43,7 +49,7 @@ class TestMessageToolSuppressLogic:
mt.set_send_callback(AsyncMock(side_effect=lambda m: sent.append(m)))
msg = InboundMessage(channel="feishu", sender_id="user1", chat_id="chat123", content="Send")
result = await loop._process_message(msg)
result = await loop._process_message(msg, ephemeral=ephemeral)
assert len(sent) == 1
assert result is None # suppressed
@@ -87,6 +93,34 @@ class TestMessageToolSuppressLogic:
assert result is not None
assert "Hello" in result.content
@pytest.mark.asyncio
async def test_internal_message_check_keeps_final_response(self, tmp_path: Path) -> None:
loop = _make_loop(tmp_path)
tool_call = ToolCallRequest(
id="call1", name="message",
arguments={"content": "all clear", "channel": "feishu", "chat_id": "chat123"},
)
calls = iter([
LLMResponse(content="", tool_calls=[tool_call]),
LLMResponse(content="Heartbeat summary", tool_calls=[]),
])
loop.provider.chat_with_retry = AsyncMock(side_effect=lambda *a, **kw: next(calls))
loop.tools.get_definitions = MagicMock(return_value=[])
mt = loop.tools.get("message")
assert isinstance(mt, MessageTool)
token = mt.set_suppress_delivery(True)
try:
msg = InboundMessage(
channel="feishu", sender_id="user1", chat_id="chat123", content="Check",
)
result = await loop._process_message(msg)
finally:
mt.reset_suppress_delivery(token)
assert result is not None
assert result.content == "Heartbeat summary"
@pytest.mark.asyncio
async def test_injected_followup_with_message_tool_does_not_emit_empty_fallback(
self, tmp_path: Path
@@ -145,7 +179,9 @@ class TestMessageToolSuppressLogic:
progress.append((content, tool_hint))
result = await loop._run_agent_loop(
[], runtime=loop.llm_runtime(), on_progress=on_progress
TranscriptInput(history=[], current_message=None),
runtime=loop.llm_runtime(),
on_progress=on_progress,
)
assert result.final_content == "Done"
@@ -154,22 +190,7 @@ class TestMessageToolSuppressLogic:
('read foo.txt', True),
]
class TestMessageToolTurnTracking:
def test_sent_in_turn_tracks_same_target(self) -> None:
tool = MessageTool()
from nanobot.agent.tools.context import RequestContext, request_context
with request_context(RequestContext(channel="feishu", chat_id="chat1")):
assert not tool._sent_in_turn
tool._sent_in_turn = True
assert tool._sent_in_turn
def test_start_turn_resets(self) -> None:
tool = MessageTool()
tool._sent_in_turn = True
tool.start_turn()
assert not tool._sent_in_turn
class TestMessageToolSchema:
def test_schema_discourages_current_chat_replies(self) -> None:
tool = MessageTool()
+60
View File
@@ -183,6 +183,66 @@ async def test_rate_limit_is_per_source_session_and_uses_a_rolling_minute(
)
@pytest.mark.asyncio
async def test_rate_limit_releases_expired_source_state_and_keeps_recent_sources(
tmp_path: Path,
) -> None:
sessions = SessionManager(tmp_path)
_persist(
sessions,
"websocket:a",
"websocket:b",
"websocket:c",
"websocket:target",
)
now = 0.0
tool = SendSessionMessageTool(
sessions=sessions,
bus=MessageBus(),
max_messages_per_minute=2,
clock=lambda: now,
)
target = _handle(sessions, "websocket:target").name
for source in ("websocket:a", "websocket:b"):
await tool.enqueue(
source_session_key=source,
target_handle=target,
content="initial",
expect_reply=False,
)
now = 30.0
await tool.enqueue(
source_session_key="websocket:a",
target_handle=target,
content="recent",
expect_reply=False,
)
now = 61.0
await tool.enqueue(
source_session_key="websocket:c",
target_handle=target,
content="trigger cleanup",
expect_reply=False,
)
assert set(tool._sent_at) == {"websocket:a", "websocket:c"}
await tool.enqueue(
source_session_key="websocket:a",
target_handle=target,
content="within rolling window",
expect_reply=False,
)
with pytest.raises(SessionMessageError, match="rate limit"):
await tool.enqueue(
source_session_key="websocket:a",
target_handle=target,
content="over limit",
expect_reply=False,
)
@pytest.mark.asyncio
async def test_reply_timeout_injects_a_user_input_back_into_the_source(
tmp_path: Path,

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