mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-04 08:28:36 +00:00
Merge remote-tracking branch 'origin/main' into codex/unified-extension-platform
# Conflicts: # docs/configuration.md
This commit is contained in:
commit
d68857bb2d
@ -14,9 +14,9 @@ Shell execution (`ExecTool`, `agent/tools/shell.py`) also respects `restrict_to_
|
||||
|
||||
## SSRF Protection
|
||||
|
||||
All outbound HTTP requests from agent tools must pass through `validate_url_target` (`security/network.py`). By default it blocks loopback, RFC1918 private addresses, CGNAT ranges, link-local ranges, and cloud metadata endpoints (including `169.254.169.254`).
|
||||
All outbound HTTP requests from agent tools must pass through the shared URL guards in `security/network.py` (`validate_url_target` or `resolve_url_target`). By default they block loopback, RFC1918 private addresses, CGNAT ranges, link-local ranges, and cloud metadata endpoints (including `169.254.169.254`).
|
||||
|
||||
The only escape hatch is `configure_ssrf_whitelist(cidrs)`, which reads from `config.tools.ssrf_whitelist` at load time.
|
||||
For direct requests, the only escape hatch is `configure_ssrf_whitelist(cidrs)`, which reads from `config.tools.ssrf_whitelist` at load time. An explicitly configured `providers.<name>.proxy` is a separate user-authorized trust boundary for provider requests and provider-returned image URL downloads. Those downloads still reject malformed URLs and locally identifiable private/internal targets on every redirect, but hostnames unavailable to local DNS are delegated to the trusted proxy. The user-selected proxy owns final DNS resolution and network egress policy.
|
||||
|
||||
HTTP/SSE MCP transports are part of this boundary: validate configured MCP URLs before probing or constructing clients, and validate each outgoing HTTP request before redirects are followed. Local/private HTTP MCP endpoints are allowed only through the explicit SSRF whitelist. Stdio MCP servers are not part of the HTTP SSRF path.
|
||||
|
||||
|
||||
@ -618,7 +618,7 @@ async def send(self, msg: OutboundMessage) -> None:
|
||||
await self._send_message(msg.chat_id, msg.content, media=msg.media)
|
||||
```
|
||||
|
||||
Tool hints are off by default for most channels. Users can enable them globally or per channel:
|
||||
Tool hints are on by default. Users can disable them globally or per channel:
|
||||
|
||||
```json
|
||||
{
|
||||
@ -626,7 +626,7 @@ Tool hints are off by default for most channels. Users can enable them globally
|
||||
"sendToolHints": true,
|
||||
"webhook": {
|
||||
"enabled": true,
|
||||
"sendToolHints": true
|
||||
"sendToolHints": false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -1558,7 +1558,7 @@ Global settings that apply to all channels. Configure under the `channels` secti
|
||||
{
|
||||
"channels": {
|
||||
"sendProgress": true,
|
||||
"sendToolHints": false,
|
||||
"sendToolHints": true,
|
||||
"extractDocumentText": true,
|
||||
"sendMaxRetries": 3,
|
||||
"telegram": {
|
||||
@ -1571,7 +1571,7 @@ Global settings that apply to all channels. Configure under the `channels` secti
|
||||
| Setting | Default | Description |
|
||||
|---------|---------|-------------|
|
||||
| `sendProgress` | `true` | Stream agent's text progress to the channel |
|
||||
| `sendToolHints` | `false` | Stream tool-call hints (e.g. `read_file("…")`) |
|
||||
| `sendToolHints` | `true` | Stream tool-call hints (e.g. `read_file("…")`) |
|
||||
| `showReasoning` | `true` | Allow channels to surface model reasoning/thinking content (DeepSeek-R1 `reasoning_content`, Anthropic `thinking_blocks`, inline `<think>` tags). Reasoning flows as a dedicated stream with `_reasoning_delta` / `_reasoning_end` markers — channels override `send_reasoning_delta` / `send_reasoning_end` to render in-place updates. Even with `true`, channels without those overrides stay no-op silently. Currently surfaced on CLI and WebSocket/WebUI (italic shimmer header, auto-collapses after the stream ends); Telegram / Slack / Discord / Feishu / WeChat / Matrix / Mattermost keep the base no-op until their bubble UI is adapted. Independent of `sendProgress`. |
|
||||
| `extractDocumentText` | `true` | Extract supported document/text attachments into the model prompt. PDF, DOCX, XLSX, and PPTX readers are included in the standard installation. Set to `false` to keep document content out of the prompt and include attachment path references instead. |
|
||||
| `sendMaxRetries` | `3` | Max delivery attempts per outbound message, including the initial send (0-10 configured, minimum 1 actual attempt) |
|
||||
@ -1584,10 +1584,11 @@ Global settings that apply to all channels. Configure under the `channels` secti
|
||||
{
|
||||
"channels": {
|
||||
"sendProgress": true,
|
||||
"sendToolHints": false,
|
||||
"sendToolHints": true,
|
||||
"telegram": {
|
||||
"enabled": true,
|
||||
"sendProgress": false
|
||||
"sendProgress": false,
|
||||
"sendToolHints": false
|
||||
},
|
||||
"websocket": {
|
||||
"enabled": true,
|
||||
@ -1997,7 +1998,9 @@ For API keys, tokens, and other secrets, see [Environment Variables for Secrets]
|
||||
| `tools.exec.timeout` | `60` | Default hard timeout in seconds for shell commands. Config values may exceed the per-call tool cap; set `0` to disable the hard timeout for trusted long-running commands. |
|
||||
| `tools.exec.pathPrepend` | `""` | Extra directories to prepend to `PATH` when running shell commands. Use this when configured tools should win executable lookup precedence, such as a Python virtual environment's `bin` or `Scripts` directory. |
|
||||
| `tools.exec.pathAppend` | `""` | Extra directories to append to `PATH` when running shell commands (e.g. `/usr/sbin` for `ufw`). |
|
||||
| `tools.webuiAllowRemotePackageInstall` | `false` | When `false`, the WebUI can install optional support or extension packages only from a browser opened on the same machine as nanobot. Set to `true` only when a trusted remote admin may place packages into this installation. Remotely installed extensions remain untrusted and inactive; trust, permission, activation, disabling, and removal stay local-only. |
|
||||
| `tools.exec.sandboxRoBinds` | `[]` | Extra absolute paths to read-only bind into the `"bwrap"` sandbox with `--ro-bind-try`, such as `/home/user/.local/bin` or `/home/user/.cargo/bin` when those paths are also in `pathPrepend`/`pathAppend`. These roots are also accepted by the shell absolute-path guard only while bwrap is active. Bind only directories whose contents are safe for agent commands to read; paths equal to or containing the active workspace are ignored so they cannot uncover its masked parent directory. |
|
||||
| `tools.exec.sandboxRwBinds` | `[]` | Extra absolute paths to read-write bind into the `"bwrap"` sandbox with `--bind-try`, for trusted tool caches or scratch directories. Use sparingly: paths listed here are intentionally writable by shell commands inside the sandbox. Paths equal to or containing the active workspace are ignored. |
|
||||
| `tools.webuiAllowRemotePackageInstall` | `false` | When `false`, the WebUI can install optional support or extension packages only from a browser opened on the same machine as nanobot. Set to `true` only when a trusted remote admin may install packages into this environment. |
|
||||
| `tools.ssrfWhitelist` | `[]` | CIDR ranges exempted from the shared SSRF guard used by web fetches and HTTP/SSE MCP connections. Prefer exact host CIDRs such as `192.168.1.50/32`; broad ranges increase SSRF exposure. |
|
||||
| `channels.*.allowFrom` | omitted | Access control per channel. Omit to use pairing-only mode; set `["*"]` to allow everyone; or list specific user IDs. See [Pairing](#pairing) for details. |
|
||||
|
||||
@ -2158,7 +2161,8 @@ When a user is idle for longer than a configured threshold, nanobot **proactivel
|
||||
{
|
||||
"agents": {
|
||||
"defaults": {
|
||||
"idleCompactAfterMinutes": 15
|
||||
"idleCompactAfterMinutes": 15,
|
||||
"idleCompactCheckIntervalSeconds": 60
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -2167,11 +2171,12 @@ When a user is idle for longer than a configured threshold, nanobot **proactivel
|
||||
| Option | Default | Description |
|
||||
|--------|---------|-------------|
|
||||
| `agents.defaults.idleCompactAfterMinutes` | `15` | Minutes of idle time before auto-compaction starts. Set to `0` to disable. The default is close to a typical LLM KV cache expiry window, so stale sessions get compacted before the user returns. |
|
||||
| `agents.defaults.idleCompactCheckIntervalSeconds` | `60` | Minimum number of seconds between scans for idle sessions. Set to `0` to scan on every idle tick (~1 s). |
|
||||
|
||||
`sessionTtlMinutes` remains accepted as a legacy alias for backward compatibility, but `idleCompactAfterMinutes` is the preferred config key going forward.
|
||||
|
||||
How it works:
|
||||
1. **Idle detection**: On each idle tick (~1 s), checks all sessions for expiration.
|
||||
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.
|
||||
|
||||
@ -70,6 +70,9 @@ Provider settings reuse normal provider config fields:
|
||||
| `providers.<name>.apiBase` | Optional custom base URL |
|
||||
| `providers.<name>.extraHeaders` | Headers merged into provider requests |
|
||||
| `providers.<name>.extraBody` | Extra JSON fields merged into provider request bodies |
|
||||
| `providers.<name>.proxy` | Explicit trusted HTTP proxy for provider requests and returned image URL downloads |
|
||||
|
||||
For providers that return image URLs, direct downloads use DNS pinning. When an explicit provider `proxy` is configured, nanobot rejects malformed URLs and locally identifiable private/internal targets on the initial URL and every redirect. Hostnames unavailable to local DNS are delegated to that trusted proxy, which owns final DNS resolution and network egress. Process-wide proxy environment variables are not used for these downloads.
|
||||
|
||||
Both camelCase and snake_case config keys are accepted, but docs use camelCase to match `config.json`.
|
||||
|
||||
|
||||
@ -25,6 +25,7 @@ class AgentHookContext:
|
||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||
streamed_content: bool = False
|
||||
streamed_reasoning: bool = False
|
||||
stream_continues_current_message: bool = False
|
||||
final_content: str | None = None
|
||||
stop_reason: str | None = None
|
||||
error: str | None = None
|
||||
|
||||
@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import dataclasses
|
||||
import inspect
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
@ -73,7 +74,7 @@ from nanobot.session.goal_state import (
|
||||
sustained_goal_active,
|
||||
)
|
||||
from nanobot.session.history_visibility import HIDDEN_HISTORY_META
|
||||
from nanobot.session.keys import UNIFIED_SESSION_KEY
|
||||
from nanobot.session.keys import UNIFIED_SESSION_KEY, remember_last_channel
|
||||
from nanobot.session.manager import (
|
||||
Session,
|
||||
SessionManager,
|
||||
@ -296,6 +297,7 @@ class AgentLoop:
|
||||
runtime_model_publisher: Callable[[str, str | None], None] | None = None,
|
||||
restart_mode: str = "auto",
|
||||
local_trigger_store: Any | None = None,
|
||||
idle_compact_check_interval_seconds: int = 0,
|
||||
):
|
||||
from nanobot.config.schema import ToolsConfig
|
||||
|
||||
@ -444,6 +446,8 @@ class AgentLoop:
|
||||
consolidator=self.consolidator,
|
||||
session_ttl_minutes=session_ttl_minutes,
|
||||
)
|
||||
self._idle_compact_check_interval_s = idle_compact_check_interval_seconds
|
||||
self._next_idle_compact_check_at = time.monotonic()
|
||||
if model_preset:
|
||||
self.set_model_preset(model_preset, publish_update=False)
|
||||
self._register_default_tools(provider_snapshot_loader=provider_snapshot_loader)
|
||||
@ -499,6 +503,7 @@ class AgentLoop:
|
||||
unified_session=defaults.unified_session,
|
||||
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),
|
||||
@ -745,14 +750,23 @@ class AgentLoop:
|
||||
self,
|
||||
ctx: TurnContext,
|
||||
) -> list[RuntimeContextBlock]:
|
||||
tools = ctx.tools or self.tools
|
||||
assert ctx.request_context is not None
|
||||
return await self._resolve_runtime_context_for_request(
|
||||
ctx.request_context,
|
||||
ctx.tools or self.tools,
|
||||
)
|
||||
|
||||
async def _resolve_runtime_context_for_request(
|
||||
self,
|
||||
request: RequestContext,
|
||||
tools: ToolRegistry,
|
||||
) -> list[RuntimeContextBlock]:
|
||||
providers = [
|
||||
*tools.get_runtime_context_providers(),
|
||||
*self._runtime_context_providers,
|
||||
]
|
||||
assert ctx.request_context is not None
|
||||
blocks = runtime_context_blocks_from_metadata(ctx.request_context.metadata)
|
||||
blocks.extend(await resolve_runtime_context(providers, ctx.request_context))
|
||||
blocks = runtime_context_blocks_from_metadata(request.metadata)
|
||||
blocks.extend(await resolve_runtime_context(providers, request))
|
||||
return blocks
|
||||
|
||||
async def _dispatch_command_inline(
|
||||
@ -789,6 +803,27 @@ class AgentLoop:
|
||||
return UNIFIED_SESSION_KEY
|
||||
return msg.session_key
|
||||
|
||||
def _remember_unified_session_route(
|
||||
self,
|
||||
session: Session,
|
||||
msg: InboundMessage,
|
||||
*,
|
||||
is_user_turn: bool,
|
||||
) -> None:
|
||||
"""Remember the latest user-facing route for unified-session delivery."""
|
||||
if (
|
||||
not self._unified_session
|
||||
or session.key != UNIFIED_SESSION_KEY
|
||||
or not is_user_turn
|
||||
or msg.channel in {"cli", "system"}
|
||||
or msg.sender_id == "subagent"
|
||||
):
|
||||
return
|
||||
_, automation_metadata = automation_history_overrides(msg.metadata)
|
||||
if automation_metadata:
|
||||
return
|
||||
remember_last_channel(session.metadata, msg.channel, msg.chat_id)
|
||||
|
||||
@staticmethod
|
||||
def _replay_token_budget(runtime: LLMRuntime) -> int:
|
||||
"""Derive a token budget for session history replay from the context window."""
|
||||
@ -830,9 +865,9 @@ class AgentLoop:
|
||||
"""Run the agent iteration loop.
|
||||
|
||||
*on_stream*: called with each content delta during streaming.
|
||||
*on_stream_end(resuming)*: called when a streaming session finishes.
|
||||
``resuming=True`` means tool calls follow (spinner should restart);
|
||||
``resuming=False`` means this is the final response.
|
||||
*on_stream_end(resuming, merge_next)*: called when a streaming session finishes.
|
||||
``resuming=True`` means the active turn continues. ``merge_next=True`` means
|
||||
the next text segment belongs to the same user-visible assistant message.
|
||||
|
||||
Returns (final_content, tools_used, messages, stop_reason, had_injections).
|
||||
"""
|
||||
@ -855,7 +890,7 @@ class AgentLoop:
|
||||
if pending_queue is None:
|
||||
return []
|
||||
|
||||
def _to_user_message(pending_msg: InboundMessage) -> dict[str, Any]:
|
||||
async def _to_user_message(pending_msg: InboundMessage) -> dict[str, Any]:
|
||||
content = pending_msg.content
|
||||
media = pending_msg.media if pending_msg.media else None
|
||||
if media:
|
||||
@ -864,6 +899,31 @@ class AgentLoop:
|
||||
user_content = self.context._build_user_content(content, media)
|
||||
row: dict[str, Any] = {"role": "user", "content": user_content}
|
||||
metadata = pending_msg.metadata if isinstance(pending_msg.metadata, dict) else {}
|
||||
if pending_msg.channel != "system":
|
||||
scope = self.workspace_scopes.for_turn(
|
||||
channel=pending_msg.channel,
|
||||
message_metadata=metadata,
|
||||
session_metadata=session.metadata if session is not None else None,
|
||||
)
|
||||
pending_request = RequestContext(
|
||||
channel=pending_msg.channel,
|
||||
chat_id=pending_msg.chat_id,
|
||||
message_id=metadata.get("message_id"),
|
||||
session_key=active_session_key,
|
||||
original_user_text=pending_msg.content,
|
||||
runtime=runtime,
|
||||
metadata=dict(metadata),
|
||||
sender_id=pending_msg.sender_id,
|
||||
turn_id=request_ctx.turn_id,
|
||||
workspace=scope.project_path,
|
||||
)
|
||||
blocks = await self._resolve_runtime_context_for_request(
|
||||
pending_request,
|
||||
effective_tools,
|
||||
)
|
||||
row["content"], marker = append_runtime_context(user_content, blocks)
|
||||
if marker is not None:
|
||||
row["_meta"] = {RUNTIME_CONTEXT_MESSAGE_META: marker}
|
||||
if (
|
||||
pending_msg.sender_id == "subagent"
|
||||
and metadata.get("injected_event") == "subagent_result"
|
||||
@ -880,7 +940,7 @@ class AgentLoop:
|
||||
items: list[dict[str, Any]] = []
|
||||
while len(items) < limit:
|
||||
try:
|
||||
items.append(_to_user_message(pending_queue.get_nowait()))
|
||||
items.append(await _to_user_message(pending_queue.get_nowait()))
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
@ -898,10 +958,10 @@ class AgentLoop:
|
||||
session.key,
|
||||
)
|
||||
return items
|
||||
items.append(_to_user_message(msg))
|
||||
items.append(await _to_user_message(msg))
|
||||
while len(items) < limit:
|
||||
try:
|
||||
items.append(_to_user_message(pending_queue.get_nowait()))
|
||||
items.append(await _to_user_message(pending_queue.get_nowait()))
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
@ -1014,12 +1074,29 @@ class AgentLoop:
|
||||
# Push final content through stream so streaming channels (e.g. Feishu)
|
||||
# update the card instead of leaving it empty.
|
||||
if on_stream and on_stream_end and should_stream:
|
||||
await on_stream(result.final_content or "")
|
||||
stream_content = (
|
||||
result.pending_stream_content
|
||||
if result.pending_stream_content is not None
|
||||
else result.final_content or ""
|
||||
)
|
||||
await on_stream(stream_content)
|
||||
await on_stream_end(resuming=False)
|
||||
elif result.stop_reason == "error":
|
||||
logger.error("LLM returned error: {}", (result.final_content or "")[:200])
|
||||
return result.final_content, result.tools_used, result.messages, result.stop_reason, result.had_injections
|
||||
|
||||
def _check_expired_sessions_if_due(self) -> None:
|
||||
"""Scan idle sessions no more often than the configured interval."""
|
||||
now = time.monotonic()
|
||||
if now < self._next_idle_compact_check_at:
|
||||
return
|
||||
self._next_idle_compact_check_at = now + self._idle_compact_check_interval_s
|
||||
self.auto_compact.check_expired(
|
||||
self._schedule_background,
|
||||
self.runtime_for_session,
|
||||
active_session_keys=self._pending_queues.keys(),
|
||||
)
|
||||
|
||||
async def run(self) -> None:
|
||||
"""Run the agent loop, dispatching messages as tasks to stay responsive to /stop."""
|
||||
self._running = True
|
||||
@ -1031,11 +1108,7 @@ class AgentLoop:
|
||||
try:
|
||||
msg = await asyncio.wait_for(self.bus.consume_inbound(), timeout=1.0)
|
||||
except asyncio.TimeoutError:
|
||||
self.auto_compact.check_expired(
|
||||
self._schedule_background,
|
||||
self.runtime_for_session,
|
||||
active_session_keys=self._pending_queues.keys(),
|
||||
)
|
||||
self._check_expired_sessions_if_due()
|
||||
continue
|
||||
except asyncio.CancelledError:
|
||||
# Preserve real task cancellation so shutdown can complete cleanly.
|
||||
@ -1161,6 +1234,14 @@ class AgentLoop:
|
||||
for _, coordinator in self._automation_turn_coordinators:
|
||||
coordinator.complete(msg, error=asyncio.CancelledError())
|
||||
logger.info("Task cancelled for session {}", session_key)
|
||||
try:
|
||||
await delivery.abort_stream()
|
||||
except Exception:
|
||||
logger.debug(
|
||||
"Could not close stream for cancelled session {}",
|
||||
session_key,
|
||||
exc_info=True,
|
||||
)
|
||||
# Preserve partial context from the interrupted turn so
|
||||
# the user does not lose tool results and assistant
|
||||
# messages accumulated before /stop. The checkpoint was
|
||||
@ -1330,6 +1411,19 @@ class AgentLoop:
|
||||
if ctx.on_stream is not None:
|
||||
stream_callback = ctx.on_stream
|
||||
stream_end_callback = ctx.on_stream_end
|
||||
stream_end_accepts_merge_next = False
|
||||
if stream_end_callback is not None:
|
||||
try:
|
||||
stream_end_signature = inspect.signature(stream_end_callback)
|
||||
stream_end_accepts_merge_next = (
|
||||
"merge_next" in stream_end_signature.parameters
|
||||
or any(
|
||||
parameter.kind is inspect.Parameter.VAR_KEYWORD
|
||||
for parameter in stream_end_signature.parameters.values()
|
||||
)
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
segment_streamed_content = False
|
||||
|
||||
async def _tracked_stream(delta: str) -> None:
|
||||
@ -1338,12 +1432,19 @@ class AgentLoop:
|
||||
segment_streamed_content = True
|
||||
await stream_callback(delta)
|
||||
|
||||
async def _tracked_stream_end(*, resuming: bool = False) -> None:
|
||||
async def _tracked_stream_end(
|
||||
*,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
nonlocal segment_streamed_content
|
||||
ctx.streamed_content = segment_streamed_content
|
||||
segment_streamed_content = False
|
||||
if stream_end_callback is not None:
|
||||
await stream_end_callback(resuming=resuming)
|
||||
if merge_next and stream_end_accepts_merge_next:
|
||||
await stream_end_callback(resuming=resuming, merge_next=True)
|
||||
else:
|
||||
await stream_end_callback(resuming=resuming)
|
||||
|
||||
ctx.on_stream = _tracked_stream
|
||||
ctx.on_stream_end = _tracked_stream_end
|
||||
@ -1456,6 +1557,11 @@ class AgentLoop:
|
||||
# ensure it exists in case this handler is invoked independently.
|
||||
if ctx.session is None:
|
||||
ctx.session = self.sessions.get_or_create(ctx.session_key)
|
||||
self._remember_unified_session_route(
|
||||
ctx.session,
|
||||
msg,
|
||||
is_user_turn=ctx.original_user_text is not None,
|
||||
)
|
||||
await ctx.delivery.started()
|
||||
if ctx.kind is TurnKind.USER:
|
||||
self.workspace_scopes.persist_message_scope(ctx.session, msg)
|
||||
|
||||
@ -43,13 +43,33 @@ if TYPE_CHECKING:
|
||||
# MemoryStore — pure file I/O layer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class DreamRunProgress:
|
||||
"""Track tool failures that make a nominally completed Dream run unsafe to advance."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.had_tool_errors = False
|
||||
|
||||
async def __call__(
|
||||
self,
|
||||
*_args: Any,
|
||||
tool_events: list[dict[str, Any]] | None = None,
|
||||
**_kwargs: Any,
|
||||
) -> None:
|
||||
if any(
|
||||
isinstance(event, dict) and event.get("phase") == "error"
|
||||
for event in tool_events or ()
|
||||
):
|
||||
self.had_tool_errors = True
|
||||
|
||||
|
||||
class MemoryStore:
|
||||
"""Pure file I/O for memory files: MEMORY.md, history.jsonl, SOUL.md, USER.md."""
|
||||
|
||||
_DEFAULT_MAX_HISTORY = 1000
|
||||
# Durable files whose real working-tree delta grounds Dream commit messages
|
||||
# and the cursor-advance gate. Deliberately excludes memory/.dream_cursor so
|
||||
# that advancing the cursor itself is never mistaken for a productive edit.
|
||||
# Durable files whose real working-tree delta grounds Dream commit messages.
|
||||
# Deliberately excludes memory/.dream_cursor so progress bookkeeping never
|
||||
# appears as a durable-memory edit in the audit record.
|
||||
_DREAM_CONTENT_PATHS = ("SOUL.md", "USER.md", "memory/MEMORY.md")
|
||||
# Per-file cap when embedding current contents into the Dream prompt. The
|
||||
# durable files are tiny in practice (~5 KB total), but a runaway file must
|
||||
@ -586,8 +606,7 @@ class MemoryStore:
|
||||
"""Structured summary of uncommitted changes to the durable memory files.
|
||||
|
||||
Returns "" when git is unavailable or no content file changed. This is
|
||||
the ground-truth input for diff-grounded Dream commit messages and for
|
||||
gating cursor advance on real edits (never on LLM self-report).
|
||||
the ground-truth input for diff-grounded Dream commit messages.
|
||||
"""
|
||||
if not self._git.is_initialized():
|
||||
return ""
|
||||
@ -636,10 +655,18 @@ class MemoryStore:
|
||||
return tools
|
||||
|
||||
@staticmethod
|
||||
def dream_run_completed(resp: object | None) -> bool:
|
||||
"""Return True only when an ephemeral Dream agent turn completed cleanly."""
|
||||
def dream_run_completed(
|
||||
resp: object | None,
|
||||
*,
|
||||
had_tool_errors: bool = False,
|
||||
) -> bool:
|
||||
"""Return True only when a Dream turn completed without tool failures."""
|
||||
metadata = getattr(resp, "metadata", None)
|
||||
return isinstance(metadata, dict) and metadata.get("_stop_reason") == "completed"
|
||||
return (
|
||||
not had_tool_errors
|
||||
and isinstance(metadata, dict)
|
||||
and metadata.get("_stop_reason") == "completed"
|
||||
)
|
||||
|
||||
# -- message formatting utility ------------------------------------------
|
||||
|
||||
|
||||
@ -85,7 +85,13 @@ class AgentProgressHook(AgentHook):
|
||||
async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None:
|
||||
await self.emit_reasoning_end()
|
||||
if self._on_stream_end:
|
||||
await self._on_stream_end(resuming=resuming)
|
||||
kwargs: dict[str, bool] = {"resuming": resuming}
|
||||
if (
|
||||
context.stream_continues_current_message
|
||||
and self._on_progress_accepts(self._on_stream_end, "merge_next")
|
||||
):
|
||||
kwargs["merge_next"] = True
|
||||
await self._on_stream_end(**kwargs)
|
||||
self._stream_buf = ""
|
||||
self._think_extractor.reset()
|
||||
|
||||
|
||||
@ -19,6 +19,11 @@ from nanobot.agent.context_governance import (
|
||||
from nanobot.agent.hook import AgentHook, AgentHookContext, AgentRunHookContext
|
||||
from nanobot.agent.tools.registry import ToolRegistry, is_tool_error_result
|
||||
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
|
||||
from nanobot.runtime_context import (
|
||||
RUNTIME_CONTEXT_MESSAGE_META,
|
||||
detach_runtime_context,
|
||||
reattach_runtime_context,
|
||||
)
|
||||
from nanobot.session.history_visibility import is_hidden_history_message
|
||||
from nanobot.utils.helpers import (
|
||||
IncrementalThinkExtractor,
|
||||
@ -55,6 +60,18 @@ _MAX_LENGTH_RECOVERIES = 3
|
||||
_MAX_INJECTIONS_PER_TURN = 3
|
||||
_MAX_INJECTION_CYCLES = 5
|
||||
|
||||
|
||||
def _restore_outer_whitespace(content: str, original: str | None) -> str:
|
||||
"""Restore boundary whitespace stripped while cleaning one recovered segment."""
|
||||
if not original:
|
||||
return content
|
||||
leading_size = len(original) - len(original.lstrip())
|
||||
trailing_size = len(original) - len(original.rstrip())
|
||||
leading = original[:leading_size]
|
||||
trailing = original[-trailing_size:] if trailing_size else ""
|
||||
return f"{leading}{content}{trailing}"
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class AgentRunSpec:
|
||||
"""Configuration for a single agent execution."""
|
||||
@ -96,6 +113,8 @@ class AgentRunResult:
|
||||
error: str | None = None
|
||||
tool_events: list[dict[str, str]] = field(default_factory=list)
|
||||
had_injections: bool = False
|
||||
# Terminal tail to emit when the preceding final-content prefix was already streamed.
|
||||
pending_stream_content: str | None = None
|
||||
|
||||
|
||||
class AgentRunner:
|
||||
@ -137,10 +156,51 @@ class AgentRunner:
|
||||
and not is_hidden_history_message(messages[-1])
|
||||
):
|
||||
merged = dict(messages[-1])
|
||||
merged["content"] = cls._merge_message_content(
|
||||
merged.get("content"),
|
||||
injection.get("content"),
|
||||
left_meta = merged.get("_meta")
|
||||
right_meta = injection.get("_meta")
|
||||
left_marker = (
|
||||
left_meta.get(RUNTIME_CONTEXT_MESSAGE_META)
|
||||
if isinstance(left_meta, dict)
|
||||
else None
|
||||
)
|
||||
right_marker = (
|
||||
right_meta.get(RUNTIME_CONTEXT_MESSAGE_META)
|
||||
if isinstance(right_meta, dict)
|
||||
else None
|
||||
)
|
||||
detached_left = (
|
||||
detach_runtime_context(merged.get("content"), left_marker)
|
||||
if isinstance(left_marker, dict)
|
||||
else (merged.get("content"), [], [])
|
||||
)
|
||||
detached_right = (
|
||||
detach_runtime_context(injection.get("content"), right_marker)
|
||||
if isinstance(right_marker, dict)
|
||||
else (injection.get("content"), [], [])
|
||||
)
|
||||
if detached_left is not None and detached_right is not None:
|
||||
left_content, left_sources, left_blocks = detached_left
|
||||
right_content, right_sources, right_blocks = detached_right
|
||||
merged_content = cls._merge_message_content(left_content, right_content)
|
||||
context_blocks = [*left_blocks, *right_blocks]
|
||||
if context_blocks:
|
||||
merged_content, marker = reattach_runtime_context(
|
||||
merged_content,
|
||||
[*left_sources, *right_sources],
|
||||
context_blocks,
|
||||
)
|
||||
internal_meta = dict(left_meta) if isinstance(left_meta, dict) else {}
|
||||
if isinstance(right_meta, dict):
|
||||
for key, value in right_meta.items():
|
||||
internal_meta.setdefault(key, value)
|
||||
internal_meta[RUNTIME_CONTEXT_MESSAGE_META] = marker
|
||||
merged["_meta"] = internal_meta
|
||||
merged["content"] = merged_content
|
||||
else:
|
||||
merged["content"] = cls._merge_message_content(
|
||||
merged.get("content"),
|
||||
injection.get("content"),
|
||||
)
|
||||
messages[-1] = merged
|
||||
continue
|
||||
messages.append(injection)
|
||||
@ -334,10 +394,13 @@ class AgentRunner:
|
||||
# Per-turn throttle for repeated attempts against the same outside target.
|
||||
workspace_violation_counts: dict[str, int] = {}
|
||||
empty_content_retries = 0
|
||||
length_recovery_count = 0
|
||||
# Segments from one uninterrupted length-recovery chain. Tool work or
|
||||
# injected user input starts a new logical answer and clears the chain.
|
||||
length_recovery_parts: list[str] = []
|
||||
had_injections = False
|
||||
injection_cycles = 0
|
||||
compacted_tool_call_ids: set[str] = set()
|
||||
pending_stream_content: str | None = None
|
||||
governance_config = ContextGovernanceConfig(
|
||||
provider=spec.runtime.provider,
|
||||
model=spec.runtime.model,
|
||||
@ -372,6 +435,7 @@ class AgentRunner:
|
||||
context.response = response
|
||||
context.tool_calls = list(response.tool_calls)
|
||||
|
||||
original_content = response.content
|
||||
reasoning_text, cleaned_content = extract_reasoning(
|
||||
response.reasoning_content,
|
||||
response.thinking_blocks,
|
||||
@ -458,6 +522,7 @@ class AgentRunner:
|
||||
)
|
||||
if should_continue:
|
||||
had_injections = True
|
||||
length_recovery_parts.clear()
|
||||
continue
|
||||
break
|
||||
await self._emit_checkpoint(
|
||||
@ -472,7 +537,7 @@ class AgentRunner:
|
||||
},
|
||||
)
|
||||
empty_content_retries = 0
|
||||
length_recovery_count = 0
|
||||
length_recovery_parts.clear()
|
||||
# Checkpoint 1: drain injections after tools, before next LLM call
|
||||
_drained, injection_cycles = await self._try_drain_injections(
|
||||
spec, messages, None, injection_cycles,
|
||||
@ -521,29 +586,50 @@ class AgentRunner:
|
||||
context.response = response
|
||||
context.usage = dict(raw_usage)
|
||||
context.tool_calls = list(response.tool_calls)
|
||||
original_content = response.content
|
||||
clean = hook.finalize_content(context, response.content)
|
||||
|
||||
if response.finish_reason == "length" and not is_blank_text(clean):
|
||||
length_recovery_count += 1
|
||||
if length_recovery_count <= _MAX_LENGTH_RECOVERIES:
|
||||
if len(length_recovery_parts) < _MAX_LENGTH_RECOVERIES:
|
||||
length_recovery_parts.append(
|
||||
_restore_outer_whitespace(clean, original_content)
|
||||
)
|
||||
logger.info(
|
||||
"Output truncated on turn {} for {} ({}/{}); continuing",
|
||||
iteration,
|
||||
spec.session_key or "default",
|
||||
length_recovery_count,
|
||||
len(length_recovery_parts),
|
||||
_MAX_LENGTH_RECOVERIES,
|
||||
)
|
||||
if hook.wants_streaming():
|
||||
context.stream_continues_current_message = True
|
||||
await hook.on_stream_end(context, resuming=True)
|
||||
messages.append(build_assistant_message(
|
||||
clean,
|
||||
reasoning_content=response.reasoning_content,
|
||||
thinking_blocks=response.thinking_blocks,
|
||||
))
|
||||
messages.append(build_length_recovery_message())
|
||||
messages.append(build_length_recovery_message(clean))
|
||||
await hook.after_iteration(context)
|
||||
continue
|
||||
|
||||
# Some streaming providers recover with a complete response but no
|
||||
# content deltas. When an earlier length segment is already visible,
|
||||
# emit this terminal segment into the same stream; otherwise the
|
||||
# regular full response would duplicate the visible prefix.
|
||||
if (
|
||||
length_recovery_parts
|
||||
and hook.wants_streaming()
|
||||
and not context.streamed_content
|
||||
and response.finish_reason != "error"
|
||||
and not is_blank_text(clean)
|
||||
):
|
||||
await hook.on_stream(
|
||||
context,
|
||||
_restore_outer_whitespace(clean, original_content),
|
||||
)
|
||||
context.streamed_content = True
|
||||
|
||||
assistant_message: dict[str, Any] | None = None
|
||||
if response.finish_reason != "error" and not is_blank_text(clean):
|
||||
assistant_message = build_assistant_message(
|
||||
@ -568,6 +654,7 @@ class AgentRunner:
|
||||
await hook.on_stream_end(context, resuming=should_continue)
|
||||
|
||||
if should_continue:
|
||||
length_recovery_parts.clear()
|
||||
await hook.after_iteration(context)
|
||||
continue
|
||||
|
||||
@ -589,6 +676,7 @@ class AgentRunner:
|
||||
)
|
||||
if should_continue:
|
||||
had_injections = True
|
||||
length_recovery_parts.clear()
|
||||
continue
|
||||
break
|
||||
if is_blank_text(clean):
|
||||
@ -606,6 +694,7 @@ class AgentRunner:
|
||||
)
|
||||
if should_continue:
|
||||
had_injections = True
|
||||
length_recovery_parts.clear()
|
||||
continue
|
||||
break
|
||||
|
||||
@ -625,7 +714,13 @@ class AgentRunner:
|
||||
"pending_tool_calls": [],
|
||||
},
|
||||
)
|
||||
final_content = clean
|
||||
if length_recovery_parts:
|
||||
final_content = (
|
||||
"".join(length_recovery_parts)
|
||||
+ _restore_outer_whitespace(clean, original_content)
|
||||
).strip()
|
||||
else:
|
||||
final_content = clean
|
||||
context.final_content = final_content
|
||||
context.stop_reason = stop_reason
|
||||
await hook.after_iteration(context)
|
||||
@ -643,17 +738,25 @@ class AgentRunner:
|
||||
)
|
||||
if drained_after_max_iterations:
|
||||
had_injections = True
|
||||
final_content = None
|
||||
terminal_content = None
|
||||
if spec.finalize_on_max_iterations:
|
||||
final_content = await self._try_finalize_after_max_iterations(
|
||||
terminal_content = await self._try_finalize_after_max_iterations(
|
||||
spec,
|
||||
hook,
|
||||
messages,
|
||||
usage,
|
||||
)
|
||||
if final_content is None:
|
||||
final_content = self._max_iterations_fallback(spec)
|
||||
self._append_final_message(messages, final_content)
|
||||
if terminal_content is None:
|
||||
terminal_content = self._max_iterations_fallback(spec)
|
||||
if length_recovery_parts:
|
||||
terminal_tail = f"\n\n{terminal_content.lstrip()}"
|
||||
final_content = (
|
||||
"".join(length_recovery_parts).rstrip() + terminal_tail
|
||||
).strip()
|
||||
pending_stream_content = terminal_tail
|
||||
else:
|
||||
final_content = terminal_content
|
||||
self._append_final_message(messages, terminal_content)
|
||||
|
||||
return AgentRunResult(
|
||||
final_content=final_content,
|
||||
@ -664,6 +767,7 @@ class AgentRunner:
|
||||
error=error,
|
||||
tool_events=tool_events,
|
||||
had_injections=had_injections,
|
||||
pending_stream_content=pending_stream_content,
|
||||
)
|
||||
|
||||
def _build_request_kwargs(
|
||||
|
||||
@ -315,13 +315,87 @@ def _extract_nullable_branch(options: Any) -> tuple[dict[str, Any], bool] | None
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]:
|
||||
"""Normalize only nullable JSON Schema patterns for tool definitions."""
|
||||
if not isinstance(schema, dict):
|
||||
return {"type": "object", "properties": {}}
|
||||
def _resolve_local_schema_ref(root: dict[str, Any], ref: str) -> Any:
|
||||
"""Resolve a local JSON Pointer without accepting remote references."""
|
||||
if not ref.startswith("#"):
|
||||
raise ValueError("not a local JSON Pointer")
|
||||
|
||||
pointer = urllib.parse.unquote(ref[1:], errors="strict")
|
||||
if not pointer:
|
||||
return root
|
||||
if not pointer.startswith("/"):
|
||||
raise ValueError("not a local JSON Pointer")
|
||||
|
||||
current: Any = root
|
||||
for raw_part in pointer[1:].split("/"):
|
||||
part = raw_part.replace("~1", "/").replace("~0", "~")
|
||||
if isinstance(current, dict):
|
||||
current = current[part]
|
||||
elif isinstance(current, list):
|
||||
current = current[int(part)]
|
||||
else:
|
||||
raise KeyError(part)
|
||||
return current
|
||||
|
||||
|
||||
def _rewrite_local_schema_refs(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Hoist arbitrary local JSON-Pointer refs into provider-compatible ``$defs``."""
|
||||
rewritten_refs: dict[str, str] = {}
|
||||
generated_defs: dict[str, Any] = {}
|
||||
|
||||
def rewrite(value: Any) -> Any:
|
||||
if isinstance(value, list):
|
||||
return [rewrite(item) for item in value]
|
||||
if not isinstance(value, dict):
|
||||
return value
|
||||
|
||||
rewritten = dict(value)
|
||||
ref = rewritten.get("$ref")
|
||||
is_rewritable_ref = False
|
||||
if isinstance(ref, str) and not ref.startswith("#/$defs/"):
|
||||
try:
|
||||
pointer = urllib.parse.unquote(ref[1:], errors="strict")
|
||||
except (UnicodeDecodeError, ValueError):
|
||||
pass
|
||||
else:
|
||||
is_rewritable_ref = ref.startswith("#") and (
|
||||
not pointer or pointer.startswith("/")
|
||||
)
|
||||
if is_rewritable_ref:
|
||||
name = rewritten_refs.get(ref)
|
||||
if name is None:
|
||||
try:
|
||||
target = _resolve_local_schema_ref(schema, ref)
|
||||
except (KeyError, IndexError, TypeError, UnicodeDecodeError, ValueError):
|
||||
logger.warning("MCP tool schema contains an unresolved local $ref: {}", ref)
|
||||
else:
|
||||
assert isinstance(ref, str)
|
||||
name = f"ref_{hashlib.sha256(ref.encode()).hexdigest()[:12]}"
|
||||
existing_defs = schema.get("$defs")
|
||||
while isinstance(existing_defs, dict) and name in existing_defs:
|
||||
name += "_"
|
||||
rewritten_refs[ref] = name
|
||||
# Reserve the name before descending so recursive refs terminate.
|
||||
generated_defs[name] = {}
|
||||
generated_defs[name] = rewrite(target)
|
||||
if name is not None:
|
||||
rewritten["$ref"] = f"#/$defs/{name}"
|
||||
|
||||
return {key: rewrite(item) for key, item in rewritten.items()}
|
||||
|
||||
result = rewrite(schema)
|
||||
if generated_defs:
|
||||
existing_defs = result.get("$defs")
|
||||
result["$defs"] = {
|
||||
**(existing_defs if isinstance(existing_defs, dict) else {}),
|
||||
**generated_defs,
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def _normalize_nullable_schema(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Normalize nullable forms in structural subschemas only."""
|
||||
normalized = dict(schema)
|
||||
|
||||
raw_type = normalized.get("type")
|
||||
if isinstance(raw_type, list):
|
||||
non_null = [item for item in raw_type if item != "null"]
|
||||
@ -339,23 +413,34 @@ def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]:
|
||||
normalized["nullable"] = True
|
||||
break
|
||||
|
||||
if "properties" in normalized and isinstance(normalized["properties"], dict):
|
||||
if isinstance(normalized.get("properties"), dict):
|
||||
normalized["properties"] = {
|
||||
name: _normalize_schema_for_openai(prop) if isinstance(prop, dict) else prop
|
||||
name: _normalize_nullable_schema(prop) if isinstance(prop, dict) else prop
|
||||
for name, prop in normalized["properties"].items()
|
||||
}
|
||||
if isinstance(normalized.get("items"), dict):
|
||||
normalized["items"] = _normalize_nullable_schema(normalized["items"])
|
||||
if isinstance(normalized.get("$defs"), dict):
|
||||
normalized["$defs"] = {
|
||||
name: _normalize_nullable_schema(definition)
|
||||
if isinstance(definition, dict)
|
||||
else definition
|
||||
for name, definition in normalized["$defs"].items()
|
||||
}
|
||||
|
||||
if "items" in normalized and isinstance(normalized["items"], dict):
|
||||
normalized["items"] = _normalize_schema_for_openai(normalized["items"])
|
||||
|
||||
if normalized.get("type") != "object":
|
||||
return normalized
|
||||
|
||||
normalized.setdefault("properties", {})
|
||||
normalized.setdefault("required", [])
|
||||
if normalized.get("type") == "object":
|
||||
normalized.setdefault("properties", {})
|
||||
normalized.setdefault("required", [])
|
||||
return normalized
|
||||
|
||||
|
||||
def _normalize_schema_for_openai(schema: Any) -> dict[str, Any]:
|
||||
"""Normalize MCP JSON Schema patterns for tool definitions."""
|
||||
if not isinstance(schema, dict):
|
||||
return {"type": "object", "properties": {}}
|
||||
return _normalize_nullable_schema(_rewrite_local_schema_refs(schema))
|
||||
|
||||
|
||||
class _MCPWrapperBase(Tool):
|
||||
"""Common reconnect handling for wrappers bound to one MCP server session."""
|
||||
|
||||
|
||||
@ -5,13 +5,54 @@ To add a new backend, implement a function with the signature:
|
||||
and register it in _BACKENDS below.
|
||||
"""
|
||||
|
||||
import os
|
||||
import shlex
|
||||
from pathlib import Path
|
||||
from typing import Iterable
|
||||
|
||||
from nanobot.config.paths import get_media_dir
|
||||
|
||||
|
||||
def _bwrap(command: str, workspace: str, cwd: str) -> str:
|
||||
def _normalize_bind_paths(
|
||||
paths: Iterable[str] | None,
|
||||
*,
|
||||
workspace: Path | None = None,
|
||||
) -> list[str]:
|
||||
out: list[str] = []
|
||||
seen: set[str] = set()
|
||||
for raw in paths or []:
|
||||
value = str(raw).strip()
|
||||
if not value:
|
||||
continue
|
||||
path = Path(os.path.expandvars(value)).expanduser()
|
||||
if not path.is_absolute():
|
||||
continue
|
||||
resolved_path = path.resolve(strict=False)
|
||||
if workspace is not None:
|
||||
try:
|
||||
workspace.relative_to(resolved_path)
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
# A later bind of the workspace or one of its parents could
|
||||
# cover the tmpfs that hides the config directory.
|
||||
continue
|
||||
resolved = str(resolved_path)
|
||||
if resolved in seen:
|
||||
continue
|
||||
seen.add(resolved)
|
||||
out.append(resolved)
|
||||
return out
|
||||
|
||||
|
||||
def _bwrap(
|
||||
command: str,
|
||||
workspace: str,
|
||||
cwd: str,
|
||||
*,
|
||||
sandbox_ro_binds: Iterable[str] | None = None,
|
||||
sandbox_rw_binds: Iterable[str] | None = None,
|
||||
) -> str:
|
||||
"""Wrap command in a bubblewrap sandbox (requires bwrap in container).
|
||||
|
||||
Only the workspace is bind-mounted read-write; its parent dir (which holds
|
||||
@ -51,17 +92,34 @@ def _bwrap(command: str, workspace: str, cwd: str) -> str:
|
||||
"--dir", str(ws), # recreate workspace mount point
|
||||
"--bind", str(ws), str(ws),
|
||||
"--ro-bind-try", str(media), str(media), # read-only access to media
|
||||
"--chdir", sandbox_cwd,
|
||||
"--", "sh", "-c", command,
|
||||
]
|
||||
for p in _normalize_bind_paths(sandbox_ro_binds, workspace=ws):
|
||||
args += ["--ro-bind-try", p, p]
|
||||
for p in _normalize_bind_paths(sandbox_rw_binds, workspace=ws):
|
||||
args += ["--bind-try", p, p]
|
||||
args += ["--chdir", sandbox_cwd, "--", "sh", "-c", command]
|
||||
return shlex.join(args)
|
||||
|
||||
|
||||
_BACKENDS = {"bwrap": _bwrap}
|
||||
|
||||
|
||||
def wrap_command(sandbox: str, command: str, workspace: str, cwd: str) -> str:
|
||||
def wrap_command(
|
||||
sandbox: str,
|
||||
command: str,
|
||||
workspace: str,
|
||||
cwd: str,
|
||||
*,
|
||||
sandbox_ro_binds: Iterable[str] | None = None,
|
||||
sandbox_rw_binds: Iterable[str] | None = None,
|
||||
) -> str:
|
||||
"""Wrap *command* using the named sandbox backend."""
|
||||
if backend := _BACKENDS.get(sandbox):
|
||||
return backend(command, workspace, cwd)
|
||||
return backend(
|
||||
command,
|
||||
workspace,
|
||||
cwd,
|
||||
sandbox_ro_binds=sandbox_ro_binds,
|
||||
sandbox_rw_binds=sandbox_rw_binds,
|
||||
)
|
||||
raise ValueError(f"Unknown sandbox backend {sandbox!r}. Available: {list(_BACKENDS)}")
|
||||
|
||||
@ -84,6 +84,8 @@ class ExecToolConfig(Base):
|
||||
path_prepend: str = ""
|
||||
path_append: str = ""
|
||||
sandbox: str = ""
|
||||
sandbox_ro_binds: list[str] = Field(default_factory=list)
|
||||
sandbox_rw_binds: list[str] = Field(default_factory=list)
|
||||
allowed_env_keys: list[str] = Field(default_factory=list)
|
||||
allow_patterns: list[str] = Field(default_factory=list)
|
||||
deny_patterns: list[str] = Field(default_factory=list)
|
||||
@ -187,6 +189,8 @@ class ExecTool(Tool):
|
||||
sandbox=cfg.sandbox,
|
||||
path_prepend=cfg.path_prepend,
|
||||
path_append=cfg.path_append,
|
||||
sandbox_ro_binds=cfg.sandbox_ro_binds,
|
||||
sandbox_rw_binds=cfg.sandbox_rw_binds,
|
||||
allowed_env_keys=cfg.allowed_env_keys,
|
||||
allow_patterns=cfg.allow_patterns,
|
||||
deny_patterns=cfg.deny_patterns,
|
||||
@ -205,6 +209,8 @@ class ExecTool(Tool):
|
||||
sandbox: str = "",
|
||||
path_prepend: str = "",
|
||||
path_append: str = "",
|
||||
sandbox_ro_binds: list[str] | None = None,
|
||||
sandbox_rw_binds: list[str] | None = None,
|
||||
allowed_env_keys: list[str] | None = None,
|
||||
session_manager: Any | None = None,
|
||||
):
|
||||
@ -237,6 +243,8 @@ class ExecTool(Tool):
|
||||
self.webui_allow_local_service_access = webui_allow_local_service_access
|
||||
self.path_prepend = path_prepend
|
||||
self.path_append = path_append
|
||||
self.sandbox_ro_binds = self._normalize_bind_roots(sandbox_ro_binds)
|
||||
self.sandbox_rw_binds = self._normalize_bind_roots(sandbox_rw_binds)
|
||||
self.allowed_env_keys = allowed_env_keys or []
|
||||
self._session_manager = session_manager or DEFAULT_EXEC_SESSION_MANAGER
|
||||
|
||||
@ -464,7 +472,14 @@ class ExecTool(Tool):
|
||||
)
|
||||
else:
|
||||
workspace = workspace_root or cwd
|
||||
command = wrap_command(self.sandbox, command, workspace, cwd)
|
||||
command = wrap_command(
|
||||
self.sandbox,
|
||||
command,
|
||||
workspace,
|
||||
cwd,
|
||||
sandbox_ro_binds=[str(p) for p in self.sandbox_ro_binds],
|
||||
sandbox_rw_binds=[str(p) for p in self.sandbox_rw_binds],
|
||||
)
|
||||
cwd = str(Path(workspace).resolve())
|
||||
|
||||
effective_timeout = self._resolve_timeout(timeout)
|
||||
@ -794,6 +809,9 @@ class ExecTool(Tool):
|
||||
if workspace_root
|
||||
else None
|
||||
)
|
||||
sandbox_bind_roots = self._active_sandbox_bind_roots(
|
||||
resolved_workspace or cwd_path
|
||||
)
|
||||
|
||||
for raw in self._extract_absolute_paths(cmd):
|
||||
try:
|
||||
@ -817,6 +835,8 @@ class ExecTool(Tool):
|
||||
)
|
||||
if not allowed and resolved_workspace is not None:
|
||||
allowed = is_path_within(p, resolved_workspace)
|
||||
if not allowed and sandbox_bind_roots:
|
||||
allowed = any(is_path_within(p, root) for root in sandbox_bind_roots)
|
||||
if p.is_absolute() and not allowed:
|
||||
return ToolResult.error(
|
||||
"Error: Command blocked by safety guard (path outside working dir)"
|
||||
@ -921,3 +941,38 @@ class ExecTool(Tool):
|
||||
posix_paths = re.findall(r"(?:^|[\s|>='\"])(/[^\s\"'>;|<]+)", command) # POSIX: /absolute only
|
||||
home_paths = re.findall(r"(?:^|[\s>='\"])(~[/+][^\s\"'>;|<]*)", command) # POSIX/Windows home shortcut: ~/ or ~+
|
||||
return win_paths + posix_paths + home_paths
|
||||
|
||||
@staticmethod
|
||||
def _normalize_bind_roots(paths: list[str] | None) -> list[Path]:
|
||||
roots: list[Path] = []
|
||||
seen: set[str] = set()
|
||||
for raw in paths or []:
|
||||
value = str(raw).strip()
|
||||
if not value:
|
||||
continue
|
||||
path = Path(os.path.expandvars(value)).expanduser()
|
||||
if not path.is_absolute():
|
||||
continue
|
||||
with suppress(OSError, RuntimeError, ValueError):
|
||||
resolved = path.resolve(strict=False)
|
||||
key = os.path.normcase(os.fspath(resolved))
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
roots.append(resolved)
|
||||
return roots
|
||||
|
||||
def _active_sandbox_bind_roots(
|
||||
self,
|
||||
workspace_root: Path | None = None,
|
||||
) -> list[Path]:
|
||||
if self.sandbox != "bwrap" or _IS_WINDOWS:
|
||||
return []
|
||||
roots = [*self.sandbox_ro_binds, *self.sandbox_rw_binds]
|
||||
if workspace_root is None:
|
||||
return roots
|
||||
return [
|
||||
root
|
||||
for root in roots
|
||||
if not is_path_within(workspace_root, root)
|
||||
]
|
||||
|
||||
@ -126,6 +126,7 @@ class TurnDelivery:
|
||||
lifecycle_message: InboundMessage = field(init=False)
|
||||
_stream_base_id: str | None = field(init=False, default=None)
|
||||
_stream_segment: int = field(init=False, default=0)
|
||||
_stream_open: bool = field(init=False, default=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.delivery_message = dataclasses.replace(
|
||||
@ -284,8 +285,14 @@ class TurnDelivery:
|
||||
metadata=self.delivery_message.metadata,
|
||||
)
|
||||
)
|
||||
self._stream_open = True
|
||||
|
||||
async def _publish_stream_end(self, *, resuming: bool = False) -> None:
|
||||
async def _publish_stream_end(
|
||||
self,
|
||||
*,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
await self.bus.publish_outbound(
|
||||
outbound_message_for_event(
|
||||
channel=self.delivery_message.channel,
|
||||
@ -293,8 +300,16 @@ class TurnDelivery:
|
||||
event=StreamEndEvent(
|
||||
stream_id=self._stream_id(),
|
||||
resuming=resuming,
|
||||
merge_next=merge_next,
|
||||
),
|
||||
metadata=self.delivery_message.metadata,
|
||||
)
|
||||
)
|
||||
self._stream_segment += 1
|
||||
self._stream_open = merge_next
|
||||
if not merge_next:
|
||||
self._stream_segment += 1
|
||||
|
||||
async def abort_stream(self) -> None:
|
||||
"""Close an interrupted stream so stateful channels can release its buffer."""
|
||||
if self._stream_open:
|
||||
await self._publish_stream_end()
|
||||
|
||||
@ -46,6 +46,7 @@ class StreamEndEvent(OutboundEvent):
|
||||
content: str = ""
|
||||
stream_id: str | None = None
|
||||
resuming: bool = False
|
||||
merge_next: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@ -176,6 +177,7 @@ def _legacy_event_from_metadata(msg: OutboundMessage) -> OutboundEvent | None:
|
||||
content=msg.content,
|
||||
stream_id=_metadata_str(meta, "_stream_id"),
|
||||
resuming=bool(meta.get("_resuming")),
|
||||
merge_next=bool(meta.get("_merge_next")),
|
||||
)
|
||||
if meta.get("_stream_delta"):
|
||||
return StreamDeltaEvent(
|
||||
|
||||
@ -29,7 +29,7 @@ class BaseChannel(ABC):
|
||||
name: str = "base"
|
||||
display_name: str = "Base"
|
||||
send_progress: bool = True
|
||||
send_tool_hints: bool = False
|
||||
send_tool_hints: bool = True
|
||||
show_reasoning: bool = True
|
||||
|
||||
def __init__(self, config: Any, bus: MessageBus):
|
||||
@ -110,6 +110,7 @@ class BaseChannel(ABC):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Deliver a streaming text chunk.
|
||||
|
||||
@ -118,6 +119,9 @@ class BaseChannel(ABC):
|
||||
|
||||
Stateful implementations should key buffers by ``stream_id`` rather
|
||||
than only by ``chat_id`` when it is provided.
|
||||
|
||||
``merge_next`` marks a resumable provider boundary whose next text
|
||||
segment belongs to the same user-visible message.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@ -24,6 +24,17 @@ from nanobot.security.network import validate_resolved_url, validate_url_target
|
||||
|
||||
DINGTALK_MAX_REMOTE_MEDIA_BYTES = 20 * 1024 * 1024
|
||||
DINGTALK_MAX_REMOTE_MEDIA_REDIRECTS = 3
|
||||
_DINGTALK_MARKDOWN_INLINE_SPECIALS = frozenset(r"\`*_{}[]()<>#+-.!|~")
|
||||
_DINGTALK_SENDER_NAME_MAX_CHARS = 80
|
||||
|
||||
|
||||
def _escape_markdown_sender_name(value: str) -> str:
|
||||
"""Render an untrusted display name as one bounded Markdown-safe line."""
|
||||
normalized = " ".join(value.split())[:_DINGTALK_SENDER_NAME_MAX_CHARS]
|
||||
return "".join(
|
||||
f"\\{char}" if char in _DINGTALK_MARKDOWN_INLINE_SPECIALS else char
|
||||
for char in normalized
|
||||
)
|
||||
|
||||
try:
|
||||
from dingtalk_stream import (
|
||||
@ -175,6 +186,7 @@ class DingTalkConfig(Base):
|
||||
allow_remote_media_redirects: bool = False
|
||||
remote_media_redirect_allowed_hosts: list[str] = Field(default_factory=list)
|
||||
group_user_isolation: bool = False # If True, each user in group chat gets their own session
|
||||
disable_private_chat: bool = False # If True, reject 1:1 DMs with a notice; group chats only
|
||||
|
||||
|
||||
class DingTalkChannel(BaseChannel):
|
||||
@ -712,8 +724,20 @@ class DingTalkChannel(BaseChannel):
|
||||
if not token:
|
||||
raise RuntimeError("DingTalk access token unavailable")
|
||||
|
||||
if msg.content and msg.content.strip():
|
||||
if not await self._send_markdown_text(token, msg.chat_id, msg.content.strip()):
|
||||
content = msg.content.strip() if msg.content else ""
|
||||
if content:
|
||||
# In group chats, prefix the reply with a markdown header naming the
|
||||
# sender so the addressed user can spot the reply. Visual only —
|
||||
# DingTalk's markdown robot messages do not push real @ notifications.
|
||||
sender_name = msg.metadata.get("sender_name") if msg.metadata else None
|
||||
safe_sender_name = (
|
||||
_escape_markdown_sender_name(sender_name)
|
||||
if isinstance(sender_name, str)
|
||||
else ""
|
||||
)
|
||||
if msg.chat_id.startswith("group:") and safe_sender_name:
|
||||
content = f"# @{safe_sender_name}\n\n{content}"
|
||||
if not await self._send_markdown_text(token, msg.chat_id, content):
|
||||
raise RuntimeError("DingTalk text message was not delivered")
|
||||
|
||||
for media_ref in msg.media or []:
|
||||
@ -733,7 +757,7 @@ class DingTalkChannel(BaseChannel):
|
||||
async def _on_message(
|
||||
self,
|
||||
content: str,
|
||||
sender_id: str,
|
||||
sender_id: str | None,
|
||||
sender_name: str,
|
||||
conversation_type: str | None = None,
|
||||
conversation_id: str | None = None,
|
||||
@ -745,11 +769,30 @@ class DingTalkChannel(BaseChannel):
|
||||
"""
|
||||
try:
|
||||
self.logger.info("inbound: {} from {}", content, sender_name)
|
||||
if not sender_id:
|
||||
self.logger.warning("dropping DingTalk message without a sender ID")
|
||||
return
|
||||
is_group = conversation_type == "2" and conversation_id
|
||||
chat_id = f"group:{conversation_id}" if is_group else sender_id
|
||||
session_key = None
|
||||
if is_group and self.config.group_user_isolation:
|
||||
session_key = f"{self.name}:group:{conversation_id}:{sender_id}"
|
||||
|
||||
if not is_group and self.config.disable_private_chat:
|
||||
# Group-only kill switch: drop DMs with a notice *before* any
|
||||
# allow_from / pairing check, so even allowlisted senders are
|
||||
# redirected — intentional, this is a hard private-chat guard
|
||||
# rather than an authorization decision. No session is created.
|
||||
self.logger.info("private chat disabled; rejecting DM from {}", sender_name)
|
||||
await self.send(
|
||||
OutboundMessage(
|
||||
channel=self.name,
|
||||
chat_id=chat_id,
|
||||
content="该机器人未开启私聊,请在群聊中与我对话。",
|
||||
)
|
||||
)
|
||||
return
|
||||
|
||||
await self._handle_message(
|
||||
sender_id=sender_id,
|
||||
chat_id=chat_id,
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import json
|
||||
import zipfile
|
||||
from io import BytesIO
|
||||
from types import SimpleNamespace
|
||||
@ -9,15 +10,15 @@ import pytest
|
||||
|
||||
# Check optional dingtalk dependencies before running tests
|
||||
try:
|
||||
from nanobot.channels import dingtalk
|
||||
DINGTALK_AVAILABLE = getattr(dingtalk, "DINGTALK_AVAILABLE", False)
|
||||
import nanobot.channels.dingtalk.runtime as dingtalk_module
|
||||
|
||||
DINGTALK_AVAILABLE = dingtalk_module.DINGTALK_AVAILABLE
|
||||
except ImportError:
|
||||
DINGTALK_AVAILABLE = False
|
||||
|
||||
if not DINGTALK_AVAILABLE:
|
||||
pytest.skip("DingTalk dependencies not installed (dingtalk-stream)", allow_module_level=True)
|
||||
|
||||
import nanobot.channels.dingtalk.runtime as dingtalk_module
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.dingtalk.runtime import (
|
||||
@ -153,6 +154,92 @@ async def test_group_user_isolation_true_separates_sessions() -> None:
|
||||
assert msg1.chat_id == msg2.chat_id == "group:conv123"
|
||||
|
||||
|
||||
def test_disable_private_chat_uses_camel_case_config_key() -> None:
|
||||
config = DingTalkConfig.model_validate({"disablePrivateChat": True})
|
||||
|
||||
assert config.disable_private_chat is True
|
||||
assert config.model_dump(mode="json", by_alias=True)["disablePrivateChat"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dm_rejected_when_private_chat_disabled(monkeypatch) -> None:
|
||||
"""With disable_private_chat=True, a 1:1 DM is rejected: nothing reaches the
|
||||
bus (no session is created) and the bot replies with a notice directing the
|
||||
user to group chat. Even allowlisted senders are blocked in DMs."""
|
||||
config = DingTalkConfig(
|
||||
client_id="app",
|
||||
client_secret="secret",
|
||||
allow_from=["*"], # even allowlisted senders are blocked in DMs
|
||||
disable_private_chat=True,
|
||||
)
|
||||
bus = MessageBus()
|
||||
channel = DingTalkChannel(config, bus)
|
||||
|
||||
async def fake_get_token():
|
||||
return "test-token"
|
||||
|
||||
monkeypatch.setattr(channel, "_get_access_token", fake_get_token)
|
||||
channel._http = _FakeHttp()
|
||||
|
||||
await channel._on_message(
|
||||
"hello",
|
||||
sender_id="user1",
|
||||
sender_name="Alice",
|
||||
conversation_type="1",
|
||||
)
|
||||
|
||||
# No inbound message was published -> no session created
|
||||
assert bus.inbound.empty()
|
||||
|
||||
# A notice was sent back to the DM user via the private-chat API
|
||||
assert len(channel._http.calls) == 1
|
||||
call = channel._http.calls[0]
|
||||
assert call["url"] == "https://api.dingtalk.com/v1.0/robot/oToMessages/batchSend"
|
||||
assert call["json"]["msgKey"] == "sampleMarkdown"
|
||||
assert call["json"]["userIds"] == ["user1"]
|
||||
assert "该机器人未开启私聊,请在群聊中与我对话。" in call["json"]["msgParam"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dm_allowed_when_private_chat_not_disabled() -> None:
|
||||
"""By default (disable_private_chat=False), a 1:1 DM still reaches the bus."""
|
||||
config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"])
|
||||
bus = MessageBus()
|
||||
channel = DingTalkChannel(config, bus)
|
||||
|
||||
await channel._on_message(
|
||||
"hello",
|
||||
sender_id="user1",
|
||||
sender_name="Alice",
|
||||
conversation_type="1",
|
||||
)
|
||||
|
||||
msg = await bus.consume_inbound()
|
||||
assert msg.chat_id == "user1"
|
||||
assert msg.metadata["conversation_type"] == "1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_message_allowed_when_private_chat_disabled() -> None:
|
||||
"""Disabling private chat must not affect group messages."""
|
||||
config = DingTalkConfig(
|
||||
client_id="app", client_secret="secret", allow_from=["*"], disable_private_chat=True
|
||||
)
|
||||
bus = MessageBus()
|
||||
channel = DingTalkChannel(config, bus)
|
||||
|
||||
await channel._on_message(
|
||||
"hello",
|
||||
sender_id="user1",
|
||||
sender_name="Alice",
|
||||
conversation_type="2",
|
||||
conversation_id="conv123",
|
||||
)
|
||||
|
||||
msg = await bus.consume_inbound()
|
||||
assert msg.chat_id == "group:conv123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_send_uses_group_messages_api() -> None:
|
||||
config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"])
|
||||
@ -173,6 +260,105 @@ async def test_group_send_uses_group_messages_api() -> None:
|
||||
assert call["json"]["msgKey"] == "sampleMarkdown"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_send_prepends_sender_mention(monkeypatch) -> None:
|
||||
"""Group replies are prefixed with a markdown header naming the sender."""
|
||||
config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"])
|
||||
channel = DingTalkChannel(config, MessageBus())
|
||||
channel._http = _FakeHttp()
|
||||
|
||||
async def _fake_token() -> str:
|
||||
return "token"
|
||||
|
||||
monkeypatch.setattr(channel, "_get_access_token", _fake_token)
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(
|
||||
channel="dingtalk",
|
||||
chat_id="group:conv123",
|
||||
content="hello",
|
||||
metadata={"sender_name": "Alice"},
|
||||
)
|
||||
)
|
||||
|
||||
sent_text = json.loads(channel._http.calls[0]["json"]["msgParam"])["text"]
|
||||
assert sent_text == "# @Alice\n\nhello"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_group_send_escapes_untrusted_sender_name(monkeypatch) -> None:
|
||||
"""A sender nickname cannot inject extra Markdown blocks into the reply."""
|
||||
config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"])
|
||||
channel = DingTalkChannel(config, MessageBus())
|
||||
channel._http = _FakeHttp()
|
||||
|
||||
async def _fake_token() -> str:
|
||||
return "token"
|
||||
|
||||
monkeypatch.setattr(channel, "_get_access_token", _fake_token)
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(
|
||||
channel="dingtalk",
|
||||
chat_id="group:conv123",
|
||||
content="hello",
|
||||
metadata={"sender_name": "Alice\n# [click](https://evil) *admin*"},
|
||||
)
|
||||
)
|
||||
|
||||
sent_text = json.loads(channel._http.calls[0]["json"]["msgParam"])["text"]
|
||||
assert sent_text == r"# @Alice \# \[click\]\(https://evil\) \*admin\*" + "\n\nhello"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_private_send_does_not_prepend_mention(monkeypatch) -> None:
|
||||
"""Private replies are sent verbatim, without the sender header."""
|
||||
config = DingTalkConfig(client_id="app", client_secret="secret", allow_from=["*"])
|
||||
channel = DingTalkChannel(config, MessageBus())
|
||||
channel._http = _FakeHttp()
|
||||
|
||||
async def _fake_token() -> str:
|
||||
return "token"
|
||||
|
||||
monkeypatch.setattr(channel, "_get_access_token", _fake_token)
|
||||
|
||||
await channel.send(
|
||||
OutboundMessage(
|
||||
channel="dingtalk",
|
||||
chat_id="user1", # private chat: no "group:" prefix
|
||||
content="hello",
|
||||
metadata={"sender_name": "Alice"},
|
||||
)
|
||||
)
|
||||
|
||||
sent_text = json.loads(channel._http.calls[0]["json"]["msgParam"])["text"]
|
||||
assert sent_text == "hello"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_message_without_sender_id_is_dropped() -> None:
|
||||
"""Malformed inbound events must not publish or attempt an invalid reply."""
|
||||
config = DingTalkConfig(
|
||||
client_id="app",
|
||||
client_secret="secret",
|
||||
allow_from=["*"],
|
||||
disable_private_chat=True,
|
||||
)
|
||||
bus = MessageBus()
|
||||
channel = DingTalkChannel(config, bus)
|
||||
channel._http = _FakeHttp()
|
||||
|
||||
await channel._on_message(
|
||||
"hello",
|
||||
sender_id=None,
|
||||
sender_name="Unknown",
|
||||
conversation_type="1",
|
||||
)
|
||||
|
||||
assert bus.inbound.empty()
|
||||
assert channel._http.calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handler_uses_voice_recognition_text_when_text_is_empty(monkeypatch) -> None:
|
||||
bus = MessageBus()
|
||||
|
||||
@ -489,6 +489,7 @@ class DiscordChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Progressive Discord delivery: send once, then edit until the stream ends."""
|
||||
client = self._client
|
||||
@ -496,6 +497,10 @@ class DiscordChannel(BaseChannel):
|
||||
self.logger.warning("client not ready; dropping stream delta")
|
||||
return
|
||||
|
||||
if stream_end and merge_next:
|
||||
if not delta:
|
||||
return
|
||||
stream_end = False
|
||||
if stream_end:
|
||||
buf = self._stream_bufs.get(chat_id)
|
||||
if not buf or buf.message is None or not buf.text:
|
||||
|
||||
@ -754,6 +754,36 @@ async def test_send_delta_streams_by_editing_message(monkeypatch) -> None:
|
||||
assert owner._stream_bufs == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_merge_next_keeps_one_message(monkeypatch) -> None:
|
||||
owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus())
|
||||
client = _FakeDiscordClient(owner, intents=None)
|
||||
owner._client = client
|
||||
owner._running = True
|
||||
target = _FakeChannel(channel_id=123)
|
||||
client.channels[123] = target
|
||||
|
||||
times = iter([1.0, 3.0, 5.0])
|
||||
monkeypatch.setattr("nanobot.channels.discord.runtime.time.monotonic", lambda: next(times, 5.0))
|
||||
|
||||
await owner.send_delta(
|
||||
"123",
|
||||
"first-",
|
||||
stream_id="s1",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
await owner.send_delta("123", "second", stream_id="s1")
|
||||
await owner.send_delta("123", "", stream_id="s1", stream_end=True)
|
||||
|
||||
assert target.sent_payloads == [{"content": "first-"}]
|
||||
assert target.sent_messages[0].edits == [
|
||||
{"content": "first-second"},
|
||||
{"content": "first-second"},
|
||||
]
|
||||
assert owner._stream_bufs == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_stream_end_splits_oversized_reply(monkeypatch) -> None:
|
||||
owner = DiscordChannel(DiscordConfig(enabled=True, allow_from=["*"]), MessageBus())
|
||||
|
||||
@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import secrets
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
@ -41,6 +42,7 @@ class FeishuConnectStore:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._sessions: dict[str, FeishuConnectSession] = {}
|
||||
self._completion_lock = threading.Lock()
|
||||
|
||||
async def handle(self, action: str, query: QueryParams) -> dict[str, Any]:
|
||||
"""Handle one generic settings connection action."""
|
||||
@ -58,7 +60,7 @@ class FeishuConnectStore:
|
||||
if action == "poll":
|
||||
return await asyncio.to_thread(self.poll, session_id)
|
||||
if action == "cancel":
|
||||
return self.cancel(session_id)
|
||||
return await asyncio.to_thread(self.cancel, session_id)
|
||||
raise ChannelConnectError(f"unsupported Feishu connect action: {action}", status=404)
|
||||
|
||||
def start(
|
||||
@ -127,24 +129,33 @@ class FeishuConnectStore:
|
||||
session.last_error = str(exc)
|
||||
return _pending_payload(session)
|
||||
|
||||
session.domain = str(result.get("domain") or session.domain)
|
||||
status = result.get("status")
|
||||
if status == "succeeded":
|
||||
session.instance_id = feishu.save_registration_result(
|
||||
result,
|
||||
instance_id=session.instance_id,
|
||||
name=session.instance_name,
|
||||
)
|
||||
self._sessions.pop(session_id, None)
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"instance_id": session.instance_id,
|
||||
"status": "succeeded",
|
||||
"message": "Feishu is connected.",
|
||||
"domain": session.domain,
|
||||
"app_id": result.get("app_id"),
|
||||
}
|
||||
with self._completion_lock:
|
||||
if self._sessions.get(session_id) is not session:
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"instance_id": session.instance_id,
|
||||
"status": "cancelled",
|
||||
"message": "Feishu connection cancelled.",
|
||||
}
|
||||
session.domain = str(result.get("domain") or session.domain)
|
||||
session.instance_id = feishu.save_registration_result(
|
||||
result,
|
||||
instance_id=session.instance_id,
|
||||
name=session.instance_name,
|
||||
)
|
||||
self._sessions.pop(session_id, None)
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"instance_id": session.instance_id,
|
||||
"status": "succeeded",
|
||||
"message": "Feishu is connected.",
|
||||
"domain": session.domain,
|
||||
"app_id": result.get("app_id"),
|
||||
}
|
||||
|
||||
session.domain = str(result.get("domain") or session.domain)
|
||||
if status == "failed":
|
||||
self._sessions.pop(session_id, None)
|
||||
return {
|
||||
@ -158,7 +169,8 @@ class FeishuConnectStore:
|
||||
return _pending_payload(session)
|
||||
|
||||
def cancel(self, session_id: str) -> dict[str, Any]:
|
||||
session = self._sessions.pop(session_id, None)
|
||||
with self._completion_lock:
|
||||
session = self._sessions.pop(session_id, None)
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"instance_id": session.instance_id if session else DEFAULT_INSTANCE_ID,
|
||||
|
||||
@ -269,7 +269,7 @@ def _extract_element_content(element: dict) -> list[str]:
|
||||
parts.append(text_content)
|
||||
elif isinstance(text, str):
|
||||
parts.append(text)
|
||||
for field in element.get("fields", []):
|
||||
for field in element.get("fields") or []:
|
||||
if isinstance(field, dict):
|
||||
field_text = field.get("text", {})
|
||||
if isinstance(field_text, dict):
|
||||
@ -291,7 +291,10 @@ def _extract_element_content(element: dict) -> list[str]:
|
||||
c = text.get("content", "")
|
||||
if c:
|
||||
parts.append(c)
|
||||
url = element.get("url", "") or element.get("multi_url", {}).get("url", "")
|
||||
multi_url = element.get("multi_url") or {}
|
||||
url = element.get("url", "") or (
|
||||
multi_url.get("url", "") if isinstance(multi_url, dict) else ""
|
||||
)
|
||||
if url:
|
||||
parts.append(f"link: {url}")
|
||||
|
||||
@ -300,12 +303,14 @@ def _extract_element_content(element: dict) -> list[str]:
|
||||
parts.append(alt.get("content", "[image]") if isinstance(alt, dict) else "[image]")
|
||||
|
||||
elif tag == "note":
|
||||
for ne in element.get("elements", []):
|
||||
for ne in element.get("elements") or []:
|
||||
parts.extend(_extract_element_content(ne))
|
||||
|
||||
elif tag == "column_set":
|
||||
for col in element.get("columns", []):
|
||||
for ce in col.get("elements", []):
|
||||
for col in element.get("columns") or []:
|
||||
if not isinstance(col, dict):
|
||||
continue
|
||||
for ce in col.get("elements") or []:
|
||||
parts.extend(_extract_element_content(ce))
|
||||
|
||||
elif tag == "plain_text":
|
||||
@ -319,7 +324,7 @@ def _extract_element_content(element: dict) -> list[str]:
|
||||
for column in (element.get("columns") or [])
|
||||
if isinstance(column, dict) and column.get("name")
|
||||
]
|
||||
rows = element.get("rows", [])
|
||||
rows = element.get("rows") or []
|
||||
if columns:
|
||||
parts.append(" | ".join(header for _, header in columns))
|
||||
if isinstance(rows, list):
|
||||
@ -337,7 +342,7 @@ def _extract_element_content(element: dict) -> list[str]:
|
||||
parts.append(row_text)
|
||||
|
||||
else:
|
||||
for ne in element.get("elements", []):
|
||||
for ne in element.get("elements") or []:
|
||||
parts.extend(_extract_element_content(ne))
|
||||
|
||||
return parts
|
||||
@ -356,7 +361,8 @@ def _extract_post_content(content_json: dict) -> tuple[str, list[str]]:
|
||||
if not isinstance(block, dict) or not isinstance(block.get("content"), list):
|
||||
return None, []
|
||||
texts, images = [], []
|
||||
if title := block.get("title"):
|
||||
title = block.get("title")
|
||||
if isinstance(title, str) and title:
|
||||
texts.append(title)
|
||||
for row in block["content"]:
|
||||
if not isinstance(row, list):
|
||||
@ -366,12 +372,19 @@ def _extract_post_content(content_json: dict) -> tuple[str, list[str]]:
|
||||
continue
|
||||
tag = el.get("tag")
|
||||
if tag in ("text", "a"):
|
||||
texts.append(el.get("text", ""))
|
||||
text = el.get("text", "")
|
||||
if isinstance(text, str):
|
||||
texts.append(text)
|
||||
elif tag == "at":
|
||||
texts.append(f"@{el.get('user_name', 'user')}")
|
||||
user = el.get("user_name", "user")
|
||||
texts.append(f"@{user if isinstance(user, str) and user else 'user'}")
|
||||
elif tag == "code_block":
|
||||
lang = el.get("language", "")
|
||||
code_text = el.get("text", "")
|
||||
if not isinstance(lang, str):
|
||||
lang = ""
|
||||
if not isinstance(code_text, str):
|
||||
code_text = ""
|
||||
texts.append(f"\n```{lang}\n{code_text}\n```\n")
|
||||
elif tag == "img" and (key := el.get("image_key")):
|
||||
images.append(key)
|
||||
@ -2203,6 +2216,7 @@ class FeishuChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Progressive streaming via CardKit: create card on first delta, stream-update on subsequent.
|
||||
|
||||
@ -2218,6 +2232,10 @@ class FeishuChannel(BaseChannel):
|
||||
rid_type = "chat_id" if chat_id.startswith("oc_") else "open_id"
|
||||
|
||||
# --- stream end: final update or fallback ---
|
||||
if stream_end and merge_next:
|
||||
if not delta:
|
||||
return
|
||||
stream_end = False
|
||||
if stream_end:
|
||||
message_id = meta.get("message_id")
|
||||
# Only finalize the OnIt -> DONE reaction transition on the truly
|
||||
|
||||
122
nanobot/channels/feishu/tests/test_connect.py
Normal file
122
nanobot/channels/feishu/tests/test_connect.py
Normal file
@ -0,0 +1,122 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.channels.feishu import runtime as feishu
|
||||
from nanobot.channels.feishu.connect import FeishuConnectStore
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_feishu_cancel_wins_over_inflight_confirmation(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
poll_started = threading.Event()
|
||||
release_poll = threading.Event()
|
||||
saved_results: list[dict[str, Any]] = []
|
||||
|
||||
monkeypatch.setattr(feishu, "_init_registration", lambda _domain: None)
|
||||
monkeypatch.setattr(
|
||||
feishu,
|
||||
"_begin_registration",
|
||||
lambda _domain: {
|
||||
"device_code": "device-cancel",
|
||||
"qr_url": "https://qr.example/cancel",
|
||||
"expire_in": 600,
|
||||
"interval": 2,
|
||||
},
|
||||
)
|
||||
|
||||
def fake_poll_registration_once(**_kwargs: Any) -> dict[str, str]:
|
||||
poll_started.set()
|
||||
assert release_poll.wait(timeout=5)
|
||||
return {
|
||||
"status": "succeeded",
|
||||
"domain": "feishu",
|
||||
"app_id": "late-app",
|
||||
"app_secret": "late-secret",
|
||||
}
|
||||
|
||||
def fake_save_registration_result(
|
||||
result: dict[str, Any],
|
||||
**_kwargs: Any,
|
||||
) -> str:
|
||||
saved_results.append(result)
|
||||
return "default"
|
||||
|
||||
monkeypatch.setattr(feishu, "poll_registration_once", fake_poll_registration_once)
|
||||
monkeypatch.setattr(feishu, "save_registration_result", fake_save_registration_result)
|
||||
|
||||
store = FeishuConnectStore()
|
||||
started = await store.handle("start", {})
|
||||
query = {"session_id": [started["session_id"]]}
|
||||
poll_task = asyncio.create_task(store.handle("poll", query))
|
||||
assert await asyncio.to_thread(poll_started.wait, 5)
|
||||
|
||||
cancelled = await store.handle("cancel", query)
|
||||
release_poll.set()
|
||||
completed = await poll_task
|
||||
|
||||
assert cancelled["status"] == "cancelled"
|
||||
assert completed["status"] == "cancelled"
|
||||
assert saved_results == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_feishu_cancel_does_not_interleave_with_registration_save(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
save_started = threading.Event()
|
||||
release_save = threading.Event()
|
||||
|
||||
monkeypatch.setattr(feishu, "_init_registration", lambda _domain: None)
|
||||
monkeypatch.setattr(
|
||||
feishu,
|
||||
"_begin_registration",
|
||||
lambda _domain: {
|
||||
"device_code": "device-lock",
|
||||
"qr_url": "https://qr.example/lock",
|
||||
"expire_in": 600,
|
||||
"interval": 2,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
feishu,
|
||||
"poll_registration_once",
|
||||
lambda **_kwargs: {
|
||||
"status": "succeeded",
|
||||
"domain": "feishu",
|
||||
"app_id": "saved-app",
|
||||
"app_secret": "saved-secret",
|
||||
},
|
||||
)
|
||||
|
||||
def fake_save_registration_result(
|
||||
_result: dict[str, Any],
|
||||
**_kwargs: Any,
|
||||
) -> str:
|
||||
save_started.set()
|
||||
assert release_save.wait(timeout=5)
|
||||
return "default"
|
||||
|
||||
monkeypatch.setattr(feishu, "save_registration_result", fake_save_registration_result)
|
||||
|
||||
store = FeishuConnectStore()
|
||||
started = await store.handle("start", {})
|
||||
query = {"session_id": [started["session_id"]]}
|
||||
poll_task = asyncio.create_task(store.handle("poll", query))
|
||||
assert await asyncio.to_thread(save_started.wait, 5)
|
||||
|
||||
cancel_task = asyncio.create_task(store.handle("cancel", query))
|
||||
await asyncio.sleep(0)
|
||||
assert not cancel_task.done()
|
||||
|
||||
release_save.set()
|
||||
completed = await poll_task
|
||||
cancelled = await cancel_task
|
||||
|
||||
assert completed["status"] == "succeeded"
|
||||
assert cancelled["status"] == "cancelled"
|
||||
@ -1,6 +1,10 @@
|
||||
import json
|
||||
|
||||
from nanobot.channels.feishu.runtime import _extract_share_card_content
|
||||
from nanobot.channels.feishu.runtime import (
|
||||
_extract_element_content,
|
||||
_extract_post_content,
|
||||
_extract_share_card_content,
|
||||
)
|
||||
|
||||
|
||||
def test_extract_interactive_card_reads_user_dsl_body_elements() -> None:
|
||||
@ -37,3 +41,48 @@ def test_extract_interactive_card_reads_table_rows() -> None:
|
||||
}
|
||||
|
||||
assert _extract_share_card_content(content, "interactive") == "Name | Score\nAlice | 98"
|
||||
|
||||
|
||||
def test_extract_post_content_tolerates_null_fields() -> None:
|
||||
text, images = _extract_post_content(
|
||||
{
|
||||
"title": None,
|
||||
"content": [
|
||||
[
|
||||
{"tag": "text", "text": None},
|
||||
{"tag": "a", "text": None},
|
||||
{"tag": "at", "user_name": None},
|
||||
{"tag": "text", "text": "ok"},
|
||||
{"tag": "code_block", "language": None, "text": None},
|
||||
]
|
||||
],
|
||||
}
|
||||
)
|
||||
assert "@user" in text
|
||||
assert "ok" in text
|
||||
assert images == []
|
||||
|
||||
|
||||
def test_extract_button_tolerates_null_multi_url() -> None:
|
||||
element = {"tag": "button", "text": {"content": "Go"}, "multi_url": None}
|
||||
assert _extract_element_content(element) == ["Go"]
|
||||
|
||||
|
||||
def test_extract_column_set_tolerates_null_columns_and_elements() -> None:
|
||||
assert _extract_element_content({"tag": "column_set", "columns": None}) == []
|
||||
assert _extract_element_content(
|
||||
{"tag": "column_set", "columns": [{"elements": None}]}
|
||||
) == []
|
||||
|
||||
|
||||
def test_extract_div_tolerates_null_fields() -> None:
|
||||
assert _extract_element_content(
|
||||
{"tag": "div", "text": {"content": "hi"}, "fields": None}
|
||||
) == ["hi"]
|
||||
|
||||
|
||||
def test_interactive_card_button_null_multi_url() -> None:
|
||||
content = {
|
||||
"elements": [{"tag": "button", "text": {"content": "Go"}, "multi_url": None}]
|
||||
}
|
||||
assert _extract_share_card_content(content, "interactive") == "Go"
|
||||
|
||||
@ -285,6 +285,27 @@ class TestSendDelta:
|
||||
settings_call = ch._client.cardkit.v1.card.settings.call_args[0][0]
|
||||
assert settings_call.body.sequence == 5 # after final content seq 4
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_merge_next_preserves_buffer(self):
|
||||
ch = _make_channel()
|
||||
ch._stream_bufs["oc_chat1"] = _FeishuStreamBuf(
|
||||
text="first-",
|
||||
card_id="card_1",
|
||||
sequence=3,
|
||||
last_edit=time.monotonic(),
|
||||
)
|
||||
|
||||
await ch.send_delta(
|
||||
"oc_chat1",
|
||||
"boundary",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
|
||||
assert ch._stream_bufs["oc_chat1"].text == "first-boundary"
|
||||
ch._client.cardkit.v1.card_element.content.assert_not_called()
|
||||
ch._client.cardkit.v1.card.settings.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_fallback_when_no_card_id(self):
|
||||
"""If card creation failed, stream_end falls back to a plain card message."""
|
||||
|
||||
@ -4,6 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import inspect
|
||||
from collections.abc import Callable, Iterable
|
||||
from contextlib import suppress
|
||||
from pathlib import Path
|
||||
@ -769,13 +770,29 @@ class ChannelManager:
|
||||
msg: OutboundMessage,
|
||||
event: StreamDeltaEvent | StreamEndEvent,
|
||||
) -> None:
|
||||
kwargs: dict[str, Any] = {
|
||||
"stream_id": event.stream_id,
|
||||
"stream_end": isinstance(event, StreamEndEvent),
|
||||
"resuming": event.resuming if isinstance(event, StreamEndEvent) else False,
|
||||
}
|
||||
if isinstance(event, StreamEndEvent) and event.merge_next:
|
||||
try:
|
||||
signature = inspect.signature(channel.send_delta)
|
||||
if (
|
||||
"merge_next" in signature.parameters
|
||||
or any(
|
||||
parameter.kind is inspect.Parameter.VAR_KEYWORD
|
||||
for parameter in signature.parameters.values()
|
||||
)
|
||||
):
|
||||
kwargs["merge_next"] = True
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
await channel.send_delta(
|
||||
msg.chat_id,
|
||||
msg.content,
|
||||
msg.metadata,
|
||||
stream_id=event.stream_id,
|
||||
stream_end=isinstance(event, StreamEndEvent),
|
||||
resuming=event.resuming if isinstance(event, StreamEndEvent) else False,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@ -856,6 +873,7 @@ class ChannelManager:
|
||||
final_event = StreamEndEvent(
|
||||
stream_id=next_stream_id,
|
||||
resuming=next_event.resuming,
|
||||
merge_next=next_event.merge_next,
|
||||
)
|
||||
# Stream ended - stop coalescing this stream
|
||||
break
|
||||
|
||||
@ -598,9 +598,14 @@ class MatrixChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
relates_to = self._build_thread_relates_to(metadata)
|
||||
|
||||
if stream_end and merge_next:
|
||||
if not delta:
|
||||
return
|
||||
stream_end = False
|
||||
if stream_end:
|
||||
stream_key = _matrix_stream_key(chat_id, stream_id)
|
||||
buf = self._stream_bufs.pop(stream_key, None)
|
||||
|
||||
@ -1937,6 +1937,29 @@ async def test_send_delta_stream_end_replaces_existing_message() -> None:
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_merge_next_preserves_buffer() -> None:
|
||||
channel = MatrixChannel(_make_config(), MessageBus())
|
||||
client = _FakeAsyncClient("", "", "", None)
|
||||
channel.client = client
|
||||
channel._stream_bufs["!room:matrix.org"] = matrix_module._StreamBuf(
|
||||
text="first-",
|
||||
event_id="event-1",
|
||||
last_edit=100.0,
|
||||
)
|
||||
channel.monotonic_time = lambda: 100.1
|
||||
|
||||
await channel.send_delta(
|
||||
"!room:matrix.org",
|
||||
"boundary",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
|
||||
assert channel._stream_bufs["!room:matrix.org"].text == "first-boundary"
|
||||
assert client.room_send_calls == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_keeps_same_room_stream_ids_independent(monkeypatch) -> None:
|
||||
channel = MatrixChannel(_make_config(), MessageBus())
|
||||
|
||||
@ -56,7 +56,7 @@ class MattermostConfig(Base):
|
||||
react_emoji: str = "eyes"
|
||||
done_emoji: str = "white_check_mark"
|
||||
send_progress: bool = True
|
||||
send_tool_hints: bool = False
|
||||
send_tool_hints: bool = True
|
||||
dm: MattermostDMConfig = Field(default_factory=MattermostDMConfig)
|
||||
|
||||
|
||||
@ -515,6 +515,7 @@ class MattermostChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
if not self._http_client:
|
||||
return
|
||||
@ -532,7 +533,11 @@ class MattermostChannel(BaseChannel):
|
||||
final += delta
|
||||
|
||||
if resuming:
|
||||
self._clear_stream_state(stream_id)
|
||||
if merge_next:
|
||||
self._stream_buffers[stream_id] = final
|
||||
self._stream_committed[stream_id] = final
|
||||
else:
|
||||
self._clear_stream_state(stream_id)
|
||||
return
|
||||
|
||||
if final and not meta.get("_progress"):
|
||||
|
||||
@ -119,6 +119,7 @@ def test_config_defaults():
|
||||
assert config.token == ""
|
||||
assert config.streaming is True
|
||||
assert config.streaming_max_chars == 16000
|
||||
assert config.send_tool_hints is True
|
||||
assert config.dm.enabled is True
|
||||
assert config.dm.policy == "open"
|
||||
assert config.reply_in_thread is True
|
||||
@ -131,6 +132,7 @@ def test_config_camelcase_aliases():
|
||||
"allowFromMatchMode": "username",
|
||||
"streamingMaxChars": 8000,
|
||||
"replyInThread": False,
|
||||
"sendToolHints": False,
|
||||
}
|
||||
config = MattermostConfig.model_validate(raw)
|
||||
assert config.server_url == "https://mm.example.com"
|
||||
@ -138,11 +140,13 @@ def test_config_camelcase_aliases():
|
||||
assert config.allow_from_match_mode == "username"
|
||||
assert config.streaming_max_chars == 8000
|
||||
assert config.reply_in_thread is False
|
||||
assert config.send_tool_hints is False
|
||||
|
||||
|
||||
def test_config_default_config_classmethod():
|
||||
d = MattermostChannel.default_config()
|
||||
assert d["enabled"] is False
|
||||
assert d["sendToolHints"] is True
|
||||
assert d["serverUrl"] == ""
|
||||
assert d["token"] == ""
|
||||
|
||||
@ -578,6 +582,33 @@ async def test_stream_end_keyword_resuming_does_not_post_or_mark_done():
|
||||
assert "s1" not in channel._stream_buffers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_merge_next_preserves_buffer_until_final_end():
|
||||
channel, fake = _make_channel()
|
||||
channel._self_id = "bot_id"
|
||||
fake.set_post_response("/api/v4/posts", {"id": "stream_post_1"})
|
||||
await channel.send_delta("chan_1", "first ", stream_id="s1")
|
||||
|
||||
await channel.send_delta(
|
||||
"chan_1",
|
||||
"boundary ",
|
||||
stream_id="s1",
|
||||
stream_end=True,
|
||||
resuming=True,
|
||||
merge_next=True,
|
||||
)
|
||||
|
||||
assert channel._stream_buffers["s1"] == "first boundary "
|
||||
|
||||
await channel.send_delta("chan_1", "second", stream_id="s1")
|
||||
await channel.send_delta("chan_1", "", stream_id="s1", stream_end=True)
|
||||
|
||||
posts = [call for call in fake.post_calls if call["path"] == "/api/v4/posts"]
|
||||
assert len(posts) == 1
|
||||
assert posts[0]["json"]["message"] == "first boundary second"
|
||||
assert "s1" not in channel._stream_buffers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_failure_keeps_buffer_for_retry():
|
||||
channel, fake = _make_channel()
|
||||
|
||||
@ -923,6 +923,7 @@ class TelegramChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Progressive message editing: send on first delta, edit on subsequent ones."""
|
||||
if not self._app:
|
||||
@ -930,6 +931,10 @@ class TelegramChannel(BaseChannel):
|
||||
meta = metadata or {}
|
||||
int_chat_id = int(chat_id)
|
||||
|
||||
if stream_end and merge_next:
|
||||
if not delta:
|
||||
return
|
||||
stream_end = False
|
||||
if stream_end:
|
||||
buf = self._stream_bufs.get(chat_id)
|
||||
if not buf or not buf.message_id or not buf.text:
|
||||
|
||||
@ -675,6 +675,33 @@ async def test_send_delta_stream_end_raises_and_keeps_buffer_on_failure() -> Non
|
||||
assert "123" in channel._stream_bufs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_merge_next_preserves_buffer() -> None:
|
||||
channel = TelegramChannel(
|
||||
TelegramConfig(enabled=True, token="123:abc", allow_from=["*"]),
|
||||
MessageBus(),
|
||||
)
|
||||
channel._app = _FakeApp(lambda: None)
|
||||
channel._app.bot.edit_message_text = AsyncMock()
|
||||
channel._stream_bufs["123"] = _StreamBuf(
|
||||
text="first-",
|
||||
message_id=7,
|
||||
last_edit=float("inf"),
|
||||
stream_id="s:0",
|
||||
)
|
||||
|
||||
await channel.send_delta(
|
||||
"123",
|
||||
"boundary",
|
||||
stream_id="s:0",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
|
||||
assert channel._stream_bufs["123"].text == "first-boundary"
|
||||
channel._app.bot.edit_message_text.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_stream_end_treats_not_modified_as_success() -> None:
|
||||
from telegram.error import BadRequest
|
||||
|
||||
@ -995,13 +995,18 @@ class WebSocketChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
conns = list(self._subs.get(chat_id, ()))
|
||||
meta = metadata or {}
|
||||
stream_key = (chat_id, str(stream_id or ""))
|
||||
if stream_end:
|
||||
body: dict[str, Any] = {"event": "stream_end", "chat_id": chat_id}
|
||||
buffered = self._stream_text_buffers.pop(stream_key, [])
|
||||
buffered = (
|
||||
self._stream_text_buffers.setdefault(stream_key, [])
|
||||
if merge_next
|
||||
else self._stream_text_buffers.pop(stream_key, [])
|
||||
)
|
||||
if delta:
|
||||
buffered.append(delta)
|
||||
full_text = "".join(buffered)
|
||||
@ -1019,6 +1024,8 @@ class WebSocketChannel(BaseChannel):
|
||||
body["stream_id"] = stream_id
|
||||
if stream_end and resuming:
|
||||
body["resuming"] = True
|
||||
if stream_end and merge_next:
|
||||
body["merge_next"] = True
|
||||
self._transcripts.prepare_and_append(
|
||||
chat_id,
|
||||
body,
|
||||
|
||||
@ -1350,6 +1350,39 @@ async def test_send_delta_marks_resuming_stream_end() -> None:
|
||||
assert payload["resuming"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_keeps_buffer_across_merged_stream_boundary() -> None:
|
||||
bus = MagicMock()
|
||||
channel = WebSocketChannel(
|
||||
{"enabled": True, "allowFrom": ["*"], "streaming": True},
|
||||
bus,
|
||||
gateway=_basic_handler(bus),
|
||||
)
|
||||
mock_ws = AsyncMock()
|
||||
channel._attach(mock_ws, "chat-1")
|
||||
|
||||
await channel.send_delta("chat-1", "first ", stream_id="sid")
|
||||
await channel.send_delta(
|
||||
"chat-1",
|
||||
"",
|
||||
stream_id="sid",
|
||||
stream_end=True,
|
||||
resuming=True,
|
||||
merge_next=True,
|
||||
)
|
||||
await channel.send_delta("chat-1", "second", stream_id="sid")
|
||||
await channel.send_delta("chat-1", "", stream_id="sid", stream_end=True)
|
||||
|
||||
payloads = [json.loads(call.args[0]) for call in mock_ws.send.await_args_list]
|
||||
assert payloads[1]["merge_next"] is True
|
||||
assert payloads[1]["resuming"] is True
|
||||
assert [payload["text"] for payload in payloads if payload["event"] == "delta"] == [
|
||||
"first ",
|
||||
"second",
|
||||
]
|
||||
assert ("chat-1", "sid") not in channel._stream_text_buffers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_delta_stream_end_includes_inline_final_text() -> None:
|
||||
bus = MagicMock()
|
||||
|
||||
@ -130,6 +130,12 @@ class WeixinConnectStore:
|
||||
|
||||
status = status_data.get("status", "")
|
||||
if status == "confirmed":
|
||||
if self._sessions.get(session_id) is not session:
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"status": "cancelled",
|
||||
"message": "WeChat login cancelled.",
|
||||
}
|
||||
token = str(status_data.get("bot_token", "") or "")
|
||||
if not token:
|
||||
self._sessions.pop(session_id, None)
|
||||
|
||||
@ -1243,6 +1243,7 @@ class WeixinChannel(BaseChannel):
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
"""Deliver a streamed reply to WeChat.
|
||||
|
||||
@ -1256,6 +1257,10 @@ class WeixinChannel(BaseChannel):
|
||||
return
|
||||
is_end = stream_end or bool(meta.get("_stream_end"))
|
||||
buffer_key = stream_id or chat_id
|
||||
if is_end and merge_next:
|
||||
if delta:
|
||||
self._stream_buffers.setdefault(buffer_key, []).append(delta)
|
||||
return
|
||||
# Accumulate intermediate deltas. The stream_end message's own content
|
||||
# (present when the manager coalesces deltas into the end message) is
|
||||
# folded into `full` below instead of appended here, so a send retry
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
@ -97,3 +98,52 @@ async def test_weixin_reconnect_keeps_existing_account_until_scan_succeeds(
|
||||
cancelled = await store.cancel(started["session_id"])
|
||||
assert cancelled["status"] == "cancelled"
|
||||
assert json.loads(state_file.read_text(encoding="utf-8")) == existing
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_weixin_cancel_wins_over_inflight_confirmation(
|
||||
tmp_path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
state_dir = tmp_path / "weixin-state"
|
||||
config_path = tmp_path / "config.json"
|
||||
save_config(
|
||||
Config.model_validate({"channels": {"weixin": {"stateDir": str(state_dir)}}}),
|
||||
config_path,
|
||||
)
|
||||
monkeypatch.setattr("nanobot.config.loader._current_config_path", config_path)
|
||||
|
||||
poll_started = asyncio.Event()
|
||||
release_poll = asyncio.Event()
|
||||
|
||||
async def fake_fetch_qr_code(self: WeixinChannel) -> tuple[str, str]:
|
||||
return "qr-cancel", "https://qr.example/cancel"
|
||||
|
||||
async def fake_api_get_with_base(
|
||||
self: WeixinChannel,
|
||||
**_kwargs: Any,
|
||||
) -> dict[str, str]:
|
||||
poll_started.set()
|
||||
await release_poll.wait()
|
||||
return {
|
||||
"status": "confirmed",
|
||||
"bot_token": "late-token",
|
||||
"ilink_user_id": "late-user",
|
||||
}
|
||||
|
||||
monkeypatch.setattr(WeixinChannel, "_fetch_qr_code", fake_fetch_qr_code)
|
||||
monkeypatch.setattr(WeixinChannel, "_api_get_with_base", fake_api_get_with_base)
|
||||
|
||||
store = WeixinConnectStore()
|
||||
started = await store.handle("start", {})
|
||||
query = {"session_id": [started["session_id"]]}
|
||||
poll_task = asyncio.create_task(store.handle("poll", query))
|
||||
await asyncio.wait_for(poll_started.wait(), timeout=5)
|
||||
|
||||
cancelled = await store.handle("cancel", query)
|
||||
release_poll.set()
|
||||
completed = await poll_task
|
||||
|
||||
assert cancelled["status"] == "cancelled"
|
||||
assert completed["status"] == "cancelled"
|
||||
assert not (state_dir / "account.json").exists()
|
||||
|
||||
@ -1824,6 +1824,29 @@ async def test_stream_end_flushes_buffered_answer() -> None:
|
||||
assert "wx-user" not in channel._stream_buffers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_merge_next_preserves_buffer_until_final_end() -> None:
|
||||
channel, _bus = _make_channel()
|
||||
channel._client = object()
|
||||
channel._token = "token"
|
||||
channel._context_tokens["wx-user"] = "ctx-1"
|
||||
channel._context_token_at["wx-user"] = time.time()
|
||||
channel._send_text = AsyncMock()
|
||||
|
||||
await channel.send_delta(
|
||||
"wx-user",
|
||||
"first-",
|
||||
stream_id="s1",
|
||||
stream_end=True,
|
||||
merge_next=True,
|
||||
)
|
||||
await channel.send_delta("wx-user", "second", stream_id="s1")
|
||||
await channel.send_delta("wx-user", "", stream_id="s1", stream_end=True)
|
||||
|
||||
channel._send_text.assert_awaited_once_with("wx-user", "first-second", "ctx-1")
|
||||
assert "s1" not in channel._stream_buffers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_send_failure_keeps_buffer_for_retry() -> None:
|
||||
channel, _bus = _make_channel()
|
||||
|
||||
@ -78,6 +78,10 @@ from nanobot.cli.stream import StreamRenderer, ThinkingSpinner # noqa: E402
|
||||
from nanobot.config.paths import get_workspace_path, is_default_workspace # noqa: E402
|
||||
from nanobot.config.schema import Config # noqa: E402
|
||||
from nanobot.security.network import is_loopback_host # noqa: E402
|
||||
from nanobot.session.keys import ( # noqa: E402
|
||||
UNIFIED_SESSION_KEY,
|
||||
last_channel_from_metadata,
|
||||
)
|
||||
from nanobot.utils.evaluator import evaluate_response, resolve_evaluator_prompt # noqa: E402
|
||||
from nanobot.utils.helpers import ( # noqa: E402
|
||||
sanitize_surrogates as _sanitize_surrogates,
|
||||
@ -265,6 +269,7 @@ def _pick_heartbeat_target_from_sessions(
|
||||
enabled_channels: Iterable[str],
|
||||
sessions: Iterable[dict[str, Any]],
|
||||
archived_keys: Iterable[str],
|
||||
unified_session_metadata: dict[str, Any] | None = None,
|
||||
) -> tuple[str, str]:
|
||||
enabled = set(enabled_channels)
|
||||
archived = set(archived_keys)
|
||||
@ -272,6 +277,13 @@ def _pick_heartbeat_target_from_sessions(
|
||||
key = item.get("key") or ""
|
||||
if key in archived:
|
||||
continue
|
||||
if key == UNIFIED_SESSION_KEY:
|
||||
route = last_channel_from_metadata(unified_session_metadata)
|
||||
if route is not None:
|
||||
channel, chat_id = route
|
||||
if channel not in {"cli", "system"} and channel in enabled:
|
||||
return channel, chat_id
|
||||
continue
|
||||
if ":" not in key:
|
||||
continue
|
||||
channel, chat_id = key.split(":", 1)
|
||||
@ -1822,12 +1834,13 @@ def _run_gateway(
|
||||
|
||||
# Dream is an internal job — run directly, not through the agent loop.
|
||||
if job.name == "dream":
|
||||
from nanobot.agent.memory import MemoryStore
|
||||
from nanobot.agent.memory import DreamRunProgress, MemoryStore
|
||||
|
||||
dream_session_key = MemoryStore.dream_session_key
|
||||
prune_dream_sessions = MemoryStore.prune_dream_sessions
|
||||
|
||||
store = agent.context.memory
|
||||
progress = DreamRunProgress()
|
||||
resp = None
|
||||
diff_body = ""
|
||||
try:
|
||||
@ -1842,22 +1855,28 @@ def _run_gateway(
|
||||
session_key=key,
|
||||
ephemeral=True,
|
||||
tools=store.build_dream_tools(),
|
||||
on_progress=_silent,
|
||||
on_progress=progress,
|
||||
)
|
||||
# Ground truth: the real file delta, not the LLM's self-report.
|
||||
# The real file delta grounds the audit record; clean completion
|
||||
# decides whether this history batch has finished processing.
|
||||
diff_body = store.dream_content_diff()
|
||||
productive = bool(diff_body) or (
|
||||
not store.git.is_initialized()
|
||||
and MemoryStore.dream_run_completed(resp)
|
||||
completed = MemoryStore.dream_run_completed(
|
||||
resp,
|
||||
had_tool_errors=progress.had_tool_errors,
|
||||
)
|
||||
if productive:
|
||||
if completed:
|
||||
store.set_last_dream_cursor(last_cursor)
|
||||
logger.info("Dream cron job completed, cursor advanced to {}", last_cursor)
|
||||
elif MemoryStore.dream_run_completed(resp):
|
||||
logger.info(
|
||||
"Dream cron job completed with no memory changes; "
|
||||
"cursor not advanced",
|
||||
)
|
||||
if diff_body:
|
||||
logger.info(
|
||||
"Dream cron job completed, cursor advanced to {}",
|
||||
last_cursor,
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Dream cron job completed with no memory changes; "
|
||||
"cursor advanced to {}",
|
||||
last_cursor,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Dream cron job did not complete; cursor remains at {}",
|
||||
@ -1995,10 +2014,16 @@ def _run_gateway(
|
||||
def _pick_heartbeat_target() -> tuple[str, str]:
|
||||
"""Pick a routable channel/chat target for heartbeat-triggered messages."""
|
||||
sidebar_state = read_webui_sidebar_state()
|
||||
unified_metadata = None
|
||||
if config.agents.defaults.unified_session:
|
||||
record = session_manager.read_session_metadata(UNIFIED_SESSION_KEY)
|
||||
if isinstance(record, dict) and isinstance(record.get("metadata"), dict):
|
||||
unified_metadata = record["metadata"]
|
||||
return _pick_heartbeat_target_from_sessions(
|
||||
enabled_channels=channels.enabled_channels,
|
||||
sessions=session_manager.list_sessions(),
|
||||
archived_keys=sidebar_state.get("archived_keys", []),
|
||||
unified_session_metadata=unified_metadata,
|
||||
)
|
||||
|
||||
if channels.enabled_channels:
|
||||
|
||||
@ -3,6 +3,7 @@
|
||||
import asyncio
|
||||
import json
|
||||
import types
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import Any, Literal, NamedTuple, get_args, get_origin
|
||||
@ -14,6 +15,7 @@ except ModuleNotFoundError: # pragma: no cover - exercised in environments with
|
||||
from loguru import logger
|
||||
from pydantic import BaseModel
|
||||
from rich.console import Console
|
||||
from rich.markup import escape
|
||||
from rich.panel import Panel
|
||||
from rich.table import Table
|
||||
|
||||
@ -22,7 +24,7 @@ from nanobot.cli.models import (
|
||||
get_model_context_limit,
|
||||
get_model_suggestions,
|
||||
)
|
||||
from nanobot.config.loader import get_config_path, load_config
|
||||
from nanobot.config.loader import get_config_path, load_config, resolve_config_env_vars
|
||||
from nanobot.config.schema import Config, ModelPresetConfig
|
||||
|
||||
console = Console()
|
||||
@ -44,6 +46,8 @@ class _QuickStartProviderInfo(NamedTuple):
|
||||
default_api_base: str
|
||||
backend: str
|
||||
is_direct: bool
|
||||
is_oauth: bool
|
||||
default_model: str
|
||||
|
||||
|
||||
class _QuickStartEndpointChoice(NamedTuple):
|
||||
@ -73,6 +77,7 @@ _BACK_PRESSED = object() # Sentinel value for back navigation
|
||||
_MODEL_PRESET_CACHE: set[str] = set()
|
||||
|
||||
_QUICK_START_CUSTOM_PROVIDER_CHOICE = "Other OpenAI-compatible"
|
||||
_QUICK_START_OAUTH_PROVIDERS = {"openai_codex"}
|
||||
|
||||
_CLEAR_CHOICE = "Clear value"
|
||||
_QUICK_START_MENU_CHOICE = "[Q] Quick Start"
|
||||
@ -1576,7 +1581,11 @@ def _get_quick_start_provider_info() -> dict[str, _QuickStartProviderInfo]:
|
||||
|
||||
result: dict[str, _QuickStartProviderInfo] = {}
|
||||
for spec in PROVIDERS:
|
||||
if spec.name == "custom" or spec.is_oauth or spec.is_transcription_only:
|
||||
if (
|
||||
spec.name == "custom"
|
||||
or spec.is_transcription_only
|
||||
or (spec.is_oauth and spec.name not in _QUICK_START_OAUTH_PROVIDERS)
|
||||
):
|
||||
continue
|
||||
result[spec.name] = _QuickStartProviderInfo(
|
||||
display_name=spec.display_name or spec.name,
|
||||
@ -1584,6 +1593,8 @@ def _get_quick_start_provider_info() -> dict[str, _QuickStartProviderInfo]:
|
||||
default_api_base=spec.default_api_base,
|
||||
backend=spec.backend,
|
||||
is_direct=spec.is_direct,
|
||||
is_oauth=spec.is_oauth,
|
||||
default_model=spec.builtin_models[0].id if spec.builtin_models else "",
|
||||
)
|
||||
return result
|
||||
|
||||
@ -1599,7 +1610,71 @@ def _get_quick_start_provider_choices() -> dict[str, str]:
|
||||
|
||||
def _quick_start_requires_api_key(provider_name: str, info: _QuickStartProviderInfo | None) -> bool:
|
||||
"""Return whether Quick Start should ask for an API key."""
|
||||
return provider_name == "custom" or not (info and info.is_local)
|
||||
return provider_name == "custom" or not (info and (info.is_local or info.is_oauth))
|
||||
|
||||
|
||||
def _quick_start_codex_proxy(config: Config) -> str | None:
|
||||
"""Resolve only the Codex proxy without validating unrelated provider secrets."""
|
||||
proxy_config = Config()
|
||||
proxy_config.providers.openai_codex.proxy = config.providers.openai_codex.proxy
|
||||
return resolve_config_env_vars(proxy_config).providers.openai_codex.proxy or None
|
||||
|
||||
|
||||
def _quick_start_oauth_login(config: Config, provider_name: str) -> bool:
|
||||
"""Authenticate an OAuth provider supported by Quick Start."""
|
||||
if provider_name != "openai_codex":
|
||||
console.print(f"[red]OAuth login is not supported for {provider_name}[/red]")
|
||||
return False
|
||||
|
||||
try:
|
||||
from oauth_cli_kit import get_token, login_oauth_interactive
|
||||
except ImportError:
|
||||
console.print("[red]oauth_cli_kit not installed. Run: pip install oauth-cli-kit[/red]")
|
||||
return False
|
||||
|
||||
try:
|
||||
proxy = _quick_start_codex_proxy(config)
|
||||
except ValueError as exc:
|
||||
console.print(f"[red]{escape(str(exc))}[/red]")
|
||||
return False
|
||||
|
||||
token = None
|
||||
with suppress(Exception):
|
||||
token = get_token(proxy=proxy)
|
||||
if not getattr(token, "access", None):
|
||||
console.print("[cyan]Starting interactive OAuth login...[/cyan]\n")
|
||||
try:
|
||||
token = login_oauth_interactive(
|
||||
print_fn=lambda message: console.print(message, markup=False),
|
||||
prompt_fn=lambda prompt: _get_questionary().text(prompt).ask() or "",
|
||||
proxy=proxy,
|
||||
)
|
||||
except Exception as exc:
|
||||
console.print(f"[red]OAuth login failed: {escape(str(exc))}[/red]")
|
||||
return False
|
||||
|
||||
if not getattr(token, "access", None):
|
||||
console.print("[red]OAuth login failed[/red]")
|
||||
return False
|
||||
|
||||
account = getattr(token, "account_id", None)
|
||||
suffix = f" [dim]{escape(str(account))}[/dim]" if account else ""
|
||||
console.print(f"[green]Authenticated with OpenAI Codex[/green]{suffix}")
|
||||
return True
|
||||
|
||||
|
||||
def _quick_start_oauth_is_authenticated(config: Config, provider_name: str) -> bool:
|
||||
"""Return whether Quick Start can load a usable OAuth token."""
|
||||
if provider_name != "openai_codex":
|
||||
return False
|
||||
try:
|
||||
from oauth_cli_kit import get_token
|
||||
|
||||
proxy = _quick_start_codex_proxy(config)
|
||||
token = get_token(proxy=proxy)
|
||||
except Exception:
|
||||
return False
|
||||
return bool(getattr(token, "access", None))
|
||||
|
||||
|
||||
def _quick_start_requires_base_url(provider_name: str, info: _QuickStartProviderInfo | None) -> bool:
|
||||
@ -1710,7 +1785,11 @@ def _configure_quick_start_provider(config: Config) -> bool | object:
|
||||
console.print(f"[red]Unknown provider: {provider_name}[/red]")
|
||||
return False
|
||||
|
||||
model = _input_model_with_autocomplete("Model ID", "", provider_name)
|
||||
model = _input_model_with_autocomplete(
|
||||
"Model ID",
|
||||
provider_info.default_model if provider_info else "",
|
||||
provider_name,
|
||||
)
|
||||
if model is _BACK_PRESSED:
|
||||
continue
|
||||
model = (model or "").strip()
|
||||
@ -1718,6 +1797,10 @@ def _configure_quick_start_provider(config: Config) -> bool | object:
|
||||
console.print("[yellow]! Model ID is required for Quick Start[/yellow]")
|
||||
return False
|
||||
|
||||
if provider_info and provider_info.is_oauth:
|
||||
if not _quick_start_oauth_login(config, provider_name):
|
||||
return False
|
||||
|
||||
if api_key is not None:
|
||||
provider_config.api_key = api_key
|
||||
if api_base:
|
||||
@ -1784,17 +1867,27 @@ def _show_quick_start_summary(config: Config) -> None:
|
||||
_show_quick_start_progress(3)
|
||||
preset = config.model_presets.get("primary")
|
||||
provider_label = "AI provider"
|
||||
has_api_key = True
|
||||
credentials_ready = True
|
||||
credential_name = "API key"
|
||||
if preset:
|
||||
provider_config = getattr(config.providers, preset.provider, None)
|
||||
provider_label, _is_gateway, is_local, _api_base = _get_provider_info().get(
|
||||
preset.provider, (preset.provider, False, False, "")
|
||||
)
|
||||
has_api_key = is_local or bool(provider_config and provider_config.api_key)
|
||||
provider_info = _get_quick_start_provider_info().get(preset.provider)
|
||||
if provider_info:
|
||||
provider_label = provider_info.display_name
|
||||
if provider_info.is_oauth:
|
||||
credential_name = "OAuth login"
|
||||
credentials_ready = _quick_start_oauth_is_authenticated(config, preset.provider)
|
||||
else:
|
||||
credentials_ready = provider_info.is_local or bool(
|
||||
provider_config and provider_config.api_key
|
||||
)
|
||||
else:
|
||||
provider_label = _get_provider_names().get(preset.provider, preset.provider)
|
||||
credentials_ready = bool(provider_config and provider_config.api_key)
|
||||
|
||||
status = "Ready"
|
||||
if not has_api_key:
|
||||
status = f"{provider_label} API key missing"
|
||||
if not credentials_ready:
|
||||
status = f"{provider_label} {credential_name} missing"
|
||||
|
||||
rows = [
|
||||
("Status", status),
|
||||
|
||||
@ -404,16 +404,14 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage:
|
||||
msg = ctx.msg
|
||||
|
||||
async def _run_dream():
|
||||
async def _silent(*_args, **_kwargs):
|
||||
pass
|
||||
|
||||
from nanobot.agent.memory import MemoryStore
|
||||
from nanobot.agent.memory import DreamRunProgress, MemoryStore
|
||||
|
||||
dream_session_key = MemoryStore.dream_session_key
|
||||
build_dream_commit_message = MemoryStore.build_dream_commit_message
|
||||
prune_dream_sessions = MemoryStore.prune_dream_sessions
|
||||
|
||||
store = loop.context.memory
|
||||
progress = DreamRunProgress()
|
||||
content = ""
|
||||
resp = None
|
||||
diff_body = ""
|
||||
@ -434,20 +432,22 @@ async def cmd_dream(ctx: CommandContext) -> OutboundMessage:
|
||||
session_key=key,
|
||||
ephemeral=True,
|
||||
tools=store.build_dream_tools(),
|
||||
on_progress=_silent,
|
||||
on_progress=progress,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
# Ground truth: the real file delta, not the LLM's self-report.
|
||||
# The real file delta grounds the audit record; clean completion
|
||||
# decides whether this history batch has finished processing.
|
||||
diff_body = store.dream_content_diff()
|
||||
productive = bool(diff_body) or (
|
||||
not store.git.is_initialized()
|
||||
and MemoryStore.dream_run_completed(resp)
|
||||
completed = MemoryStore.dream_run_completed(
|
||||
resp,
|
||||
had_tool_errors=progress.had_tool_errors,
|
||||
)
|
||||
if productive:
|
||||
if completed:
|
||||
store.set_last_dream_cursor(last_cursor)
|
||||
content = f"Dream completed in {elapsed:.1f}s."
|
||||
elif MemoryStore.dream_run_completed(resp):
|
||||
content = f"Dream completed in {elapsed:.1f}s; no memory changes."
|
||||
if diff_body:
|
||||
content = f"Dream completed in {elapsed:.1f}s."
|
||||
else:
|
||||
content = f"Dream completed in {elapsed:.1f}s; no memory changes."
|
||||
else:
|
||||
content = (
|
||||
f"Dream did not complete after {elapsed:.1f}s; "
|
||||
|
||||
@ -30,7 +30,7 @@ class ChannelsConfig(Base):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
send_progress: bool = True # stream agent's text progress to the channel
|
||||
send_tool_hints: bool = False # stream tool-call hints (e.g. read_file("…"))
|
||||
send_tool_hints: bool = True # stream tool-call hints (e.g. read_file("…"))
|
||||
show_reasoning: bool = True # surface model reasoning when channel implements it
|
||||
extract_document_text: bool = True # extract text from document attachments before sending to the model
|
||||
send_max_retries: int = Field(default=3, ge=0, le=10) # Max delivery attempts (initial send included)
|
||||
@ -154,6 +154,10 @@ class AgentDefaults(Base):
|
||||
validation_alias=AliasChoices("idleCompactAfterMinutes", "sessionTtlMinutes"),
|
||||
serialization_alias="idleCompactAfterMinutes",
|
||||
) # Auto-compact idle threshold in minutes (0 = disabled)
|
||||
idle_compact_check_interval_seconds: int = Field(
|
||||
default=60,
|
||||
ge=0,
|
||||
) # Minimum interval in seconds between scans for idle sessions
|
||||
consolidation_ratio: float = Field(
|
||||
default=0.5,
|
||||
ge=0.1,
|
||||
@ -195,7 +199,7 @@ class ProviderConfig(Base):
|
||||
extra_headers: dict[str, str] | None = None # Custom headers (e.g. APP-Code for AiHubMix)
|
||||
extra_body: dict[str, Any] | None = None # Extra provider request fields; shape depends on provider/API surface
|
||||
extra_query: dict[str, str] | None = None # Extra query params (e.g. api-version for Azure-style gateways)
|
||||
proxy: str | None = None # OpenAI-compatible/Codex HTTP proxy URL
|
||||
proxy: str | None = None # Explicit HTTP proxy; image downloads trust its DNS and egress
|
||||
thinking_style: str | None = None # Thinking/reasoning style for custom providers
|
||||
|
||||
# Valid values mirror the keys of _THINKING_STYLE_MAP in
|
||||
|
||||
@ -43,9 +43,22 @@ def _load() -> dict[str, Any]:
|
||||
except (json.JSONDecodeError, OSError):
|
||||
logger.warning("Corrupted pairing store, resetting")
|
||||
return {"approved": {}, "pending": {}}
|
||||
if not isinstance(data, dict):
|
||||
logger.warning("Corrupted pairing store, resetting")
|
||||
return {"approved": {}, "pending": {}}
|
||||
|
||||
# JSON stores may contain null maps after partial edits; treat like {}.
|
||||
approved = data.get("approved") or {}
|
||||
if not isinstance(approved, dict):
|
||||
approved = {}
|
||||
data["approved"] = approved
|
||||
pending = data.get("pending") or {}
|
||||
if not isinstance(pending, dict):
|
||||
pending = {}
|
||||
data["pending"] = pending
|
||||
|
||||
# Convert approved lists to str sets for O(1) lookup.
|
||||
for channel, users in data.get("approved", {}).items():
|
||||
for channel, users in approved.items():
|
||||
if not isinstance(users, list):
|
||||
users = []
|
||||
data["approved"][channel] = {str(u) for u in users}
|
||||
@ -56,9 +69,15 @@ def _save(data: dict[str, Any]) -> None:
|
||||
path = _store_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Convert sets back to lists for JSON serialization
|
||||
approved = data.get("approved") or {}
|
||||
pending = data.get("pending") or {}
|
||||
if not isinstance(approved, dict):
|
||||
approved = {}
|
||||
if not isinstance(pending, dict):
|
||||
pending = {}
|
||||
payload = {
|
||||
"approved": {ch: sorted(list(users)) for ch, users in data.get("approved", {}).items()},
|
||||
"pending": dict(data.get("pending", {})),
|
||||
"approved": {ch: sorted(list(users)) for ch, users in approved.items()},
|
||||
"pending": dict(pending),
|
||||
}
|
||||
_write_text_atomic(path, json.dumps(payload, indent=2, ensure_ascii=False))
|
||||
|
||||
@ -66,10 +85,26 @@ def _save(data: dict[str, Any]) -> None:
|
||||
def _gc_pending(data: dict[str, Any]) -> None:
|
||||
"""Remove expired pending entries in-place."""
|
||||
now = time.time()
|
||||
pending: dict[str, Any] = data.get("pending", {})
|
||||
expired = [code for code, info in pending.items() if info.get("expires_at", 0) < now]
|
||||
pending: dict[str, Any] = data.get("pending") or {}
|
||||
if not isinstance(pending, dict):
|
||||
data["pending"] = {}
|
||||
return
|
||||
expired = [
|
||||
code
|
||||
for code, info in pending.items()
|
||||
if (
|
||||
not isinstance(info, dict)
|
||||
or not isinstance(info.get("channel"), str)
|
||||
or not info.get("channel")
|
||||
or info.get("sender_id") is None
|
||||
or isinstance(info.get("expires_at"), bool)
|
||||
or not isinstance(info.get("expires_at"), (int, float))
|
||||
or info["expires_at"] < now
|
||||
)
|
||||
]
|
||||
for code in expired:
|
||||
del pending[code]
|
||||
data["pending"] = pending
|
||||
|
||||
|
||||
def generate_code(
|
||||
@ -152,6 +187,7 @@ def list_pending() -> list[dict[str, Any]]:
|
||||
return [
|
||||
{"code": code, **info}
|
||||
for code, info in data.get("pending", {}).items()
|
||||
if isinstance(info, dict)
|
||||
]
|
||||
|
||||
|
||||
@ -195,6 +231,7 @@ def clear_channel(channel: str) -> dict[str, int]:
|
||||
"""Remove approved senders and pending requests for *channel*."""
|
||||
with _LOCK:
|
||||
data = _load()
|
||||
_gc_pending(data)
|
||||
approved: dict[str, set[str]] = data.get("approved", {})
|
||||
approved_users = approved.pop(channel, set())
|
||||
|
||||
|
||||
@ -10,11 +10,17 @@ from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import httpx
|
||||
from loguru import logger
|
||||
|
||||
from nanobot.providers.registry import find_by_name
|
||||
from nanobot.security.network import (
|
||||
PinnedDNSAsyncTransport,
|
||||
UnsafeURLRequestError,
|
||||
resolve_url_target,
|
||||
)
|
||||
from nanobot.utils.helpers import detect_image_mime
|
||||
|
||||
_OPENROUTER_ATTRIBUTION_HEADERS = {
|
||||
@ -23,6 +29,8 @@ _OPENROUTER_ATTRIBUTION_HEADERS = {
|
||||
"X-OpenRouter-Categories": "cli-agent,personal-agent",
|
||||
}
|
||||
_DEFAULT_TIMEOUT_S = 120.0
|
||||
_IMAGE_DOWNLOAD_MAX_BYTES = 32 * 1024 * 1024
|
||||
_IMAGE_DOWNLOAD_MAX_REDIRECTS = 5
|
||||
_AIHUBMIX_TIMEOUT_S = 300.0
|
||||
_AIHUBMIX_ASPECT_RATIO_SIZES = {
|
||||
"1:1": "1024x1024",
|
||||
@ -33,6 +41,23 @@ _AIHUBMIX_ASPECT_RATIO_SIZES = {
|
||||
}
|
||||
_GEMINI_DEFAULT_TIMEOUT_S = 120.0
|
||||
_GEMINI_IMAGEN_ASPECT_RATIOS = {"1:1", "9:16", "16:9", "3:4", "4:3"}
|
||||
# Aspect ratios documented for every Gemini image model using generateContent.
|
||||
_GEMINI_FLASH_COMMON_ASPECT_RATIOS = {
|
||||
"1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9",
|
||||
}
|
||||
# Gemini 3.1 Flash and Flash Lite additionally accept extreme aspect ratios.
|
||||
_GEMINI_31_FLASH_ASPECT_RATIOS = {
|
||||
*_GEMINI_FLASH_COMMON_ASPECT_RATIOS,
|
||||
"1:4",
|
||||
"4:1",
|
||||
"1:8",
|
||||
"8:1",
|
||||
}
|
||||
# Gemini 3 Pro image models accept these sizes. Gemini 3.1 Flash adds 512,
|
||||
# while Gemini 3.1 Flash Lite supports only 1K.
|
||||
_GEMINI_3_IMAGE_SIZES = {"1K", "2K", "4K"}
|
||||
_GEMINI_31_FLASH_IMAGE_SIZES = {"512", *_GEMINI_3_IMAGE_SIZES}
|
||||
_GEMINI_31_FLASH_LITE_IMAGE_SIZES = {"1K"}
|
||||
_OLLAMA_DEFAULT_SIDE = 1024
|
||||
_OLLAMA_SIZE_PRESETS = {
|
||||
"1K": 1024,
|
||||
@ -114,16 +139,81 @@ def _aihubmix_model_path(model: str) -> str:
|
||||
|
||||
|
||||
async def _download_image_data_url(
|
||||
client: httpx.AsyncClient,
|
||||
url: str,
|
||||
*,
|
||||
proxy: str | None = None,
|
||||
transport: httpx.AsyncBaseTransport | None = None,
|
||||
) -> str:
|
||||
response = await client.get(url)
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
detail = response.text[:500]
|
||||
raise ImageGenerationError(f"failed to download generated image: {detail}") from exc
|
||||
raw = response.content
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"follow_redirects": False,
|
||||
"timeout": _DEFAULT_TIMEOUT_S,
|
||||
"trust_env": False,
|
||||
}
|
||||
if proxy:
|
||||
# An explicit provider proxy is a user-selected trusted egress boundary.
|
||||
# Validate each URL locally, while the proxy owns final DNS resolution.
|
||||
client_kwargs["proxy"] = proxy
|
||||
else:
|
||||
client_kwargs["transport"] = PinnedDNSAsyncTransport(inner=transport)
|
||||
|
||||
async with httpx.AsyncClient(**client_kwargs) as client:
|
||||
current_url = url
|
||||
for _ in range(_IMAGE_DOWNLOAD_MAX_REDIRECTS + 1):
|
||||
if proxy:
|
||||
ok, error, _ = resolve_url_target(
|
||||
current_url,
|
||||
trust_remote_dns=True,
|
||||
)
|
||||
if not ok:
|
||||
raise ImageGenerationError(
|
||||
f"blocked unsafe generated image URL: {error}"
|
||||
)
|
||||
async with client.stream("GET", current_url) as response:
|
||||
if response.is_redirect:
|
||||
location = response.headers.get("location")
|
||||
if not location:
|
||||
raise ImageGenerationError(
|
||||
"generated image URL redirected without a location"
|
||||
)
|
||||
current_url = urljoin(str(response.url), location)
|
||||
continue
|
||||
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise ImageGenerationError(
|
||||
f"failed to download generated image (HTTP {response.status_code})"
|
||||
) from exc
|
||||
|
||||
declared_size = response.headers.get("content-length")
|
||||
if declared_size:
|
||||
try:
|
||||
if int(declared_size) > _IMAGE_DOWNLOAD_MAX_BYTES:
|
||||
raise ImageGenerationError(
|
||||
"generated image exceeded the 32 MiB download limit"
|
||||
)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
chunks: list[bytes] = []
|
||||
total = 0
|
||||
async for chunk in response.aiter_bytes():
|
||||
total += len(chunk)
|
||||
if total > _IMAGE_DOWNLOAD_MAX_BYTES:
|
||||
raise ImageGenerationError(
|
||||
"generated image exceeded the 32 MiB download limit"
|
||||
)
|
||||
chunks.append(chunk)
|
||||
raw = b"".join(chunks)
|
||||
break
|
||||
else:
|
||||
raise ImageGenerationError("generated image URL exceeded the redirect limit")
|
||||
except UnsafeURLRequestError as exc:
|
||||
raise ImageGenerationError(f"blocked unsafe generated image URL: {exc}") from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise ImageGenerationError(f"failed to download generated image: {exc}") from exc
|
||||
|
||||
mime = detect_image_mime(raw)
|
||||
if mime is None:
|
||||
raise ImageGenerationError("generated image URL did not return a supported image")
|
||||
@ -231,6 +321,13 @@ class ImageGenerationProvider(ABC):
|
||||
raise ImageGenerationError(f"{label} returned no images: {provider_error}")
|
||||
raise ImageGenerationError(f"{label} returned no images for this request")
|
||||
|
||||
def _http_client_kwargs(self) -> dict[str, Any]:
|
||||
kwargs: dict[str, Any] = {"timeout": self.timeout}
|
||||
if self.proxy:
|
||||
kwargs["proxy"] = self.proxy
|
||||
kwargs["trust_env"] = False
|
||||
return kwargs
|
||||
|
||||
async def _http_post(
|
||||
self,
|
||||
url: str,
|
||||
@ -243,11 +340,7 @@ class ImageGenerationProvider(ABC):
|
||||
return await client.post(url, headers=headers, json=body)
|
||||
if self._client is not None:
|
||||
return await self._client.post(url, headers=headers, json=body)
|
||||
client_kwargs: dict[str, Any] = {"timeout": self.timeout}
|
||||
if self.proxy:
|
||||
client_kwargs["proxy"] = self.proxy
|
||||
client_kwargs["trust_env"] = False
|
||||
async with httpx.AsyncClient(**client_kwargs) as c:
|
||||
async with httpx.AsyncClient(**self._http_client_kwargs()) as c:
|
||||
return await c.post(url, headers=headers, json=body)
|
||||
|
||||
|
||||
@ -375,7 +468,7 @@ class AIHubMixImageGenerationClient(ImageGenerationProvider):
|
||||
}
|
||||
size = _aihubmix_size(aspect_ratio, image_size)
|
||||
|
||||
client = self._client or httpx.AsyncClient(timeout=self.timeout)
|
||||
client = self._client or httpx.AsyncClient(**self._http_client_kwargs())
|
||||
try:
|
||||
return await self._generate_with_client(
|
||||
client,
|
||||
@ -435,7 +528,7 @@ class AIHubMixImageGenerationClient(ImageGenerationProvider):
|
||||
raise ImageGenerationError(f"AIHubMix image generation failed: {detail}") from exc
|
||||
|
||||
payload = response.json()
|
||||
images = await _aihubmix_images_from_payload(client, payload)
|
||||
images = await _aihubmix_images_from_payload(payload, proxy=self.proxy)
|
||||
|
||||
self._require_images(images, payload)
|
||||
|
||||
@ -635,7 +728,11 @@ class GeminiImageGenerationClient(ImageGenerationProvider):
|
||||
prompt=prompt, model=model, aspect_ratio=aspect_ratio
|
||||
)
|
||||
return await self._generate_gemini_flash(
|
||||
prompt=prompt, model=model, reference_images=reference_images or []
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
reference_images=reference_images or [],
|
||||
aspect_ratio=aspect_ratio,
|
||||
image_size=image_size,
|
||||
)
|
||||
|
||||
async def _generate_imagen(
|
||||
@ -691,15 +788,22 @@ class GeminiImageGenerationClient(ImageGenerationProvider):
|
||||
prompt: str,
|
||||
model: str,
|
||||
reference_images: list[str],
|
||||
aspect_ratio: str | None = None,
|
||||
image_size: str | None = None,
|
||||
) -> GeneratedImageResponse:
|
||||
parts: list[dict[str, Any]] = [
|
||||
{"inlineData": image_path_to_inline_data(path)} for path in reference_images
|
||||
]
|
||||
parts.append({"text": prompt})
|
||||
|
||||
generation_config: dict[str, Any] = {"responseModalities": ["TEXT", "IMAGE"]}
|
||||
image_config = _gemini_flash_image_config(model, aspect_ratio, image_size)
|
||||
if image_config:
|
||||
generation_config["responseFormat"] = {"image": image_config}
|
||||
|
||||
body: dict[str, Any] = {
|
||||
"contents": [{"role": "user", "parts": parts}],
|
||||
"generationConfig": {"responseModalities": ["TEXT", "IMAGE"]},
|
||||
"generationConfig": generation_config,
|
||||
}
|
||||
body.update(self.extra_body)
|
||||
|
||||
@ -748,9 +852,60 @@ class GeminiImageGenerationClient(ImageGenerationProvider):
|
||||
)
|
||||
|
||||
|
||||
def _gemini_flash_image_config(
|
||||
model: str,
|
||||
aspect_ratio: str | None,
|
||||
image_size: str | None,
|
||||
) -> dict[str, str]:
|
||||
"""Build the ``responseFormat.image`` config for Gemini Flash image models.
|
||||
|
||||
Capabilities are model-specific: Gemini 3.1 Flash variants support four
|
||||
additional extreme ratios, while configurable image sizes are limited to
|
||||
the documented Gemini 3 image model families.
|
||||
"""
|
||||
config: dict[str, str] = {}
|
||||
if aspect_ratio and aspect_ratio in _gemini_flash_supported_aspect_ratios(model):
|
||||
config["aspectRatio"] = aspect_ratio
|
||||
if image_size:
|
||||
normalized = image_size.strip().upper()
|
||||
if normalized in _gemini_flash_supported_image_sizes(model):
|
||||
config["imageSize"] = normalized
|
||||
return config
|
||||
|
||||
|
||||
def _gemini_flash_supported_aspect_ratios(model: str) -> set[str]:
|
||||
"""Return the documented aspect ratios for a generateContent image model."""
|
||||
normalized = model.lower()
|
||||
if (
|
||||
"gemini-3.1-flash-lite-image" in normalized
|
||||
or "gemini-3.1-flash-image" in normalized
|
||||
):
|
||||
return _GEMINI_31_FLASH_ASPECT_RATIOS
|
||||
if "gemini-" in normalized and "image" in normalized:
|
||||
return _GEMINI_FLASH_COMMON_ASPECT_RATIOS
|
||||
return set()
|
||||
|
||||
|
||||
def _gemini_flash_supported_image_sizes(model: str) -> set[str]:
|
||||
"""Return the ``imageSize`` values documented for a Flash-path model.
|
||||
|
||||
Earlier Flash image models (2.0, 2.5) expose no configurable size. Gemini
|
||||
3.1 Flash Lite is intentionally checked before the broader Flash match.
|
||||
"""
|
||||
normalized = model.lower()
|
||||
if "gemini-3.1-flash-lite-image" in normalized:
|
||||
return _GEMINI_31_FLASH_LITE_IMAGE_SIZES
|
||||
if "gemini-3.1-flash-image" in normalized:
|
||||
return _GEMINI_31_FLASH_IMAGE_SIZES
|
||||
if "gemini-3-pro-image" in normalized:
|
||||
return _GEMINI_3_IMAGE_SIZES
|
||||
return set()
|
||||
|
||||
|
||||
async def _aihubmix_images_from_payload(
|
||||
client: httpx.AsyncClient,
|
||||
payload: dict[str, Any],
|
||||
*,
|
||||
proxy: str | None = None,
|
||||
) -> list[str]:
|
||||
images: list[str] = []
|
||||
candidates: list[Any] = []
|
||||
@ -768,7 +923,7 @@ async def _aihubmix_images_from_payload(
|
||||
if value.startswith("data:image/"):
|
||||
images.append(value)
|
||||
elif value.startswith(("http://", "https://")):
|
||||
images.append(await _download_image_data_url(client, value))
|
||||
images.append(await _download_image_data_url(value, proxy=proxy))
|
||||
return
|
||||
if not isinstance(value, dict):
|
||||
return
|
||||
@ -969,15 +1124,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
|
||||
return model
|
||||
|
||||
async def _parse_images_response(self, payload: dict[str, Any]) -> list[str]:
|
||||
client = self._client
|
||||
owns_client = client is None
|
||||
if owns_client:
|
||||
client = httpx.AsyncClient(timeout=self.timeout)
|
||||
try:
|
||||
return await _openai_images_from_payload(client, payload)
|
||||
finally:
|
||||
if owns_client:
|
||||
await client.aclose()
|
||||
return await _openai_images_from_payload(payload, proxy=self.proxy)
|
||||
|
||||
async def _post_image_edit(
|
||||
self,
|
||||
@ -1007,7 +1154,7 @@ class OpenAIImageGenerationClient(ImageGenerationProvider):
|
||||
data=body,
|
||||
files=files,
|
||||
)
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as c:
|
||||
async with httpx.AsyncClient(**self._http_client_kwargs()) as c:
|
||||
return await c.post(
|
||||
f"{self.api_base}/images/edits",
|
||||
headers=headers,
|
||||
@ -1188,15 +1335,7 @@ class CustomImageGenerationClient(ImageGenerationProvider):
|
||||
logger.info("Custom Images API response ({}): {}", response.status_code,
|
||||
{k: v for k, v in payload.items() if k != "data"})
|
||||
|
||||
client = self._client
|
||||
owns_client = client is None
|
||||
if owns_client:
|
||||
client = httpx.AsyncClient(timeout=self.timeout)
|
||||
try:
|
||||
images = await _openai_images_from_payload(client, payload)
|
||||
finally:
|
||||
if owns_client:
|
||||
await client.aclose()
|
||||
images = await _openai_images_from_payload(payload, proxy=self.proxy)
|
||||
|
||||
self._require_images(images, payload)
|
||||
|
||||
@ -1389,8 +1528,9 @@ def _openai_explicit_size_supported(
|
||||
|
||||
|
||||
async def _openai_images_from_payload(
|
||||
client: httpx.AsyncClient,
|
||||
payload: dict[str, Any],
|
||||
*,
|
||||
proxy: str | None = None,
|
||||
) -> list[str]:
|
||||
"""Extract images from OpenAI Images API response.
|
||||
|
||||
@ -1406,7 +1546,7 @@ async def _openai_images_from_payload(
|
||||
continue
|
||||
url = item.get("url")
|
||||
if isinstance(url, str) and url:
|
||||
images.append(await _download_image_data_url(client, url))
|
||||
images.append(await _download_image_data_url(url, proxy=proxy))
|
||||
return images
|
||||
|
||||
|
||||
@ -1686,7 +1826,7 @@ class ZhipuImageGenerationClient(ImageGenerationProvider):
|
||||
|
||||
url = f"{self.api_base}/images/generations"
|
||||
|
||||
client = self._client or httpx.AsyncClient(timeout=self.timeout)
|
||||
client = self._client or httpx.AsyncClient(**self._http_client_kwargs())
|
||||
try:
|
||||
return await self._generate_with_client(
|
||||
client,
|
||||
@ -1720,7 +1860,7 @@ class ZhipuImageGenerationClient(ImageGenerationProvider):
|
||||
raise ImageGenerationError(f"Zhipu image generation failed: {detail}") from exc
|
||||
|
||||
payload = response.json()
|
||||
images = await _zhipu_images_from_payload(client, payload)
|
||||
images = await _zhipu_images_from_payload(payload, proxy=self.proxy)
|
||||
|
||||
self._require_images(images, payload)
|
||||
|
||||
@ -1744,8 +1884,9 @@ def _zhipu_size(
|
||||
|
||||
|
||||
async def _zhipu_images_from_payload(
|
||||
client: httpx.AsyncClient,
|
||||
payload: dict[str, Any],
|
||||
*,
|
||||
proxy: str | None = None,
|
||||
) -> list[str]:
|
||||
"""Extract image data URLs from Zhipu API response.
|
||||
|
||||
@ -1758,7 +1899,7 @@ async def _zhipu_images_from_payload(
|
||||
continue
|
||||
url = item.get("url")
|
||||
if isinstance(url, str) and url:
|
||||
images.append(await _download_image_data_url(client, url))
|
||||
images.append(await _download_image_data_url(url, proxy=proxy))
|
||||
return images
|
||||
|
||||
|
||||
@ -1844,7 +1985,7 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider):
|
||||
body.update(self.extra_body)
|
||||
|
||||
url = f"{self.api_base}/images/generations"
|
||||
client = self._client or httpx.AsyncClient(timeout=self.timeout)
|
||||
client = self._client or httpx.AsyncClient(**self._http_client_kwargs())
|
||||
try:
|
||||
return await self._generate_with_client(
|
||||
client,
|
||||
@ -1921,7 +2062,7 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider):
|
||||
status = data.get("task_status")
|
||||
|
||||
if status == "SUCCEED":
|
||||
return await self._collect_images(client, data)
|
||||
return await self._collect_images(data)
|
||||
if status == "FAILED":
|
||||
raise ImageGenerationError(
|
||||
f"ModelScope image generation task failed: {data}"
|
||||
@ -1934,9 +2075,8 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider):
|
||||
f"{_MODELSCOPE_POLL_MAX_ATTEMPTS} polls"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _collect_images(
|
||||
client: httpx.AsyncClient,
|
||||
self,
|
||||
data: dict[str, Any],
|
||||
) -> list[str]:
|
||||
images: list[str] = []
|
||||
@ -1945,7 +2085,9 @@ class ModelScopeImageGenerationClient(ImageGenerationProvider):
|
||||
if url.startswith("data:image/"):
|
||||
images.append(url)
|
||||
else:
|
||||
images.append(await _download_image_data_url(client, url))
|
||||
images.append(
|
||||
await _download_image_data_url(url, proxy=self.proxy)
|
||||
)
|
||||
return images
|
||||
|
||||
|
||||
|
||||
@ -139,6 +139,70 @@ def append_runtime_context(
|
||||
}
|
||||
|
||||
|
||||
def detach_runtime_context(
|
||||
content: Any,
|
||||
marker: Mapping[str, Any],
|
||||
) -> tuple[Any, list[str], list[dict[str, Any]]] | None:
|
||||
"""Detach one validated runtime-context suffix for safe message merging."""
|
||||
if marker.get("version") != 1:
|
||||
return None
|
||||
raw_sources = marker.get("sources")
|
||||
sources = [
|
||||
source
|
||||
for source in raw_sources
|
||||
if isinstance(source, str) and source
|
||||
] if isinstance(raw_sources, list) else []
|
||||
|
||||
suffix = marker.get("suffix")
|
||||
if isinstance(content, str) and isinstance(suffix, str) and suffix:
|
||||
if content == suffix:
|
||||
clean_content = ""
|
||||
elif content.endswith("\n\n" + suffix):
|
||||
clean_content = content[: -(len(suffix) + 2)]
|
||||
else:
|
||||
return None
|
||||
return clean_content, sources, [{"type": "text", "text": suffix}]
|
||||
|
||||
expected = marker.get("blocks")
|
||||
if isinstance(content, list) and isinstance(expected, list) and expected:
|
||||
count = len(expected)
|
||||
if content[-count:] != expected:
|
||||
return None
|
||||
return content[:-count], sources, deepcopy(expected)
|
||||
return None
|
||||
|
||||
|
||||
def reattach_runtime_context(
|
||||
content: Any,
|
||||
sources: Sequence[str],
|
||||
blocks: Sequence[Mapping[str, Any]],
|
||||
) -> tuple[Any, dict[str, Any]]:
|
||||
"""Append detached runtime-context blocks after visible messages are merged."""
|
||||
context_blocks = [deepcopy(dict(block)) for block in blocks]
|
||||
if isinstance(content, str) and all(
|
||||
block.get("type") == "text" and isinstance(block.get("text"), str)
|
||||
for block in context_blocks
|
||||
):
|
||||
suffix = "\n\n".join(block["text"] for block in context_blocks)
|
||||
merged = f"{content}\n\n{suffix}" if content else suffix
|
||||
return merged, {
|
||||
"version": 1,
|
||||
"sources": list(sources),
|
||||
"suffix": suffix,
|
||||
}
|
||||
|
||||
visible_blocks = (
|
||||
[*content]
|
||||
if isinstance(content, list)
|
||||
else ([] if content is None else [{"type": "text", "text": str(content)}])
|
||||
)
|
||||
return [*visible_blocks, *context_blocks], {
|
||||
"version": 1,
|
||||
"sources": list(sources),
|
||||
"blocks": context_blocks,
|
||||
}
|
||||
|
||||
|
||||
def public_history_message(message: Mapping[str, Any]) -> dict[str, Any]:
|
||||
"""Return a user-visible copy with trusted runtime context removed exactly."""
|
||||
cleaned = deepcopy(dict(message))
|
||||
|
||||
@ -20,6 +20,7 @@ _BLOCKED_NETWORKS = [
|
||||
ipaddress.ip_network("169.254.0.0/16"), # link-local / cloud metadata
|
||||
ipaddress.ip_network("172.16.0.0/12"),
|
||||
ipaddress.ip_network("192.168.0.0/16"),
|
||||
ipaddress.ip_network("::/128"), # unspecified; may route to local host
|
||||
ipaddress.ip_network("::1/128"),
|
||||
ipaddress.ip_network("fc00::/7"), # unique local
|
||||
ipaddress.ip_network("fe80::/10"), # link-local v6
|
||||
@ -73,7 +74,12 @@ def _is_private(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
|
||||
return any(normalized in net for net in _BLOCKED_NETWORKS)
|
||||
|
||||
|
||||
def resolve_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool, str, tuple[str, ...]]:
|
||||
def resolve_url_target(
|
||||
url: str,
|
||||
*,
|
||||
allow_loopback: bool = False,
|
||||
trust_remote_dns: bool = False,
|
||||
) -> tuple[bool, str, tuple[str, ...]]:
|
||||
"""Validate a URL is safe to fetch: scheme, hostname, and resolved IPs.
|
||||
|
||||
``allow_loopback`` is intentionally narrow: it only permits literal
|
||||
@ -81,8 +87,14 @@ def resolve_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool,
|
||||
loopback. It does not allow RFC1918, link-local, metadata, or public DNS
|
||||
names that happen to resolve to loopback.
|
||||
|
||||
``trust_remote_dns`` accepts ordinary hostnames unavailable to local DNS.
|
||||
This is only safe when a user-configured trusted proxy owns final DNS
|
||||
resolution and network egress. Localhost names and private/internal IP
|
||||
literals remain blocked.
|
||||
|
||||
Returns (ok, error_message, resolved_ips). When ok is True,
|
||||
resolved_ips contains the public IPs that were validated for this URL.
|
||||
resolved_ips contains the public IPs that were validated for this URL, or
|
||||
is empty when an unresolved hostname is delegated to a trusted proxy.
|
||||
"""
|
||||
try:
|
||||
p = urlparse(url)
|
||||
@ -101,7 +113,20 @@ def resolve_url_target(url: str, *, allow_loopback: bool = False) -> tuple[bool,
|
||||
try:
|
||||
infos = socket.getaddrinfo(hostname, None, socket.AF_UNSPEC, socket.SOCK_STREAM)
|
||||
except socket.gaierror:
|
||||
return False, f"Cannot resolve hostname: {hostname}", ()
|
||||
if not trust_remote_dns:
|
||||
return False, f"Cannot resolve hostname: {hostname}", ()
|
||||
|
||||
normalized_hostname = hostname.rstrip(".").lower()
|
||||
if normalized_hostname == "localhost" or normalized_hostname.endswith(".localhost"):
|
||||
return False, f"Blocked local/internal hostname: {hostname}", ()
|
||||
|
||||
try:
|
||||
literal_addr = ipaddress.ip_address(normalized_hostname)
|
||||
except ValueError:
|
||||
return True, "", ()
|
||||
if _is_private(literal_addr):
|
||||
return False, f"Blocked private/internal address: {literal_addr}", ()
|
||||
return True, "", (str(_normalize_addr(literal_addr)),)
|
||||
|
||||
addrs: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = []
|
||||
for info in infos:
|
||||
|
||||
@ -2,7 +2,11 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping, MutableMapping
|
||||
from typing import Any
|
||||
|
||||
UNIFIED_SESSION_KEY = "unified:default"
|
||||
LAST_CHANNEL_METADATA_KEY = "last_channel"
|
||||
|
||||
|
||||
def session_key_for_channel(channel: str, chat_id: str, *, unified_session: bool = False) -> str:
|
||||
@ -10,3 +14,29 @@ def session_key_for_channel(channel: str, chat_id: str, *, unified_session: bool
|
||||
if unified_session:
|
||||
return UNIFIED_SESSION_KEY
|
||||
return f"{channel}:{chat_id}"
|
||||
|
||||
|
||||
def remember_last_channel(
|
||||
metadata: MutableMapping[str, Any],
|
||||
channel: str,
|
||||
chat_id: str,
|
||||
) -> None:
|
||||
"""Persist the latest concrete delivery route in session metadata."""
|
||||
if not channel or not chat_id:
|
||||
return
|
||||
metadata[LAST_CHANNEL_METADATA_KEY] = f"{channel}:{chat_id}"
|
||||
|
||||
|
||||
def last_channel_from_metadata(
|
||||
metadata: Mapping[str, Any] | None,
|
||||
) -> tuple[str, str] | None:
|
||||
"""Return a concrete delivery route from persisted session metadata."""
|
||||
if not isinstance(metadata, Mapping):
|
||||
return None
|
||||
route = metadata.get(LAST_CHANNEL_METADATA_KEY)
|
||||
if not isinstance(route, str) or ":" not in route:
|
||||
return None
|
||||
channel, chat_id = route.split(":", 1)
|
||||
if not channel or not chat_id:
|
||||
return None
|
||||
return channel, chat_id
|
||||
|
||||
@ -16,6 +16,13 @@ def _int_or_zero(value: Any) -> int:
|
||||
return 0 if value is None or value == "" else int(value)
|
||||
|
||||
|
||||
def _optional_int(value: Any) -> int | None:
|
||||
"""Coerce a stored JSON numeric; null/blank stays None."""
|
||||
if value is None or value == "":
|
||||
return None
|
||||
return int(value)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TriggerRunRecord:
|
||||
"""A single local trigger delivery record."""
|
||||
@ -61,9 +68,10 @@ class LocalTrigger:
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "LocalTrigger":
|
||||
raw_history = data.get("runHistory", data.get("run_history", [])) or []
|
||||
history = [
|
||||
record if isinstance(record, TriggerRunRecord) else TriggerRunRecord.from_dict(record)
|
||||
for record in data.get("runHistory", data.get("run_history", []))
|
||||
for record in raw_history
|
||||
if isinstance(record, (dict, TriggerRunRecord))
|
||||
]
|
||||
return cls(
|
||||
@ -77,7 +85,7 @@ class LocalTrigger:
|
||||
origin_metadata=dict(_get(data, "originMetadata", "origin_metadata", {}) or {}),
|
||||
created_at_ms=_int_or_zero(_get(data, "createdAtMs", "created_at_ms", 0)),
|
||||
updated_at_ms=_int_or_zero(_get(data, "updatedAtMs", "updated_at_ms", 0)),
|
||||
last_run_at_ms=_get(data, "lastRunAtMs", "last_run_at_ms"),
|
||||
last_run_at_ms=_optional_int(_get(data, "lastRunAtMs", "last_run_at_ms")),
|
||||
last_status=_get(data, "lastStatus", "last_status"), # type: ignore[arg-type]
|
||||
last_error=_get(data, "lastError", "last_error"),
|
||||
run_history=history,
|
||||
|
||||
@ -14,6 +14,7 @@ _MAX_REPEAT_EXTERNAL_LOOKUPS = 2
|
||||
|
||||
# Third same-target workspace violation in a turn escalates to "stop retrying".
|
||||
_MAX_REPEAT_WORKSPACE_VIOLATIONS = 2
|
||||
_LENGTH_RECOVERY_TAIL_CHARS = 64
|
||||
|
||||
EMPTY_FINAL_RESPONSE_MESSAGE = (
|
||||
"I completed the tool steps but couldn't produce a final answer. "
|
||||
@ -33,8 +34,10 @@ BUDGET_EXHAUSTED_FINALIZATION_PROMPT = (
|
||||
)
|
||||
|
||||
LENGTH_RECOVERY_PROMPT = (
|
||||
"Output limit reached. Continue exactly where you left off "
|
||||
"— no recap, no apology. Break remaining work into smaller steps if needed."
|
||||
"The previous assistant response was cut off. Continue the same response from its "
|
||||
"exact endpoint. Output only new continuation text in the same language and style. "
|
||||
"Do not acknowledge this instruction, restart the response, repeat its title or any "
|
||||
"existing text, recap, or apologize."
|
||||
)
|
||||
|
||||
SUSTAINED_GOAL_CONTINUE_PROMPT = (
|
||||
@ -79,9 +82,19 @@ def build_budget_exhausted_finalization_message() -> dict[str, str]:
|
||||
return {"role": "user", "content": BUDGET_EXHAUSTED_FINALIZATION_PROMPT}
|
||||
|
||||
|
||||
def build_length_recovery_message() -> dict[str, str]:
|
||||
def build_length_recovery_message(content: str) -> dict[str, str]:
|
||||
"""Prompt the model to continue after hitting output token limit."""
|
||||
return {"role": "user", "content": LENGTH_RECOVERY_PROMPT}
|
||||
tail = content[-_LENGTH_RECOVERY_TAIL_CHARS:]
|
||||
prompt = (
|
||||
f"{LENGTH_RECOVERY_PROMPT}\n\n"
|
||||
"The following tail was already delivered to the user. Treat it as immutable "
|
||||
"context and do not output it again:\n"
|
||||
"<already_delivered_tail>\n"
|
||||
f"{tail}\n"
|
||||
"</already_delivered_tail>\n"
|
||||
"Begin with the text that belongs immediately after this tail."
|
||||
)
|
||||
return {"role": "user", "content": prompt}
|
||||
|
||||
|
||||
def build_goal_continue_message(custom: str | None = None) -> dict[str, str]:
|
||||
|
||||
@ -1770,6 +1770,7 @@ def replay_transcript_to_ui_messages(
|
||||
buffer_message_id = None
|
||||
buffer_parts = []
|
||||
continue
|
||||
merge_next = rec.get("resuming") is True and rec.get("merge_next") is True
|
||||
final_text = rec.get("text")
|
||||
if isinstance(final_text, str):
|
||||
if buffer_message_id is None:
|
||||
@ -1794,8 +1795,11 @@ def replay_transcript_to_ui_messages(
|
||||
**_turn_fields(rec, "answer"),
|
||||
}
|
||||
break
|
||||
buffer_message_id = None
|
||||
buffer_parts = []
|
||||
if merge_next:
|
||||
buffer_parts = [final_text]
|
||||
if not merge_next:
|
||||
buffer_message_id = None
|
||||
buffer_parts = []
|
||||
continue
|
||||
|
||||
if ev == "reasoning_delta":
|
||||
|
||||
@ -11,7 +11,7 @@ from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.command import CommandContext
|
||||
from nanobot.config.schema import AgentDefaults
|
||||
from nanobot.config.schema import AgentDefaults, Config
|
||||
from nanobot.providers.base import LLMResponse
|
||||
|
||||
|
||||
@ -180,12 +180,64 @@ class TestSessionTTLConfig:
|
||||
assert data["idleCompactAfterMinutes"] == 30
|
||||
assert "sessionTtlMinutes" not in data
|
||||
|
||||
def test_idle_scan_interval_defaults_to_sixty_seconds(self):
|
||||
"""The config default should avoid scanning all sessions every idle tick."""
|
||||
defaults = AgentDefaults()
|
||||
assert defaults.idle_compact_check_interval_seconds == 60
|
||||
|
||||
def test_idle_scan_interval_uses_camel_case_config_key(self):
|
||||
"""The JSON config should use the standard camelCase alias."""
|
||||
defaults = AgentDefaults.model_validate({"idleCompactCheckIntervalSeconds": 10})
|
||||
assert defaults.idle_compact_check_interval_seconds == 10
|
||||
data = defaults.model_dump(mode="json", by_alias=True)
|
||||
assert data["idleCompactCheckIntervalSeconds"] == 10
|
||||
|
||||
def test_session_file_cap_is_internal_constant(self):
|
||||
"""Session file cap should remain an internal constant, not a config field."""
|
||||
from nanobot.session.manager import FILE_MAX_MESSAGES
|
||||
assert FILE_MAX_MESSAGES == 2000
|
||||
|
||||
|
||||
class TestIdleScanThrottling:
|
||||
"""Test scheduling of full idle-session scans."""
|
||||
|
||||
def test_configured_idle_scan_interval_throttles_checks(self, tmp_path, monkeypatch):
|
||||
"""The configured interval should reach the loop and gate session scans."""
|
||||
ticks = iter((1_000.0, 1_000.0, 1_009.999, 1_010.0))
|
||||
monkeypatch.setattr("nanobot.agent.loop.time.monotonic", lambda: next(ticks))
|
||||
config = Config.model_validate({
|
||||
"agents": {
|
||||
"defaults": {
|
||||
"workspace": str(tmp_path),
|
||||
"idleCompactCheckIntervalSeconds": 10,
|
||||
}
|
||||
}
|
||||
})
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
loop = AgentLoop.from_config(config, provider=provider)
|
||||
loop.auto_compact.check_expired = MagicMock()
|
||||
|
||||
loop._check_expired_sessions_if_due()
|
||||
loop.auto_compact.check_expired.assert_called_once()
|
||||
loop._check_expired_sessions_if_due()
|
||||
loop.auto_compact.check_expired.assert_called_once()
|
||||
loop._check_expired_sessions_if_due()
|
||||
|
||||
assert loop.auto_compact.check_expired.call_count == 2
|
||||
|
||||
def test_zero_idle_scan_interval_checks_every_tick(self, tmp_path, monkeypatch):
|
||||
"""An explicit zero should leave each idle tick eligible to scan."""
|
||||
monkeypatch.setattr("nanobot.agent.loop.time.monotonic", lambda: 1_000.0)
|
||||
loop = _make_loop(tmp_path)
|
||||
loop.auto_compact.check_expired = MagicMock()
|
||||
|
||||
loop._check_expired_sessions_if_due()
|
||||
loop._check_expired_sessions_if_due()
|
||||
|
||||
assert loop.auto_compact.check_expired.call_count == 2
|
||||
|
||||
|
||||
class TestAgentLoopTTLParam:
|
||||
"""Test that AutoCompact receives and stores session_ttl_minutes."""
|
||||
|
||||
|
||||
@ -534,6 +534,189 @@ class TestToolEventProgress:
|
||||
assert turn_end_msgs[0].content == ""
|
||||
provider.chat_with_retry.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_length_recovery_keeps_one_user_visible_stream(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
) -> 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"),
|
||||
LLMResponse(content="second", finish_reason="stop"),
|
||||
])
|
||||
|
||||
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
|
||||
response = next(responses)
|
||||
await on_content_delta(response.content or "")
|
||||
return response
|
||||
|
||||
provider.chat_stream_with_retry = chat_stream_with_retry
|
||||
provider.chat_with_retry = AsyncMock()
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
||||
_attach_webui_runtime_events(loop, bus)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
await loop._dispatch(InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="u1",
|
||||
chat_id="chat1",
|
||||
content="give a long answer",
|
||||
metadata={"_wants_stream": True},
|
||||
))
|
||||
|
||||
outbound = []
|
||||
while bus.outbound_size > 0:
|
||||
outbound.append(await bus.consume_outbound())
|
||||
|
||||
deltas = [m.event for m in outbound if isinstance(m.event, StreamDeltaEvent)]
|
||||
endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)]
|
||||
|
||||
assert [event.content for event in deltas] == ["first-", "second"]
|
||||
assert [event.resuming for event in endings] == [True, False]
|
||||
assert [event.merge_next for event in endings] == [True, False]
|
||||
assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_length_recovery_streams_non_delta_terminal_segment(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.supports_progress_deltas = True
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
call_count = 0
|
||||
|
||||
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
await on_content_delta("first-")
|
||||
return LLMResponse(content="first-", finish_reason="length")
|
||||
return LLMResponse(content="second", finish_reason="stop")
|
||||
|
||||
provider.chat_stream_with_retry = chat_stream_with_retry
|
||||
provider.chat_with_retry = AsyncMock()
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
||||
_attach_webui_runtime_events(loop, bus)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
await loop._dispatch(InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="u1",
|
||||
chat_id="chat1",
|
||||
content="give a long answer",
|
||||
metadata={"_wants_stream": True},
|
||||
))
|
||||
|
||||
outbound = []
|
||||
while bus.outbound_size > 0:
|
||||
outbound.append(await bus.consume_outbound())
|
||||
|
||||
deltas = [m.event for m in outbound if isinstance(m.event, StreamDeltaEvent)]
|
||||
endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)]
|
||||
final = [m for m in outbound if m.content == "first-second"]
|
||||
|
||||
assert [event.content for event in deltas] == ["first-", "second"]
|
||||
assert [event.merge_next for event in endings] == [True, False]
|
||||
assert len(final) == 1
|
||||
assert isinstance(final[0].event, StreamedResponseEvent)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_length_recovery_at_max_iterations_streams_only_missing_tail(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
) -> 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):
|
||||
await on_content_delta("partial")
|
||||
return LLMResponse(content="partial", finish_reason="length")
|
||||
|
||||
provider.chat_stream_with_retry = chat_stream_with_retry
|
||||
provider.chat_with_retry = AsyncMock(
|
||||
return_value=LLMResponse(content="summary", finish_reason="stop")
|
||||
)
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
||||
_attach_webui_runtime_events(loop, bus)
|
||||
loop.max_iterations = 1
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
|
||||
await loop._dispatch(InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="u1",
|
||||
chat_id="chat1",
|
||||
content="give a long answer",
|
||||
metadata={"_wants_stream": True},
|
||||
))
|
||||
|
||||
outbound = []
|
||||
while bus.outbound_size > 0:
|
||||
outbound.append(await bus.consume_outbound())
|
||||
|
||||
deltas = [m.event for m in outbound if isinstance(m.event, StreamDeltaEvent)]
|
||||
endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)]
|
||||
final = [m for m in outbound if isinstance(m.event, StreamedResponseEvent)]
|
||||
|
||||
assert [event.content for event in deltas] == ["partial", "\n\nsummary"]
|
||||
assert [event.merge_next for event in endings] == [True, False]
|
||||
assert {event.stream_id for event in [*deltas, *endings]} == {deltas[0].stream_id}
|
||||
assert [message.content for message in final] == ["partial\n\nsummary"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_length_recovery_closes_merged_stream(
|
||||
self,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
bus = MessageBus()
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
loop = AgentLoop(bus=bus, provider=provider, workspace=tmp_path, model="test-model")
|
||||
|
||||
async def cancel_after_merge(
|
||||
_msg: InboundMessage,
|
||||
*,
|
||||
on_stream,
|
||||
on_stream_end,
|
||||
**_kwargs,
|
||||
):
|
||||
assert on_stream is not None
|
||||
assert on_stream_end is not None
|
||||
await on_stream("partial")
|
||||
await on_stream_end(resuming=True, merge_next=True)
|
||||
raise asyncio.CancelledError
|
||||
|
||||
loop._process_message = cancel_after_merge # type: ignore[method-assign]
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await loop._dispatch(InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="u1",
|
||||
chat_id="chat1",
|
||||
content="give a long answer",
|
||||
metadata={"_wants_stream": True},
|
||||
))
|
||||
|
||||
outbound = []
|
||||
while bus.outbound_size > 0:
|
||||
outbound.append(await bus.consume_outbound())
|
||||
|
||||
endings = [m.event for m in outbound if isinstance(m.event, StreamEndEvent)]
|
||||
assert [(event.resuming, event.merge_next) for event in endings] == [
|
||||
(True, True),
|
||||
(False, False),
|
||||
]
|
||||
assert endings[0].stream_id == endings[1].stream_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streamed_finalization_is_delivered_as_regular_message(
|
||||
self,
|
||||
|
||||
@ -30,6 +30,10 @@ from nanobot.runtime_context import (
|
||||
)
|
||||
from nanobot.session.automation_turns import AUTOMATION_HISTORY_META
|
||||
from nanobot.session.goal_state import GOAL_STATE_KEY
|
||||
from nanobot.session.keys import (
|
||||
LAST_CHANNEL_METADATA_KEY,
|
||||
UNIFIED_SESSION_KEY,
|
||||
)
|
||||
from nanobot.session.manager import Session, SessionManager
|
||||
from nanobot.session.turn_continuation import (
|
||||
INTERNAL_CONTINUATION_META,
|
||||
@ -682,6 +686,85 @@ async def test_process_message_persists_user_message_before_turn_completes(tmp_p
|
||||
assert persisted.updated_at >= persisted.created_at
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_message_persists_unified_session_delivery_route(tmp_path: Path) -> None:
|
||||
loop = _make_full_loop(tmp_path)
|
||||
loop._unified_session = True
|
||||
loop.consolidator.maybe_consolidate_by_tokens = AsyncMock(return_value=False) # type: ignore[method-assign]
|
||||
loop._run_agent_loop = AsyncMock(side_effect=RuntimeError("boom")) # type: ignore[method-assign]
|
||||
|
||||
msg = InboundMessage(
|
||||
channel="feishu",
|
||||
sender_id="u1",
|
||||
chat_id="oc_123",
|
||||
content="persist my route",
|
||||
session_key_override=UNIFIED_SESSION_KEY,
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
await loop._process_message(msg)
|
||||
|
||||
loop.sessions.invalidate(UNIFIED_SESSION_KEY)
|
||||
persisted = loop.sessions.get_or_create(UNIFIED_SESSION_KEY)
|
||||
assert persisted.metadata[LAST_CHANNEL_METADATA_KEY] == "feishu:oc_123"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("msg", "is_user_turn"),
|
||||
[
|
||||
(
|
||||
InboundMessage(
|
||||
channel="cli",
|
||||
sender_id="u1",
|
||||
chat_id="direct",
|
||||
content="cli input",
|
||||
),
|
||||
True,
|
||||
),
|
||||
(
|
||||
InboundMessage(
|
||||
channel="system",
|
||||
sender_id="system",
|
||||
chat_id="discord:automation",
|
||||
content="system event",
|
||||
),
|
||||
False,
|
||||
),
|
||||
(
|
||||
InboundMessage(
|
||||
channel="discord",
|
||||
sender_id="subagent",
|
||||
chat_id="subagent-result",
|
||||
content="subagent result",
|
||||
),
|
||||
True,
|
||||
),
|
||||
(
|
||||
InboundMessage(
|
||||
channel="discord",
|
||||
sender_id="u1",
|
||||
chat_id="automation",
|
||||
content="scheduled turn",
|
||||
metadata={CRON_TRIGGER_META: {"job_id": "job-1"}},
|
||||
),
|
||||
True,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_unified_session_route_ignores_non_user_destinations(
|
||||
tmp_path: Path,
|
||||
msg: InboundMessage,
|
||||
is_user_turn: bool,
|
||||
) -> None:
|
||||
loop = _make_full_loop(tmp_path)
|
||||
loop._unified_session = True
|
||||
session = loop.sessions.get_or_create(UNIFIED_SESSION_KEY)
|
||||
session.metadata[LAST_CHANNEL_METADATA_KEY] = "telegram:existing"
|
||||
|
||||
loop._remember_unified_session_route(session, msg, is_user_turn=is_user_turn)
|
||||
|
||||
assert session.metadata[LAST_CHANNEL_METADATA_KEY] == "telegram:existing"
|
||||
|
||||
|
||||
# 1x1 PNG used by the media-persistence tests. ``extract_documents`` runs
|
||||
# at the top of ``_process_message`` and filters ``msg.media`` down to
|
||||
# paths that magic-byte-sniff as images, so the test fixture needs real
|
||||
|
||||
@ -978,7 +978,14 @@ class TestMainMenuUpdate:
|
||||
expected_provider_names = set()
|
||||
seen_display_names: set[str] = set()
|
||||
for spec in PROVIDERS:
|
||||
if spec.name == "custom" or spec.is_oauth or spec.is_transcription_only:
|
||||
if (
|
||||
spec.name == "custom"
|
||||
or spec.is_transcription_only
|
||||
or (
|
||||
spec.is_oauth
|
||||
and spec.name not in onboard_wizard._QUICK_START_OAUTH_PROVIDERS
|
||||
)
|
||||
):
|
||||
continue
|
||||
if spec.display_name in seen_display_names:
|
||||
continue
|
||||
@ -988,9 +995,212 @@ class TestMainMenuUpdate:
|
||||
|
||||
assert selected_provider_names == expected_provider_names
|
||||
assert "assemblyai" not in selected_provider_names
|
||||
assert choices["OpenAI Codex"] == "openai_codex"
|
||||
assert "github_copilot" not in selected_provider_names
|
||||
assert choices["OpenCode Zen"] == "opencode"
|
||||
assert choices[onboard_wizard._QUICK_START_CUSTOM_PROVIDER_CHOICE] == "custom"
|
||||
|
||||
def test_quick_start_openai_codex_uses_oauth_and_default_model(self, monkeypatch):
|
||||
"""Codex should authenticate without asking for an API key."""
|
||||
config = Config()
|
||||
oauth_calls: list[tuple[Config, str]] = []
|
||||
model_prompts: list[tuple[str, str, str]] = []
|
||||
|
||||
monkeypatch.setattr(onboard_wizard, "_show_quick_start_progress", lambda *_args: None)
|
||||
monkeypatch.setattr(
|
||||
onboard_wizard,
|
||||
"_select_with_back",
|
||||
lambda *args, **kwargs: "OpenAI Codex",
|
||||
)
|
||||
|
||||
def fail_api_key_prompt(*_args, **_kwargs):
|
||||
raise AssertionError("OpenAI Codex Quick Start should not ask for an API key")
|
||||
|
||||
def fake_model_input(prompt, current, provider):
|
||||
model_prompts.append((prompt, current, provider))
|
||||
return current
|
||||
|
||||
monkeypatch.setattr(onboard_wizard, "_input_text", fail_api_key_prompt)
|
||||
monkeypatch.setattr(onboard_wizard, "_input_model_with_autocomplete", fake_model_input)
|
||||
monkeypatch.setattr(
|
||||
onboard_wizard,
|
||||
"_quick_start_oauth_login",
|
||||
lambda selected_config, provider: oauth_calls.append(
|
||||
(selected_config, provider)
|
||||
)
|
||||
or True,
|
||||
)
|
||||
|
||||
assert onboard_wizard._configure_quick_start_provider(config) is True
|
||||
|
||||
assert oauth_calls == [(config, "openai_codex")]
|
||||
assert model_prompts == [
|
||||
("Model ID", "openai-codex/gpt-5.6-sol", "openai_codex")
|
||||
]
|
||||
assert config.providers.openai_codex.api_key is None
|
||||
assert config.model_presets["primary"].provider == "openai_codex"
|
||||
assert config.model_presets["primary"].model == "openai-codex/gpt-5.6-sol"
|
||||
|
||||
def test_quick_start_openai_codex_login_failure_does_not_create_preset(self, monkeypatch):
|
||||
"""A failed Codex login must not leave a ready-looking model preset."""
|
||||
config = Config()
|
||||
|
||||
monkeypatch.setattr(onboard_wizard, "_show_quick_start_progress", lambda *_args: None)
|
||||
monkeypatch.setattr(
|
||||
onboard_wizard,
|
||||
"_select_with_back",
|
||||
lambda *args, **kwargs: "OpenAI Codex",
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
onboard_wizard,
|
||||
"_input_model_with_autocomplete",
|
||||
lambda *args, **kwargs: "openai-codex/gpt-5.6-sol",
|
||||
)
|
||||
monkeypatch.setattr(onboard_wizard, "_quick_start_oauth_login", lambda *args: False)
|
||||
|
||||
assert onboard_wizard._configure_quick_start_provider(config) is False
|
||||
assert "primary" not in config.model_presets
|
||||
|
||||
def test_quick_start_openai_codex_login_reuses_existing_token(self, monkeypatch):
|
||||
"""Quick Start should not open a new login flow when Codex is already authenticated."""
|
||||
import oauth_cli_kit
|
||||
|
||||
config = Config()
|
||||
config.providers.openai.api_key = "${UNRELATED_MISSING_KEY}"
|
||||
config.providers.openai_codex.proxy = "${CODEX_PROXY}"
|
||||
token = SimpleNamespace(access="existing-token", account_id="account-123")
|
||||
token_proxies: list[str | None] = []
|
||||
login_calls: list[object] = []
|
||||
|
||||
monkeypatch.setenv("CODEX_PROXY", "http://127.0.0.1:8080")
|
||||
monkeypatch.delenv("UNRELATED_MISSING_KEY", raising=False)
|
||||
monkeypatch.setattr(
|
||||
oauth_cli_kit,
|
||||
"get_token",
|
||||
lambda **kwargs: token_proxies.append(kwargs.get("proxy")) or token,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
oauth_cli_kit,
|
||||
"login_oauth_interactive",
|
||||
lambda **kwargs: login_calls.append(kwargs),
|
||||
)
|
||||
monkeypatch.setattr(onboard_wizard.console, "print", lambda *args, **kwargs: None)
|
||||
|
||||
assert onboard_wizard._quick_start_oauth_login(config, "openai_codex") is True
|
||||
assert token_proxies == ["http://127.0.0.1:8080"]
|
||||
assert login_calls == []
|
||||
assert config.providers.openai.api_key == "${UNRELATED_MISSING_KEY}"
|
||||
assert config.providers.openai_codex.proxy == "${CODEX_PROXY}"
|
||||
|
||||
def test_quick_start_openai_codex_runs_interactive_login_for_bad_cached_token(
|
||||
self, monkeypatch
|
||||
):
|
||||
"""A malformed cached token should fall back to the interactive OAuth flow."""
|
||||
import oauth_cli_kit
|
||||
|
||||
config = Config()
|
||||
config.providers.openai_codex.proxy = "http://127.0.0.1:8080"
|
||||
prompts: list[str] = []
|
||||
printed: list[tuple[tuple[object, ...], dict[str, object]]] = []
|
||||
|
||||
class FakePrompt:
|
||||
def ask(self):
|
||||
return "authorization-code"
|
||||
|
||||
def fake_login(**kwargs):
|
||||
kwargs["print_fn"]("[bold]Open the browser[/bold]")
|
||||
prompts.append(kwargs["prompt_fn"]("Paste the authorization code"))
|
||||
assert kwargs["proxy"] == "http://127.0.0.1:8080"
|
||||
return SimpleNamespace(
|
||||
access="fresh-token",
|
||||
account_id="[red]account-123[/red]",
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
oauth_cli_kit,
|
||||
"get_token",
|
||||
lambda **_kwargs: SimpleNamespace(account_id="missing-access"),
|
||||
)
|
||||
monkeypatch.setattr(oauth_cli_kit, "login_oauth_interactive", fake_login)
|
||||
monkeypatch.setattr(
|
||||
onboard_wizard,
|
||||
"_get_questionary",
|
||||
lambda: SimpleNamespace(text=lambda *_args, **_kwargs: FakePrompt()),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
onboard_wizard.console,
|
||||
"print",
|
||||
lambda *args, **kwargs: printed.append((args, kwargs)),
|
||||
)
|
||||
|
||||
assert onboard_wizard._quick_start_oauth_login(config, "openai_codex") is True
|
||||
assert prompts == ["authorization-code"]
|
||||
assert any(
|
||||
args == ("[bold]Open the browser[/bold]",) and kwargs == {"markup": False}
|
||||
for args, kwargs in printed
|
||||
)
|
||||
assert any(r"\[red]account-123\[/red]" in str(args[0]) for args, _kwargs in printed)
|
||||
|
||||
def test_quick_start_codex_auth_check_ignores_unrelated_missing_env(self, monkeypatch):
|
||||
"""OAuth readiness should depend only on the Codex proxy and token."""
|
||||
import oauth_cli_kit
|
||||
|
||||
config = Config()
|
||||
config.providers.anthropic.api_key = "${UNRELATED_MISSING_KEY}"
|
||||
monkeypatch.delenv("UNRELATED_MISSING_KEY", raising=False)
|
||||
monkeypatch.setattr(
|
||||
oauth_cli_kit,
|
||||
"get_token",
|
||||
lambda **kwargs: SimpleNamespace(access="existing-token"),
|
||||
)
|
||||
|
||||
assert (
|
||||
onboard_wizard._quick_start_oauth_is_authenticated(config, "openai_codex")
|
||||
is True
|
||||
)
|
||||
|
||||
def test_quick_start_codex_auth_check_rejects_malformed_token(self, monkeypatch):
|
||||
"""A malformed cached token should report not-ready instead of crashing."""
|
||||
import oauth_cli_kit
|
||||
|
||||
monkeypatch.setattr(
|
||||
oauth_cli_kit,
|
||||
"get_token",
|
||||
lambda **_kwargs: SimpleNamespace(account_id="missing-access"),
|
||||
)
|
||||
|
||||
assert (
|
||||
onboard_wizard._quick_start_oauth_is_authenticated(Config(), "openai_codex")
|
||||
is False
|
||||
)
|
||||
|
||||
def test_quick_start_summary_reports_missing_codex_oauth(self, monkeypatch):
|
||||
"""The review step should distinguish OAuth from an API-key setup."""
|
||||
config = Config()
|
||||
config.model_presets["primary"] = ModelPresetConfig(
|
||||
model="openai-codex/gpt-5.6-sol",
|
||||
provider="openai_codex",
|
||||
)
|
||||
captured: dict[str, list[tuple[str, str]]] = {}
|
||||
|
||||
monkeypatch.setattr(onboard_wizard, "_show_quick_start_progress", lambda *_args: None)
|
||||
monkeypatch.setattr(
|
||||
onboard_wizard,
|
||||
"_quick_start_oauth_is_authenticated",
|
||||
lambda *args: False,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
onboard_wizard,
|
||||
"_print_summary_panel",
|
||||
lambda rows, _title: captured.setdefault("rows", rows),
|
||||
)
|
||||
|
||||
onboard_wizard._show_quick_start_summary(config)
|
||||
|
||||
rows = dict(captured["rows"])
|
||||
assert rows["Status"] == "OpenAI Codex OAuth login missing"
|
||||
assert rows["WebSocket channel"] == "enabled"
|
||||
|
||||
def test_quick_start_provider_choice_skips_advanced_prompts(self, monkeypatch):
|
||||
"""The beginner path should ask for provider credentials and model."""
|
||||
config = Config()
|
||||
|
||||
@ -450,6 +450,104 @@ async def test_runner_uses_specific_message_after_empty_finalization_retry():
|
||||
assert result.stop_reason == "empty_final_response"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_length_recovery_returns_all_segments():
|
||||
"""Recovered output segments are returned together instead of only the tail."""
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content="first ", finish_reason="length"),
|
||||
LLMResponse(content="second ", finish_reason="length"),
|
||||
LLMResponse(content="third", finish_reason="stop"),
|
||||
])
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
|
||||
runner = AgentRunner()
|
||||
result = await runner.run(make_run_spec(provider,
|
||||
initial_messages=[{"role": "user", "content": "give a long answer"}],
|
||||
tools=tools,
|
||||
model="test-model",
|
||||
max_iterations=5,
|
||||
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||
))
|
||||
|
||||
assert result.final_content == "first second third"
|
||||
assert [
|
||||
message["content"]
|
||||
for message in result.messages
|
||||
if message.get("role") == "assistant"
|
||||
] == ["first", "second", "third"]
|
||||
assert provider.chat_with_retry.await_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_length_recovery_preserves_prefix_at_max_iterations():
|
||||
"""Budget exhaustion must not replace output already produced by recovery."""
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(
|
||||
return_value=LLMResponse(content="partial answer", finish_reason="length")
|
||||
)
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
|
||||
runner = AgentRunner()
|
||||
result = await runner.run(make_run_spec(
|
||||
provider,
|
||||
initial_messages=[{"role": "user", "content": "give a long answer"}],
|
||||
tools=tools,
|
||||
model="test-model",
|
||||
max_iterations=1,
|
||||
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||
finalize_on_max_iterations=False,
|
||||
max_iterations_message="limit reached",
|
||||
))
|
||||
|
||||
assert result.stop_reason == "max_iterations"
|
||||
assert result.final_content == "partial answer\n\nlimit reached"
|
||||
assert result.pending_stream_content == "\n\nlimit reached"
|
||||
assert [
|
||||
message["content"]
|
||||
for message in result.messages
|
||||
if message.get("role") == "assistant"
|
||||
] == ["partial answer", "limit reached"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_length_recovery_does_not_leak_across_tool_calls():
|
||||
"""A recovered prefix belongs only to its contiguous response chain."""
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
provider.chat_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content="working", finish_reason="length"),
|
||||
LLMResponse(
|
||||
content=None,
|
||||
tool_calls=[ToolCallRequest(id="call_1", name="read_file", arguments={"path": "x"})],
|
||||
finish_reason="tool_calls",
|
||||
),
|
||||
LLMResponse(content="final answer", finish_reason="stop"),
|
||||
])
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
tools.execute = AsyncMock(return_value="file content")
|
||||
|
||||
runner = AgentRunner()
|
||||
result = await runner.run(make_run_spec(provider,
|
||||
initial_messages=[{"role": "user", "content": "inspect a file"}],
|
||||
tools=tools,
|
||||
model="test-model",
|
||||
max_iterations=5,
|
||||
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||
))
|
||||
|
||||
assert result.final_content == "final answer"
|
||||
assert result.tools_used == ["read_file"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_empty_response_does_not_break_tool_chain():
|
||||
"""An empty intermediate response must not kill an ongoing tool chain.
|
||||
|
||||
@ -143,6 +143,58 @@ async def test_runner_streaming_hook_receives_deltas_and_end_signal():
|
||||
provider.chat_with_retry.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_length_recovery_streams_segments_once_and_returns_all_content():
|
||||
from nanobot.agent.hook import AgentHook, AgentHookContext
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
|
||||
provider = MagicMock(spec=LLMProvider)
|
||||
streamed: list[str] = []
|
||||
endings: list[bool] = []
|
||||
merge_next: list[bool] = []
|
||||
responses = iter([
|
||||
LLMResponse(content="first ", finish_reason="length"),
|
||||
LLMResponse(content="second", finish_reason="stop"),
|
||||
])
|
||||
|
||||
async def chat_stream_with_retry(*, on_content_delta, **kwargs):
|
||||
response = next(responses)
|
||||
await on_content_delta(response.content or "")
|
||||
return response
|
||||
|
||||
provider.chat_stream_with_retry = chat_stream_with_retry
|
||||
provider.chat_with_retry = AsyncMock()
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
|
||||
class StreamingHook(AgentHook):
|
||||
def wants_streaming(self) -> bool:
|
||||
return True
|
||||
|
||||
async def on_stream(self, context: AgentHookContext, delta: str) -> None:
|
||||
streamed.append(delta)
|
||||
|
||||
async def on_stream_end(self, context: AgentHookContext, *, resuming: bool) -> None:
|
||||
endings.append(resuming)
|
||||
merge_next.append(context.stream_continues_current_message)
|
||||
|
||||
runner = AgentRunner()
|
||||
result = await runner.run(make_run_spec(provider,
|
||||
initial_messages=[{"role": "user", "content": "give a long answer"}],
|
||||
tools=tools,
|
||||
model="test-model",
|
||||
max_iterations=3,
|
||||
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||
hook=StreamingHook(),
|
||||
))
|
||||
|
||||
assert result.final_content == "first second"
|
||||
assert streamed == ["first ", "second"]
|
||||
assert endings == [True, False]
|
||||
assert merge_next == [True, False]
|
||||
provider.chat_with_retry.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_passes_cached_tokens_to_hook_context():
|
||||
"""Hook context.usage should contain cached_tokens."""
|
||||
|
||||
@ -352,6 +352,45 @@ async def test_checkpoint2_injects_after_final_response_with_resuming_stream():
|
||||
assert stream_end_calls[-1] is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_injected_followup_starts_new_length_recovery_chain():
|
||||
"""A follow-up gets a fresh recovery budget and no content from the prior answer."""
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
from nanobot.bus.events import InboundMessage
|
||||
|
||||
provider = MagicMock()
|
||||
provider.chat_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content="first-1 ", finish_reason="length"),
|
||||
LLMResponse(content="first-2 ", finish_reason="length"),
|
||||
LLMResponse(content="first-3 ", finish_reason="length"),
|
||||
LLMResponse(content="first-final", finish_reason="stop"),
|
||||
LLMResponse(content="follow-up ", finish_reason="length"),
|
||||
LLMResponse(content="answer", finish_reason="stop"),
|
||||
])
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
|
||||
injection_queue = asyncio.Queue()
|
||||
inject_cb = _make_injection_callback(injection_queue)
|
||||
await injection_queue.put(
|
||||
InboundMessage(channel="cli", sender_id="u", chat_id="c", content="follow-up question")
|
||||
)
|
||||
|
||||
runner = AgentRunner()
|
||||
result = await runner.run(make_run_spec(provider,
|
||||
initial_messages=[{"role": "user", "content": "give a long answer"}],
|
||||
tools=tools,
|
||||
model="test-model",
|
||||
max_iterations=8,
|
||||
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||
injection_callback=inject_cb,
|
||||
))
|
||||
|
||||
assert result.had_injections is True
|
||||
assert result.final_content == "follow-up answer"
|
||||
assert provider.chat_with_retry.await_count == 6
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint2_preserves_final_response_in_history_before_followup():
|
||||
"""A follow-up injected after a final answer must still see that answer in history."""
|
||||
@ -468,6 +507,131 @@ async def test_loop_injected_followup_preserves_image_media(tmp_path):
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_injection_resolves_its_own_runtime_context(tmp_path):
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.runtime_context import (
|
||||
RUNTIME_CONTEXT_MESSAGE_META,
|
||||
RuntimeContextBlock,
|
||||
public_history_message,
|
||||
wrap_runtime_context_lines,
|
||||
)
|
||||
|
||||
provider = MagicMock()
|
||||
provider.get_default_model.return_value = "test-model"
|
||||
provider.chat_with_retry = AsyncMock(side_effect=[
|
||||
LLMResponse(content="first answer", tool_calls=[], usage={}),
|
||||
LLMResponse(content="second answer", tool_calls=[], usage={}),
|
||||
])
|
||||
loop = AgentLoop(
|
||||
bus=MessageBus(),
|
||||
provider=provider,
|
||||
workspace=tmp_path,
|
||||
model="test-model",
|
||||
)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
seen_contexts = []
|
||||
|
||||
async def provide_identity(request):
|
||||
seen_contexts.append((
|
||||
request.channel,
|
||||
request.chat_id,
|
||||
request.sender_id,
|
||||
request.message_id,
|
||||
request.session_key,
|
||||
request.original_user_text,
|
||||
request.metadata["sender_name"],
|
||||
request.metadata["thread_id"],
|
||||
))
|
||||
return RuntimeContextBlock(
|
||||
source="identity",
|
||||
content=wrap_runtime_context_lines([
|
||||
" | ".join(str(value) for value in seen_contexts[-1]),
|
||||
]),
|
||||
)
|
||||
|
||||
loop.register_runtime_context_provider(provide_identity)
|
||||
session = loop.sessions.get_or_create("telegram:group-1")
|
||||
pending_queue = asyncio.Queue()
|
||||
await pending_queue.put(InboundMessage(
|
||||
channel="telegram",
|
||||
sender_id="user-b",
|
||||
chat_id="group-1",
|
||||
content="follow-up from the second speaker",
|
||||
metadata={
|
||||
"message_id": "message-2",
|
||||
"sender_name": "Bob",
|
||||
"thread_id": "topic-7",
|
||||
},
|
||||
))
|
||||
await pending_queue.put(InboundMessage(
|
||||
channel="telegram",
|
||||
sender_id="user-c",
|
||||
chat_id="group-1",
|
||||
content="another follow-up",
|
||||
metadata={
|
||||
"message_id": "message-3",
|
||||
"sender_name": "Carol",
|
||||
"thread_id": "topic-7",
|
||||
},
|
||||
))
|
||||
|
||||
_, _, all_messages, _, _ = await loop._run_agent_loop(
|
||||
[{"role": "user", "content": "initial message from user A"}],
|
||||
runtime=loop.llm_runtime(),
|
||||
session=session,
|
||||
channel="telegram",
|
||||
chat_id="group-1",
|
||||
session_key=session.key,
|
||||
pending_queue=pending_queue,
|
||||
)
|
||||
|
||||
assert seen_contexts == [
|
||||
(
|
||||
"telegram",
|
||||
"group-1",
|
||||
"user-b",
|
||||
"message-2",
|
||||
session.key,
|
||||
"follow-up from the second speaker",
|
||||
"Bob",
|
||||
"topic-7",
|
||||
),
|
||||
(
|
||||
"telegram",
|
||||
"group-1",
|
||||
"user-c",
|
||||
"message-3",
|
||||
session.key,
|
||||
"another follow-up",
|
||||
"Carol",
|
||||
"topic-7",
|
||||
),
|
||||
]
|
||||
|
||||
injected = [message for message in all_messages if message.get("role") == "user"][-1]
|
||||
assert "follow-up from the second speaker" in str(injected["content"])
|
||||
model_messages = provider.chat_with_retry.await_args_list[-1].kwargs["messages"]
|
||||
assert "telegram | group-1 | user-b | message-2" in str(model_messages)
|
||||
assert "Bob | topic-7" in str(model_messages)
|
||||
assert "telegram | group-1 | user-c | message-3" in str(model_messages)
|
||||
assert "Carol | topic-7" in str(model_messages)
|
||||
assert injected["_meta"][RUNTIME_CONTEXT_MESSAGE_META]["sources"] == [
|
||||
"identity",
|
||||
"identity",
|
||||
]
|
||||
|
||||
loop._save_turn(session, all_messages, skip=1)
|
||||
persisted = [message for message in session.messages if message.get("role") == "user"][-1]
|
||||
assert "telegram | group-1 | user-b | message-2" in str(persisted["content"])
|
||||
assert "telegram | group-1 | user-c | message-3" in str(persisted["content"])
|
||||
assert public_history_message(persisted)["content"] == (
|
||||
"follow-up from the second speaker\n\nanother follow-up"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subagent_pending_injection_is_hidden_history_and_not_merged(tmp_path):
|
||||
from nanobot.agent.loop import AgentLoop
|
||||
@ -594,6 +758,58 @@ async def test_runner_merges_multiple_injected_user_messages_without_losing_medi
|
||||
)
|
||||
|
||||
|
||||
def test_runner_merge_preserves_runtime_markers_with_media() -> None:
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
from nanobot.runtime_context import (
|
||||
RUNTIME_CONTEXT_HISTORY_META,
|
||||
RUNTIME_CONTEXT_MESSAGE_META,
|
||||
RuntimeContextBlock,
|
||||
append_runtime_context,
|
||||
public_history_message,
|
||||
)
|
||||
|
||||
first_visible = [
|
||||
{"type": "text", "text": "first"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}},
|
||||
]
|
||||
first_content, first_marker = append_runtime_context(
|
||||
first_visible,
|
||||
[RuntimeContextBlock(source="first", content="private first")],
|
||||
)
|
||||
second_content, second_marker = append_runtime_context(
|
||||
"second",
|
||||
[RuntimeContextBlock(source="second", content="private second")],
|
||||
)
|
||||
messages: list[dict] = []
|
||||
|
||||
AgentRunner._append_injected_messages(messages, [
|
||||
{
|
||||
"role": "user",
|
||||
"content": first_content,
|
||||
"_meta": {RUNTIME_CONTEXT_MESSAGE_META: first_marker},
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": second_content,
|
||||
"_meta": {RUNTIME_CONTEXT_MESSAGE_META: second_marker},
|
||||
},
|
||||
])
|
||||
|
||||
assert len(messages) == 1
|
||||
merged = messages[0]
|
||||
assert "private first" in str(merged["content"])
|
||||
assert "private second" in str(merged["content"])
|
||||
persisted = {
|
||||
"role": "user",
|
||||
"content": merged["content"],
|
||||
RUNTIME_CONTEXT_HISTORY_META: merged["_meta"][RUNTIME_CONTEXT_MESSAGE_META],
|
||||
}
|
||||
assert public_history_message(persisted)["content"] == [
|
||||
*first_visible,
|
||||
{"type": "text", "text": "second"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_injection_cycles_capped_at_max():
|
||||
"""Injection cycles should be capped at _MAX_INJECTION_CYCLES."""
|
||||
@ -1135,7 +1351,7 @@ async def test_dispatch_republishes_leftover_queue_messages(tmp_path):
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_injections_on_fatal_tool_error():
|
||||
"""Pending injections should be drained even when a fatal tool error occurs."""
|
||||
"""A fatal tool error must not leak recovered content into an injected follow-up."""
|
||||
from nanobot.agent.runner import AgentRunner
|
||||
from nanobot.bus.events import InboundMessage
|
||||
|
||||
@ -1145,12 +1361,18 @@ async def test_drain_injections_on_fatal_tool_error():
|
||||
async def chat_with_retry(*, messages, **kwargs):
|
||||
call_count["n"] += 1
|
||||
if call_count["n"] == 1:
|
||||
return LLMResponse(
|
||||
content="stale prefix ",
|
||||
finish_reason="length",
|
||||
usage={},
|
||||
)
|
||||
if call_count["n"] == 2:
|
||||
return LLMResponse(
|
||||
content="",
|
||||
tool_calls=[ToolCallRequest(id="c1", name="exec", arguments={"cmd": "bad"})],
|
||||
usage={},
|
||||
)
|
||||
# Second call: respond normally to the injected follow-up
|
||||
# Third call: respond normally to the injected follow-up.
|
||||
return LLMResponse(content="reply to follow-up", tool_calls=[], usage={})
|
||||
|
||||
provider.chat_with_retry = chat_with_retry
|
||||
@ -1178,6 +1400,7 @@ async def test_drain_injections_on_fatal_tool_error():
|
||||
|
||||
assert result.had_injections is True
|
||||
assert result.final_content == "reply to follow-up"
|
||||
assert call_count["n"] == 3
|
||||
# The injection should be in the messages history
|
||||
injected = [
|
||||
m for m in result.messages
|
||||
|
||||
@ -94,7 +94,12 @@ def test_legacy_stream_metadata_flags_create_runtime_events() -> None:
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
content="",
|
||||
metadata={"_stream_end": True, "_stream_id": "s1", "_resuming": True},
|
||||
metadata={
|
||||
"_stream_end": True,
|
||||
"_stream_id": "s1",
|
||||
"_resuming": True,
|
||||
"_merge_next": True,
|
||||
},
|
||||
)
|
||||
|
||||
delta_event = outbound_event_from_message(delta)
|
||||
@ -106,6 +111,7 @@ def test_legacy_stream_metadata_flags_create_runtime_events() -> None:
|
||||
assert isinstance(end_event, StreamEndEvent)
|
||||
assert end_event.stream_id == "s1"
|
||||
assert end_event.resuming is True
|
||||
assert end_event.merge_next is True
|
||||
|
||||
|
||||
def test_legacy_webui_runtime_metadata_flags_create_runtime_events() -> None:
|
||||
@ -221,7 +227,7 @@ def test_replace_outbound_event_keeps_routing_metadata() -> None:
|
||||
|
||||
updated = replace_outbound_event(
|
||||
msg,
|
||||
StreamEndEvent(stream_id="s1", resuming=True),
|
||||
StreamEndEvent(stream_id="s1", resuming=True, merge_next=True),
|
||||
content="hello world",
|
||||
)
|
||||
|
||||
@ -230,6 +236,7 @@ def test_replace_outbound_event_keeps_routing_metadata() -> None:
|
||||
assert isinstance(updated.event, StreamEndEvent)
|
||||
assert updated.event.stream_id == "s1"
|
||||
assert updated.event.resuming is True
|
||||
assert updated.event.merge_next is True
|
||||
|
||||
|
||||
def test_streamed_response_event_keeps_final_content_outside_event_payload() -> None:
|
||||
|
||||
@ -49,6 +49,7 @@ class MockChannel(BaseChannel):
|
||||
stream_id=None,
|
||||
stream_end=False,
|
||||
resuming=False,
|
||||
merge_next=False,
|
||||
):
|
||||
return await self._send_delta_mock(
|
||||
chat_id,
|
||||
@ -57,6 +58,7 @@ class MockChannel(BaseChannel):
|
||||
stream_id=stream_id,
|
||||
stream_end=stream_end,
|
||||
resuming=resuming,
|
||||
merge_next=merge_next,
|
||||
)
|
||||
|
||||
|
||||
@ -74,7 +76,7 @@ def bus():
|
||||
@pytest.fixture
|
||||
def manager(config, bus):
|
||||
manager = ChannelManager(config, bus)
|
||||
manager.channels["mock"] = MockChannel({}, bus)
|
||||
manager.channels["mock"] = manager._build_channel("mock", MockChannel, {})
|
||||
return manager
|
||||
|
||||
|
||||
@ -92,11 +94,17 @@ def _end(
|
||||
chat_id: str = "chat1",
|
||||
stream_id: str | None = None,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
):
|
||||
return outbound_message_for_event(
|
||||
channel="mock",
|
||||
chat_id=chat_id,
|
||||
event=StreamEndEvent(content=content, stream_id=stream_id, resuming=resuming),
|
||||
event=StreamEndEvent(
|
||||
content=content,
|
||||
stream_id=stream_id,
|
||||
resuming=resuming,
|
||||
merge_next=merge_next,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@ -137,6 +145,7 @@ class TestDeltaCoalescing:
|
||||
stream_id=None,
|
||||
stream_end=False,
|
||||
resuming=False,
|
||||
merge_next=False,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -184,13 +193,19 @@ class TestDeltaCoalescing:
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_end_terminates_coalescing(self, manager, bus):
|
||||
await bus.publish_outbound(_delta("Hello"))
|
||||
await bus.publish_outbound(_end(" world"))
|
||||
await bus.publish_outbound(_end(
|
||||
" world",
|
||||
resuming=True,
|
||||
merge_next=True,
|
||||
))
|
||||
|
||||
first_msg = await bus.consume_outbound()
|
||||
merged, pending = manager._coalesce_stream_deltas(first_msg)
|
||||
|
||||
assert merged.content == "Hello world"
|
||||
assert isinstance(merged.event, StreamEndEvent)
|
||||
assert merged.event.resuming is True
|
||||
assert merged.event.merge_next is True
|
||||
assert len(pending) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -284,14 +299,17 @@ class TestProgressFiltering:
|
||||
|
||||
def test_progress_visibility_uses_global_defaults(self, manager):
|
||||
assert manager._should_send_progress("mock", tool_hint=False) is True
|
||||
assert manager._should_send_progress("mock", tool_hint=True) is False
|
||||
assert manager._should_send_progress("mock", tool_hint=True) is True
|
||||
|
||||
def test_progress_visibility_uses_channel_overrides(self, manager):
|
||||
manager.channels["mock"].send_progress = False
|
||||
manager.channels["mock"].send_tool_hints = True
|
||||
def test_progress_visibility_uses_channel_overrides(self, manager, bus):
|
||||
manager.channels["mock"] = manager._build_channel(
|
||||
"mock",
|
||||
MockChannel,
|
||||
{"sendProgress": False, "sendToolHints": False},
|
||||
)
|
||||
|
||||
assert manager._should_send_progress("mock", tool_hint=False) is False
|
||||
assert manager._should_send_progress("mock", tool_hint=True) is True
|
||||
assert manager._should_send_progress("mock", tool_hint=True) is False
|
||||
|
||||
def test_progress_visibility_returns_false_for_missing_channel(self, manager):
|
||||
assert manager._should_send_progress("nonexistent", tool_hint=False) is False
|
||||
|
||||
@ -269,9 +269,12 @@ def test_channels_config_has_no_per_channel_fields():
|
||||
cfg = ChannelsConfig()
|
||||
assert not hasattr(cfg, "telegram")
|
||||
assert cfg.send_progress is True
|
||||
assert cfg.send_tool_hints is False
|
||||
assert cfg.send_tool_hints is True
|
||||
assert cfg.extract_document_text is True
|
||||
|
||||
opted_out = ChannelsConfig.model_validate({"sendToolHints": False})
|
||||
assert opted_out.send_tool_hints is False
|
||||
|
||||
|
||||
def test_channels_config_extract_document_text_accepts_camel_alias():
|
||||
cfg = ChannelsConfig.model_validate({"extractDocumentText": False})
|
||||
@ -2815,7 +2818,7 @@ async def test_send_with_retry_no_retry_when_max_is_zero():
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_with_retry_calls_send_delta():
|
||||
"""_send_with_retry should call send_delta for stream delta events."""
|
||||
calls: list[tuple[str, str, str | None, bool, bool]] = []
|
||||
calls: list[tuple[str, str, str | None, bool, bool, bool]] = []
|
||||
|
||||
class _StreamingChannel(BaseChannel):
|
||||
name = "streaming"
|
||||
@ -2839,8 +2842,9 @@ async def test_send_with_retry_calls_send_delta():
|
||||
stream_id: str | None = None,
|
||||
stream_end: bool = False,
|
||||
resuming: bool = False,
|
||||
merge_next: bool = False,
|
||||
) -> None:
|
||||
calls.append((chat_id, delta, stream_id, stream_end, resuming))
|
||||
calls.append((chat_id, delta, stream_id, stream_end, resuming, merge_next))
|
||||
|
||||
fake_config = SimpleNamespace(
|
||||
channels=ChannelsConfig(send_max_retries=3),
|
||||
@ -2862,13 +2866,18 @@ async def test_send_with_retry_calls_send_delta():
|
||||
end = outbound_message_for_event(
|
||||
channel="streaming",
|
||||
chat_id="123",
|
||||
event=StreamEndEvent(content="", stream_id="s1", resuming=True),
|
||||
event=StreamEndEvent(
|
||||
content="",
|
||||
stream_id="s1",
|
||||
resuming=True,
|
||||
merge_next=True,
|
||||
),
|
||||
)
|
||||
await mgr._send_with_retry(mgr.channels["streaming"], end)
|
||||
|
||||
assert calls == [
|
||||
("123", "test delta", "s1", False, False),
|
||||
("123", "", "s1", True, True),
|
||||
("123", "test delta", "s1", False, False, False),
|
||||
("123", "", "s1", True, True, True),
|
||||
]
|
||||
|
||||
|
||||
|
||||
@ -1816,6 +1816,42 @@ def test_heartbeat_target_skips_archived_webui_sessions():
|
||||
assert target == ("websocket", "active")
|
||||
|
||||
|
||||
def test_heartbeat_target_uses_last_channel_for_unified_session():
|
||||
from nanobot.cli.commands import _pick_heartbeat_target_from_sessions
|
||||
from nanobot.session.keys import LAST_CHANNEL_METADATA_KEY, UNIFIED_SESSION_KEY
|
||||
|
||||
target = _pick_heartbeat_target_from_sessions(
|
||||
enabled_channels=["telegram", "discord"],
|
||||
archived_keys=[],
|
||||
sessions=[{"key": UNIFIED_SESSION_KEY}],
|
||||
unified_session_metadata={LAST_CHANNEL_METADATA_KEY: "discord:chat-42"},
|
||||
)
|
||||
|
||||
assert target == ("discord", "chat-42")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"metadata",
|
||||
[
|
||||
{"last_channel": "telegram:chat-42"},
|
||||
{"last_channel": "cli:direct"},
|
||||
{"last_channel": "invalid"},
|
||||
],
|
||||
)
|
||||
def test_heartbeat_target_rejects_unroutable_unified_metadata(metadata):
|
||||
from nanobot.cli.commands import _pick_heartbeat_target_from_sessions
|
||||
from nanobot.session.keys import UNIFIED_SESSION_KEY
|
||||
|
||||
target = _pick_heartbeat_target_from_sessions(
|
||||
enabled_channels=["discord"],
|
||||
archived_keys=[],
|
||||
sessions=[{"key": UNIFIED_SESSION_KEY}],
|
||||
unified_session_metadata=metadata,
|
||||
)
|
||||
|
||||
assert target == ("cli", "direct")
|
||||
|
||||
|
||||
def _write_instance_config(tmp_path: Path) -> Path:
|
||||
config_file = tmp_path / "instance" / "config.json"
|
||||
config_file.parent.mkdir(parents=True)
|
||||
|
||||
@ -192,6 +192,7 @@ def _build_runnable_dream(
|
||||
initialized: bool,
|
||||
content_diff: str,
|
||||
stop_reason: str = "completed",
|
||||
tool_error: bool = False,
|
||||
) -> tuple[CommandContext, _FakeStore]:
|
||||
"""Build a /dream ctx whose run is driven by a canned stop reason + diff."""
|
||||
msg = InboundMessage(channel="cli", sender_id="u1", chat_id="direct", content="/dream")
|
||||
@ -203,6 +204,15 @@ def _build_runnable_dream(
|
||||
)
|
||||
|
||||
async def process_direct(*args, **kwargs):
|
||||
if tool_error:
|
||||
await kwargs["on_progress"](
|
||||
"",
|
||||
tool_events=[{
|
||||
"phase": "error",
|
||||
"name": "edit_file",
|
||||
"error": "edit failed",
|
||||
}],
|
||||
)
|
||||
return OutboundMessage(
|
||||
channel="cli",
|
||||
chat_id="direct",
|
||||
@ -225,7 +235,7 @@ def _build_runnable_dream(
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dream_advances_cursor_when_diff_nonempty(tmp_path) -> None:
|
||||
"""A real file delta => productive run => cursor advances (Tier 3)."""
|
||||
"""A completed run with a real file delta advances the cursor."""
|
||||
ctx, store = _build_runnable_dream(tmp_path, initialized=True, content_diff="SOUL.md: +1 -0")
|
||||
await cmd_dream(ctx)
|
||||
await asyncio.sleep(0)
|
||||
@ -233,19 +243,98 @@ async def test_dream_advances_cursor_when_diff_nonempty(tmp_path) -> None:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dream_keeps_cursor_on_completed_noop(tmp_path) -> None:
|
||||
"""Completed run with no file changes must NOT advance the cursor, so the
|
||||
history batch is reconsidered next run instead of silently swallowed."""
|
||||
async def test_dream_advances_cursor_on_completed_noop(tmp_path) -> None:
|
||||
"""A completed no-op has processed the batch and must not repeat it."""
|
||||
ctx, store = _build_runnable_dream(tmp_path, initialized=True, content_diff="")
|
||||
await cmd_dream(ctx)
|
||||
await asyncio.sleep(0)
|
||||
assert store._last_dream_cursor == 5 # unchanged
|
||||
assert store._last_dream_cursor == 42
|
||||
assert "no memory changes" in ctx.loop.bus.outbound[0].content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dream_keeps_cursor_when_incomplete_with_diff(tmp_path) -> None:
|
||||
"""An incomplete run remains retryable even if it left a partial edit."""
|
||||
ctx, store = _build_runnable_dream(
|
||||
tmp_path,
|
||||
initialized=True,
|
||||
content_diff="SOUL.md: +1 -0",
|
||||
stop_reason="length",
|
||||
)
|
||||
await cmd_dream(ctx)
|
||||
await asyncio.sleep(0)
|
||||
assert store._last_dream_cursor == 5
|
||||
assert "did not complete" in ctx.loop.bus.outbound[0].content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dream_keeps_cursor_when_completed_after_tool_error(tmp_path) -> None:
|
||||
"""A soft tool failure must not masquerade as a verified no-op."""
|
||||
ctx, store = _build_runnable_dream(
|
||||
tmp_path,
|
||||
initialized=True,
|
||||
content_diff="",
|
||||
tool_error=True,
|
||||
)
|
||||
await cmd_dream(ctx)
|
||||
await asyncio.sleep(0)
|
||||
assert store._last_dream_cursor == 5
|
||||
assert "did not complete" in ctx.loop.bus.outbound[0].content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dream_noop_batch_unlocks_following_history(tmp_path) -> None:
|
||||
"""A no-op first batch must not starve later history entries."""
|
||||
workspace = tmp_path / "workspace"
|
||||
workspace.mkdir()
|
||||
store = MemoryStore(workspace)
|
||||
store.write_soul("# Soul")
|
||||
store.write_memory("# Memory")
|
||||
for index in range(1, 22):
|
||||
store.append_history(f"entry-{index:02d}")
|
||||
store.git.init()
|
||||
|
||||
processed_prompts: list[str] = []
|
||||
|
||||
async def process_direct(prompt, *args, **kwargs):
|
||||
processed_prompts.append(prompt)
|
||||
return OutboundMessage(
|
||||
channel="cli",
|
||||
chat_id="direct",
|
||||
content="done",
|
||||
metadata={"_stop_reason": "completed"},
|
||||
)
|
||||
|
||||
msg = InboundMessage(channel="cli", sender_id="u1", chat_id="direct", content="/dream")
|
||||
bus = _FakeBus()
|
||||
sessions_dir = tmp_path / "sessions"
|
||||
sessions_dir.mkdir()
|
||||
loop = SimpleNamespace(
|
||||
bus=bus,
|
||||
context=SimpleNamespace(memory=store, timezone="UTC"),
|
||||
sessions=SimpleNamespace(sessions_dir=sessions_dir),
|
||||
process_direct=process_direct,
|
||||
)
|
||||
ctx = CommandContext(msg=msg, session=None, key=msg.session_key, raw="/dream", args="", loop=loop)
|
||||
|
||||
await cmd_dream(ctx)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert len(processed_prompts) == 1
|
||||
assert "entry-20" in processed_prompts[0]
|
||||
assert "entry-21" not in processed_prompts[0]
|
||||
assert store.get_last_dream_cursor() == 20
|
||||
next_result = store.build_dream_prompt()
|
||||
assert next_result is not None
|
||||
next_prompt, next_cursor = next_result
|
||||
assert next_cursor == 21
|
||||
assert "entry-21" in next_prompt
|
||||
assert "entry-01" not in next_prompt
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dream_non_git_falls_back_to_completion_gate(tmp_path) -> None:
|
||||
"""Without git there is no diff signal; productivity falls back to the
|
||||
completion check so non-git workspaces keep working."""
|
||||
"""Non-git workspaces use the same clean-completion gate."""
|
||||
ctx, store = _build_runnable_dream(
|
||||
tmp_path, initialized=False, content_diff="", stop_reason="completed",
|
||||
)
|
||||
|
||||
@ -257,3 +257,50 @@ def test_load_treats_null_approved_channel_list_as_empty(tmp_path, monkeypatch):
|
||||
assert store.is_approved("telegram", "123") is False
|
||||
assert store.is_approved("discord", "456") is True
|
||||
assert store.get_approved("telegram") == []
|
||||
|
||||
|
||||
def test_load_treats_null_approved_and_pending_maps_as_empty(tmp_path, monkeypatch):
|
||||
"""Top-level approved/pending null must not crash pairing load or list_pending."""
|
||||
path = tmp_path / "pairing.json"
|
||||
path.write_text(
|
||||
'{"approved": null, "pending": null}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setattr(store, "_store_path", lambda: path)
|
||||
assert store.is_approved("telegram", "123") is False
|
||||
assert store.list_pending() == []
|
||||
assert store.get_approved("telegram") == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", ["null", "[]", "true"])
|
||||
def test_load_treats_non_object_store_as_empty(tmp_path, monkeypatch, payload):
|
||||
path = tmp_path / "pairing.json"
|
||||
path.write_text(payload, encoding="utf-8")
|
||||
monkeypatch.setattr(store, "_store_path", lambda: path)
|
||||
assert store.list_pending() == []
|
||||
assert store.is_approved("telegram", "123") is False
|
||||
|
||||
|
||||
def test_list_pending_skips_null_pending_entries(tmp_path, monkeypatch):
|
||||
"""Null pending entry values must be dropped instead of crashing list_pending."""
|
||||
path = tmp_path / "pairing.json"
|
||||
path.write_text(
|
||||
'{"approved": {}, "pending": {"ABCD-EFGH": null}}',
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setattr(store, "_store_path", lambda: path)
|
||||
assert store.list_pending() == []
|
||||
assert store.clear_channel("telegram") == {"approved": 0, "pending": 0}
|
||||
|
||||
|
||||
def test_pending_gc_drops_malformed_entries(tmp_path, monkeypatch):
|
||||
path = tmp_path / "pairing.json"
|
||||
path.write_text(
|
||||
'{"approved": {}, "pending": {'
|
||||
'"bad-expiry": {"channel": "telegram", "sender_id": "123", "expires_at": null},'
|
||||
'"missing-sender": {"channel": "telegram", "expires_at": 9999999999}'
|
||||
"}}",
|
||||
encoding="utf-8",
|
||||
)
|
||||
monkeypatch.setattr(store, "_store_path", lambda: path)
|
||||
assert store.list_pending() == []
|
||||
|
||||
@ -102,6 +102,22 @@ class CodexStreamingCompleteThenErrorResponse(FakeResponse):
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def generated_image_downloads(monkeypatch) -> list[tuple[str, str | None]]:
|
||||
"""Keep provider response parsing tests independent from outbound HTTP."""
|
||||
downloads: list[tuple[str, str | None]] = []
|
||||
|
||||
async def download(url: str, *, proxy: str | None = None) -> str:
|
||||
downloads.append((url, proxy))
|
||||
return PNG_DATA_URL
|
||||
|
||||
monkeypatch.setattr(
|
||||
"nanobot.providers.image_generation._download_image_data_url",
|
||||
download,
|
||||
)
|
||||
return downloads
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openrouter_image_generation_payload_and_response(tmp_path: Path) -> None:
|
||||
ref = tmp_path / "ref.png"
|
||||
@ -277,18 +293,22 @@ async def test_aihubmix_image_edit_payload_uses_reference_images(tmp_path: Path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aihubmix_image_generation_downloads_url_response() -> None:
|
||||
async def test_aihubmix_image_generation_downloads_url_response(
|
||||
generated_image_downloads: list[tuple[str, str | None]],
|
||||
) -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
|
||||
fake.get_response = FakeResponse({}, content=PNG_BYTES)
|
||||
proxy = "http://127.0.0.1:23458"
|
||||
client = AIHubMixImageGenerationClient(
|
||||
api_key="sk-ahm-test",
|
||||
proxy=proxy,
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
response = await client.generate(prompt="draw", model="gpt-image-2-free")
|
||||
|
||||
assert response.images[0].startswith("data:image/png;base64,")
|
||||
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
|
||||
assert generated_image_downloads == [("https://cdn.example/image.png", proxy)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -422,6 +442,138 @@ async def test_gemini_flash_reference_images(tmp_path: Path) -> None:
|
||||
assert parts[1] == {"text": "edit this"}
|
||||
|
||||
|
||||
def _gemini_flash_image_response() -> FakeResponse:
|
||||
return FakeResponse(
|
||||
{
|
||||
"candidates": [
|
||||
{"content": {"parts": [{"inlineData": {"mimeType": "image/png", "data": RAW_B64}}]}}
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_flash_forwards_aspect_ratio_and_image_size() -> None:
|
||||
fake = FakeClient(_gemini_flash_image_response())
|
||||
client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type]
|
||||
|
||||
await client.generate(
|
||||
prompt="draw a cat",
|
||||
model="gemini-3-pro-image",
|
||||
aspect_ratio="16:9",
|
||||
image_size="2K",
|
||||
)
|
||||
|
||||
image_config = fake.calls[0]["json"]["generationConfig"]["responseFormat"]["image"]
|
||||
assert image_config == {"aspectRatio": "16:9", "imageSize": "2K"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_flash_2_5_drops_image_size() -> None:
|
||||
fake = FakeClient(_gemini_flash_image_response())
|
||||
client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type]
|
||||
|
||||
await client.generate(
|
||||
prompt="draw a cat",
|
||||
model="gemini-2.5-flash-image",
|
||||
aspect_ratio="4:3",
|
||||
image_size="1K",
|
||||
)
|
||||
|
||||
image_config = fake.calls[0]["json"]["generationConfig"]["responseFormat"]["image"]
|
||||
assert image_config == {"aspectRatio": "4:3"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_flash_2_0_drops_image_size() -> None:
|
||||
fake = FakeClient(_gemini_flash_image_response())
|
||||
client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type]
|
||||
|
||||
await client.generate(
|
||||
prompt="draw a cat",
|
||||
model="gemini-2.0-flash-preview-image-generation",
|
||||
aspect_ratio="16:9",
|
||||
image_size="1K",
|
||||
)
|
||||
|
||||
image_config = fake.calls[0]["json"]["generationConfig"]["responseFormat"]["image"]
|
||||
assert image_config == {"aspectRatio": "16:9"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "aspect_ratio", "expected"),
|
||||
[
|
||||
("gemini-3.1-flash-image", "1:8", {"aspectRatio": "1:8"}),
|
||||
("gemini-3.1-flash-lite-image", "4:1", {"aspectRatio": "4:1"}),
|
||||
("gemini-3-pro-image", "1:8", None),
|
||||
("gemini-2.5-flash-image", "4:1", None),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_flash_scopes_extreme_aspect_ratios_by_model(
|
||||
model: str,
|
||||
aspect_ratio: str,
|
||||
expected: dict[str, str] | None,
|
||||
) -> None:
|
||||
fake = FakeClient(_gemini_flash_image_response())
|
||||
client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type]
|
||||
|
||||
await client.generate(
|
||||
prompt="draw a cat",
|
||||
model=model,
|
||||
aspect_ratio=aspect_ratio,
|
||||
)
|
||||
|
||||
response_format = fake.calls[0]["json"]["generationConfig"].get("responseFormat")
|
||||
assert response_format == ({"image": expected} if expected else None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("model", "image_size", "expected"),
|
||||
[
|
||||
("gemini-3-pro-image", "512", None),
|
||||
("gemini-3-pro", "2K", None),
|
||||
("gemini-3.1-flash-lite-image", "2K", None),
|
||||
("gemini-3.1-flash-lite-image", "1K", {"imageSize": "1K"}),
|
||||
("gemini-3.1-flash-image", "512", {"imageSize": "512"}),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_flash_scopes_image_size_by_model(
|
||||
model: str,
|
||||
image_size: str,
|
||||
expected: dict[str, str] | None,
|
||||
) -> None:
|
||||
fake = FakeClient(_gemini_flash_image_response())
|
||||
client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type]
|
||||
|
||||
await client.generate(
|
||||
prompt="draw a cat",
|
||||
model=model,
|
||||
image_size=image_size,
|
||||
)
|
||||
|
||||
response_format = fake.calls[0]["json"]["generationConfig"].get("responseFormat")
|
||||
assert response_format == ({"image": expected} if expected else None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_flash_ignores_unsupported_hints() -> None:
|
||||
fake = FakeClient(_gemini_flash_image_response())
|
||||
client = GeminiImageGenerationClient(api_key="AIza-test", client=fake) # type: ignore[arg-type]
|
||||
|
||||
# 7:5 is not a documented ratio; 1:8 is only valid for 3.1 Flash, not Pro;
|
||||
# 1024x1024 is not a valid Gemini image-size token. All are dropped.
|
||||
await client.generate(
|
||||
prompt="draw a cat",
|
||||
model="gemini-3-pro-image",
|
||||
aspect_ratio="1:8",
|
||||
image_size="1024x1024",
|
||||
)
|
||||
|
||||
assert "responseFormat" not in fake.calls[0]["json"]["generationConfig"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_requires_api_key() -> None:
|
||||
client = GeminiImageGenerationClient(api_key=None)
|
||||
@ -686,18 +838,22 @@ async def test_openai_b64_json_response_uses_detected_mime() -> None:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_url_download_fallback() -> None:
|
||||
async def test_openai_url_download_fallback(
|
||||
generated_image_downloads: list[tuple[str, str | None]],
|
||||
) -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
|
||||
fake.get_response = FakeResponse({}, content=PNG_BYTES)
|
||||
proxy = "http://127.0.0.1:23458"
|
||||
client = OpenAIImageGenerationClient(
|
||||
api_key="sk-openai-test",
|
||||
proxy=proxy,
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
response = await client.generate(prompt="draw", model="dall-e-3")
|
||||
|
||||
assert response.images[0].startswith("data:image/png;base64,")
|
||||
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
|
||||
assert generated_image_downloads == [("https://cdn.example/image.png", proxy)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -1060,13 +1216,17 @@ async def test_custom_generate_maps_one_k_to_openai_dimension() -> None:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_generate_extra_body_can_override_defaults() -> None:
|
||||
async def test_custom_generate_extra_body_can_override_defaults(
|
||||
generated_image_downloads: list[tuple[str, str | None]],
|
||||
) -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"url": "https://images.example/cat.png"}]}))
|
||||
fake.get_response = FakeResponse({}, content=PNG_BYTES)
|
||||
proxy = "http://127.0.0.1:23458"
|
||||
client = CustomImageGenerationClient(
|
||||
api_key="sk-custom-test",
|
||||
api_base="https://custom.example/v1",
|
||||
extra_body={"response_format": "url", "size": "2K"},
|
||||
proxy=proxy,
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
@ -1076,9 +1236,8 @@ async def test_custom_generate_extra_body_can_override_defaults() -> None:
|
||||
image_size="1K",
|
||||
)
|
||||
|
||||
expected_data_url = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode('ascii')}"
|
||||
assert response.images == [expected_data_url]
|
||||
assert fake.get_calls[0]["url"] == "https://images.example/cat.png"
|
||||
assert response.images == [PNG_DATA_URL]
|
||||
assert generated_image_downloads == [("https://images.example/cat.png", proxy)]
|
||||
body = fake.calls[0]["json"]
|
||||
assert body["response_format"] == "url"
|
||||
assert body["size"] == "2K"
|
||||
@ -1484,18 +1643,22 @@ async def test_zhipu_image_generation_with_explicit_size() -> None:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zhipu_image_generation_downloads_url_response() -> None:
|
||||
async def test_zhipu_image_generation_downloads_url_response(
|
||||
generated_image_downloads: list[tuple[str, str | None]],
|
||||
) -> None:
|
||||
fake = FakeClient(FakeResponse({"data": [{"url": "https://cdn.example/image.png"}]}))
|
||||
fake.get_response = FakeResponse({}, content=PNG_BYTES)
|
||||
proxy = "http://127.0.0.1:23458"
|
||||
client = ZhipuImageGenerationClient(
|
||||
api_key="sk-zhipu-test",
|
||||
proxy=proxy,
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
response = await client.generate(prompt="draw", model="glm-image")
|
||||
|
||||
assert response.images[0].startswith("data:image/png;base64,")
|
||||
assert fake.get_calls[0]["url"] == "https://cdn.example/image.png"
|
||||
assert generated_image_downloads == [("https://cdn.example/image.png", proxy)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -1575,7 +1738,9 @@ def _modelscope_fast_poll(monkeypatch) -> None:
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_modelscope_image_generation_submit_and_poll() -> None:
|
||||
async def test_modelscope_image_generation_submit_and_poll(
|
||||
generated_image_downloads: list[tuple[str, str | None]],
|
||||
) -> None:
|
||||
submit = FakeResponse({"task_id": "abc123"})
|
||||
poll_responses = [
|
||||
FakeResponse({"task_status": "PENDING"}),
|
||||
@ -1585,9 +1750,11 @@ async def test_modelscope_image_generation_submit_and_poll() -> None:
|
||||
}),
|
||||
]
|
||||
fake = ModelScopeFakeClient(submit, poll_responses)
|
||||
proxy = "http://127.0.0.1:23458"
|
||||
client = ModelScopeImageGenerationClient(
|
||||
api_key="ms-token",
|
||||
api_base="https://api-inference.modelscope.cn/v1",
|
||||
proxy=proxy,
|
||||
client=fake, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
@ -1597,6 +1764,7 @@ async def test_modelscope_image_generation_submit_and_poll() -> None:
|
||||
)
|
||||
|
||||
assert response.images[0].startswith("data:image/png;base64,")
|
||||
assert generated_image_downloads == [("https://cdn.example/image.png", proxy)]
|
||||
|
||||
# Verify POST request
|
||||
post_call = fake.calls[0]
|
||||
@ -1766,3 +1934,18 @@ async def test_modelscope_image_generation_poll_timeout(monkeypatch) -> None:
|
||||
|
||||
# Should have polled up to the (patched) attempt limit.
|
||||
assert len(fake.get_calls) == 3
|
||||
|
||||
|
||||
|
||||
def test_image_provider_http_client_kwargs_include_explicit_proxy() -> None:
|
||||
proxy = "http://127.0.0.1:23458"
|
||||
client = AIHubMixImageGenerationClient(
|
||||
api_key="sk-ahm-test",
|
||||
proxy=proxy,
|
||||
)
|
||||
|
||||
assert client._http_client_kwargs() == {
|
||||
"timeout": client.timeout,
|
||||
"proxy": proxy,
|
||||
"trust_env": False,
|
||||
}
|
||||
|
||||
227
tests/providers/test_image_generation_security.py
Normal file
227
tests/providers/test_image_generation_security.py
Normal file
@ -0,0 +1,227 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import socket
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from nanobot.providers import image_generation
|
||||
from nanobot.providers.image_generation import ImageGenerationError, _download_image_data_url
|
||||
|
||||
PNG_BYTES = (
|
||||
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01"
|
||||
b"\x00\x00\x00\x01\x08\x04\x00\x00\x00\xb5\x1c\x0c\x02"
|
||||
b"\x00\x00\x00\x0bIDATx\xdacd\xfc\xff\x1f\x00\x03\x03"
|
||||
b"\x02\x00\xef\xbf\xa7\xdb\x00\x00\x00\x00IEND\xaeB`\x82"
|
||||
)
|
||||
|
||||
|
||||
def _resolve_public(host: str, port: int | None, *args, **kwargs):
|
||||
return [
|
||||
(
|
||||
socket.AF_INET,
|
||||
socket.SOCK_STREAM,
|
||||
socket.IPPROTO_TCP,
|
||||
"",
|
||||
("93.184.216.34", port or 0),
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url",
|
||||
["http://127.0.0.1/admin", "http://[::]/admin"],
|
||||
ids=["ipv4-loopback", "ipv6-unspecified"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"proxy",
|
||||
[None, "http://127.0.0.1:23458"],
|
||||
ids=["direct", "explicit-proxy"],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_image_download_blocks_unsafe_target(
|
||||
url: str,
|
||||
proxy: str | None,
|
||||
) -> None:
|
||||
requested = False
|
||||
|
||||
async def handler(request: httpx.Request) -> httpx.Response:
|
||||
nonlocal requested
|
||||
requested = True
|
||||
return httpx.Response(200, content=PNG_BYTES)
|
||||
|
||||
with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"):
|
||||
await _download_image_data_url(
|
||||
url,
|
||||
proxy=proxy,
|
||||
transport=httpx.MockTransport(handler),
|
||||
)
|
||||
|
||||
assert requested is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_image_download_revalidates_redirects(monkeypatch) -> None:
|
||||
original_getaddrinfo = socket.getaddrinfo
|
||||
|
||||
def resolve_test_hosts(host: str, port: int | None, *args, **kwargs):
|
||||
if host == "cdn.example":
|
||||
return _resolve_public(host, port, *args, **kwargs)
|
||||
return original_getaddrinfo(host, port, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", resolve_test_hosts)
|
||||
requested: list[str] = []
|
||||
|
||||
async def handler(request: httpx.Request) -> httpx.Response:
|
||||
requested.append(str(request.url))
|
||||
return httpx.Response(302, headers={"location": "http://169.254.169.254/latest"})
|
||||
|
||||
with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"):
|
||||
await _download_image_data_url(
|
||||
"https://cdn.example/image.png",
|
||||
transport=httpx.MockTransport(handler),
|
||||
)
|
||||
|
||||
assert requested == ["https://cdn.example/image.png"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_image_download_returns_valid_data_url(monkeypatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"nanobot.security.network.socket.getaddrinfo",
|
||||
_resolve_public,
|
||||
)
|
||||
|
||||
async def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, content=PNG_BYTES)
|
||||
|
||||
result = await _download_image_data_url(
|
||||
"https://cdn.example/image.png",
|
||||
transport=httpx.MockTransport(handler),
|
||||
)
|
||||
|
||||
assert result.startswith("data:image/png;base64,")
|
||||
|
||||
|
||||
class _OversizedStream(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
yield b"12345"
|
||||
yield b"6789"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_image_download_enforces_streaming_size_limit(monkeypatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
"nanobot.security.network.socket.getaddrinfo",
|
||||
_resolve_public,
|
||||
)
|
||||
monkeypatch.setattr(image_generation, "_IMAGE_DOWNLOAD_MAX_BYTES", 8)
|
||||
|
||||
async def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, stream=_OversizedStream())
|
||||
|
||||
with pytest.raises(ImageGenerationError, match="download limit"):
|
||||
await _download_image_data_url(
|
||||
"https://cdn.example/image.png",
|
||||
transport=httpx.MockTransport(handler),
|
||||
)
|
||||
|
||||
|
||||
class _StreamContext:
|
||||
def __init__(self, response: httpx.Response) -> None:
|
||||
self.response = response
|
||||
|
||||
async def __aenter__(self) -> httpx.Response:
|
||||
return self.response
|
||||
|
||||
async def __aexit__(self, exc_type, exc, traceback) -> None:
|
||||
await self.response.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_generated_image_download_delegates_unresolved_host_to_provider_proxy(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
def fail_local_dns(host: str, port: int | None, *args, **kwargs):
|
||||
raise socket.gaierror(f"cannot resolve {host}")
|
||||
|
||||
monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", fail_local_dns)
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class FakeAsyncClient:
|
||||
def __init__(self, **kwargs) -> None:
|
||||
captured["kwargs"] = kwargs
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, traceback) -> None:
|
||||
return None
|
||||
|
||||
def stream(self, method: str, url: str) -> _StreamContext:
|
||||
captured["request"] = (method, url)
|
||||
request = httpx.Request(method, url)
|
||||
return _StreamContext(httpx.Response(200, content=PNG_BYTES, request=request))
|
||||
|
||||
monkeypatch.setattr(image_generation.httpx, "AsyncClient", FakeAsyncClient)
|
||||
proxy = "http://127.0.0.1:23458"
|
||||
|
||||
result = await _download_image_data_url(
|
||||
"https://proxy-only.example/image.png",
|
||||
proxy=proxy,
|
||||
)
|
||||
|
||||
assert result.startswith("data:image/png;base64,")
|
||||
assert captured["request"] == ("GET", "https://proxy-only.example/image.png")
|
||||
assert captured["kwargs"] == {
|
||||
"follow_redirects": False,
|
||||
"timeout": image_generation._DEFAULT_TIMEOUT_S,
|
||||
"trust_env": False,
|
||||
"proxy": proxy,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxied_generated_image_download_revalidates_redirects(
|
||||
monkeypatch,
|
||||
) -> None:
|
||||
original_getaddrinfo = socket.getaddrinfo
|
||||
|
||||
def resolve_test_hosts(host: str, port: int | None, *args, **kwargs):
|
||||
if host == "cdn.example":
|
||||
return _resolve_public(host, port, *args, **kwargs)
|
||||
return original_getaddrinfo(host, port, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr("nanobot.security.network.socket.getaddrinfo", resolve_test_hosts)
|
||||
requested: list[str] = []
|
||||
|
||||
class FakeAsyncClient:
|
||||
def __init__(self, **kwargs) -> None:
|
||||
pass
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, traceback) -> None:
|
||||
return None
|
||||
|
||||
def stream(self, method: str, url: str) -> _StreamContext:
|
||||
requested.append(url)
|
||||
request = httpx.Request(method, url)
|
||||
return _StreamContext(
|
||||
httpx.Response(
|
||||
302,
|
||||
headers={"location": "http://169.254.169.254/latest"},
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
monkeypatch.setattr(image_generation.httpx, "AsyncClient", FakeAsyncClient)
|
||||
|
||||
with pytest.raises(ImageGenerationError, match="blocked unsafe generated image URL"):
|
||||
await _download_image_data_url(
|
||||
"https://cdn.example/image.png",
|
||||
proxy="http://127.0.0.1:23458",
|
||||
)
|
||||
|
||||
assert requested == ["https://cdn.example/image.png"]
|
||||
@ -148,6 +148,7 @@ def test_blocks_sampled_addresses_from_internal_networks():
|
||||
"169.254.0.0/16",
|
||||
"172.16.0.0/12",
|
||||
"192.168.0.0/16",
|
||||
"::/128",
|
||||
"::1/128",
|
||||
"fc00::/7",
|
||||
"fe80::/10",
|
||||
@ -194,6 +195,47 @@ def test_resolve_url_target_returns_validated_public_ips():
|
||||
assert resolved_ips == ("93.184.216.34",)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("trust_remote_dns", "expected_ok"),
|
||||
[(False, False), (True, True)],
|
||||
)
|
||||
def test_resolve_url_target_only_delegates_dns_to_trusted_proxy(
|
||||
trust_remote_dns: bool,
|
||||
expected_ok: bool,
|
||||
):
|
||||
with patch(
|
||||
"nanobot.security.network.socket.getaddrinfo",
|
||||
side_effect=socket.gaierror("local DNS unavailable"),
|
||||
):
|
||||
ok, err, resolved_ips = resolve_url_target(
|
||||
"https://proxy-only.example/image.png",
|
||||
trust_remote_dns=trust_remote_dns,
|
||||
)
|
||||
|
||||
assert ok is expected_ok, err
|
||||
assert resolved_ips == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url",
|
||||
[
|
||||
"http://localhost/secret",
|
||||
"http://service.localhost/secret",
|
||||
"http://127.0.0.1/secret",
|
||||
"http://169.254.169.254/latest",
|
||||
"http://[::1]/secret",
|
||||
],
|
||||
)
|
||||
def test_resolve_url_target_does_not_delegate_local_targets(url: str):
|
||||
with patch(
|
||||
"nanobot.security.network.socket.getaddrinfo",
|
||||
side_effect=socket.gaierror("local DNS unavailable"),
|
||||
):
|
||||
ok, _, _ = resolve_url_target(url, trust_remote_dns=True)
|
||||
|
||||
assert not ok
|
||||
|
||||
|
||||
def test_pin_resolved_url_dns_prevents_second_resolution_rebind():
|
||||
def _rebinding_resolver(hostname, port, family=0, type_=0):
|
||||
return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", ("169.254.169.254", 0))]
|
||||
|
||||
@ -473,6 +473,37 @@ class TestSandboxPlatform:
|
||||
spawned_cmd = mock_spawn.call_args[0][0]
|
||||
assert "bwrap" in spawned_cmd
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bwrap_receives_configured_bind_roots(self, tmp_path):
|
||||
"""Configured bwrap bind roots should be forwarded to the sandbox wrapper."""
|
||||
mock_proc = AsyncMock()
|
||||
mock_proc.communicate.return_value = (b"sandboxed", b"")
|
||||
mock_proc.returncode = 0
|
||||
tool_bin = tmp_path / "tool-bin"
|
||||
tool_cache = tmp_path / "tool-cache"
|
||||
|
||||
with (
|
||||
patch("nanobot.agent.tools.shell._IS_WINDOWS", False),
|
||||
patch("nanobot.agent.tools.shell.wrap_command", return_value="bwrap -- sh -c ls") as mock_wrap,
|
||||
patch.object(ExecTool, "_spawn", return_value=mock_proc),
|
||||
patch.object(ExecTool, "_guard_command", return_value=None),
|
||||
):
|
||||
tool = ExecTool(
|
||||
sandbox="bwrap",
|
||||
working_dir="/workspace",
|
||||
sandbox_ro_binds=[str(tool_bin)],
|
||||
sandbox_rw_binds=[str(tool_cache)],
|
||||
)
|
||||
await tool.execute(command="ls")
|
||||
|
||||
kwargs = mock_wrap.call_args.kwargs
|
||||
assert kwargs["sandbox_ro_binds"] == [
|
||||
str(tool_bin.resolve(strict=False))
|
||||
]
|
||||
assert kwargs["sandbox_rw_binds"] == [
|
||||
str(tool_cache.resolve(strict=False))
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# end-to-end (mocked subprocess, full execute path)
|
||||
|
||||
@ -314,6 +314,103 @@ def test_exec_still_blocks_real_outside_path_via_redirect(tmp_path):
|
||||
assert "path outside working dir" in blocked
|
||||
|
||||
|
||||
def test_exec_allows_absolute_path_inside_bwrap_ro_bind(tmp_path, monkeypatch):
|
||||
workspace = tmp_path / "workspace"
|
||||
workspace.mkdir()
|
||||
tool_bin = tmp_path / "home" / ".local" / "bin"
|
||||
tool_bin.mkdir(parents=True)
|
||||
uv = tool_bin / "uv"
|
||||
uv.write_text("#!/bin/sh\n")
|
||||
monkeypatch.setattr("nanobot.agent.tools.shell._IS_WINDOWS", False)
|
||||
tool = ExecTool(
|
||||
working_dir=str(workspace),
|
||||
restrict_to_workspace=True,
|
||||
sandbox="bwrap",
|
||||
sandbox_ro_binds=[str(tool_bin)],
|
||||
)
|
||||
|
||||
blocked = tool._guard_command(
|
||||
f"{uv} --version",
|
||||
str(workspace),
|
||||
restrict_to_workspace=True,
|
||||
workspace_root=str(workspace),
|
||||
)
|
||||
|
||||
assert blocked is None
|
||||
|
||||
|
||||
def test_exec_allows_absolute_path_inside_bwrap_rw_bind(tmp_path, monkeypatch):
|
||||
workspace = tmp_path / "workspace"
|
||||
workspace.mkdir()
|
||||
cache_dir = tmp_path / "cache"
|
||||
cache_dir.mkdir()
|
||||
monkeypatch.setattr("nanobot.agent.tools.shell._IS_WINDOWS", False)
|
||||
tool = ExecTool(
|
||||
working_dir=str(workspace),
|
||||
restrict_to_workspace=True,
|
||||
sandbox="bwrap",
|
||||
sandbox_rw_binds=[str(cache_dir)],
|
||||
)
|
||||
|
||||
blocked = tool._guard_command(
|
||||
f"touch {cache_dir / 'stamp'}",
|
||||
str(workspace),
|
||||
restrict_to_workspace=True,
|
||||
workspace_root=str(workspace),
|
||||
)
|
||||
|
||||
assert blocked is None
|
||||
|
||||
|
||||
def test_exec_bind_roots_do_not_widen_guard_without_bwrap(tmp_path):
|
||||
workspace = tmp_path / "workspace"
|
||||
workspace.mkdir()
|
||||
tool_bin = tmp_path / "home" / ".local" / "bin"
|
||||
tool_bin.mkdir(parents=True)
|
||||
uv = tool_bin / "uv"
|
||||
uv.write_text("#!/bin/sh\n")
|
||||
tool = ExecTool(
|
||||
working_dir=str(workspace),
|
||||
restrict_to_workspace=True,
|
||||
sandbox="",
|
||||
sandbox_ro_binds=[str(tool_bin)],
|
||||
)
|
||||
|
||||
blocked = tool._guard_command(
|
||||
f"{uv} --version",
|
||||
str(workspace),
|
||||
restrict_to_workspace=True,
|
||||
workspace_root=str(workspace),
|
||||
)
|
||||
|
||||
assert blocked is not None
|
||||
assert "path outside working dir" in blocked
|
||||
|
||||
|
||||
def test_exec_bwrap_bind_parent_does_not_widen_workspace_guard(tmp_path, monkeypatch):
|
||||
workspace = tmp_path / "workspace"
|
||||
workspace.mkdir()
|
||||
secret = tmp_path / "config.json"
|
||||
secret.write_text("secret")
|
||||
monkeypatch.setattr("nanobot.agent.tools.shell._IS_WINDOWS", False)
|
||||
tool = ExecTool(
|
||||
working_dir=str(workspace),
|
||||
restrict_to_workspace=True,
|
||||
sandbox="bwrap",
|
||||
sandbox_ro_binds=[str(tmp_path)],
|
||||
)
|
||||
|
||||
blocked = tool._guard_command(
|
||||
f"cat {secret}",
|
||||
str(workspace),
|
||||
restrict_to_workspace=True,
|
||||
workspace_root=str(workspace),
|
||||
)
|
||||
|
||||
assert blocked is not None
|
||||
assert "path outside working dir" in blocked
|
||||
|
||||
|
||||
# --- format command blocking -----------------------------------------------
|
||||
|
||||
|
||||
|
||||
@ -236,6 +236,100 @@ def test_wrapper_normalizes_nullable_property_anyof() -> None:
|
||||
}
|
||||
|
||||
|
||||
def test_wrapper_hoists_recursive_local_refs_into_defs() -> None:
|
||||
recursive_items_ref = "#/properties/filter/properties/items"
|
||||
tool_def = SimpleNamespace(
|
||||
name="search_dataset",
|
||||
description="search tool",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"filter": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"items": {
|
||||
"type": "array",
|
||||
"items": {"$ref": recursive_items_ref},
|
||||
}
|
||||
},
|
||||
"required": ["items"],
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)
|
||||
|
||||
generated_ref = wrapper.parameters["properties"]["filter"]["properties"]["items"][
|
||||
"items"
|
||||
]["$ref"]
|
||||
assert generated_ref.startswith("#/$defs/ref_")
|
||||
generated_name = generated_ref.removeprefix("#/$defs/")
|
||||
generated_schema = wrapper.parameters["$defs"][generated_name]
|
||||
assert generated_schema["type"] == "array"
|
||||
assert generated_schema["items"]["$ref"] == generated_ref
|
||||
|
||||
|
||||
def test_wrapper_hoists_root_self_ref_into_defs() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="tree",
|
||||
description="tree tool",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"children": {"type": "array", "items": {"$ref": "#"}},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)
|
||||
|
||||
generated_ref = wrapper.parameters["properties"]["children"]["items"]["$ref"]
|
||||
assert generated_ref.startswith("#/$defs/ref_")
|
||||
generated_name = generated_ref.removeprefix("#/$defs/")
|
||||
assert wrapper.parameters["$defs"][generated_name]["properties"]["children"]["items"] == {
|
||||
"$ref": generated_ref
|
||||
}
|
||||
|
||||
|
||||
def test_wrapper_preserves_existing_defs_refs() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="demo",
|
||||
description="demo tool",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"$defs": {"value": {"type": "string"}},
|
||||
"properties": {"value": {"$ref": "#/$defs/value"}},
|
||||
},
|
||||
)
|
||||
|
||||
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)
|
||||
|
||||
assert wrapper.parameters["properties"]["value"]["$ref"] == "#/$defs/value"
|
||||
assert wrapper.parameters["$defs"]["value"]["type"] == "string"
|
||||
|
||||
|
||||
def test_wrapper_resolves_uri_encoded_json_pointer() -> None:
|
||||
tool_def = SimpleNamespace(
|
||||
name="demo",
|
||||
description="demo tool",
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"space name/value": {"type": "string"},
|
||||
"alias": {"$ref": "#/properties/space%20name~1value"},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
wrapper = MCPToolWrapper(SimpleNamespace(call_tool=None), "test", tool_def)
|
||||
|
||||
generated_ref = wrapper.parameters["properties"]["alias"]["$ref"]
|
||||
assert generated_ref.startswith("#/$defs/ref_")
|
||||
generated_name = generated_ref.removeprefix("#/$defs/")
|
||||
assert wrapper.parameters["$defs"][generated_name] == {"type": "string"}
|
||||
|
||||
|
||||
def test_normalize_windows_stdio_command_is_noop_off_windows(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
@ -150,6 +150,78 @@ class TestBwrapBackend:
|
||||
try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in try_indices}
|
||||
assert (str(fake_media), str(fake_media)) in try_pairs
|
||||
|
||||
def test_custom_read_only_binds_use_ro_bind_try(self, tmp_path):
|
||||
ws = tmp_path / "project"
|
||||
tool_bin = tmp_path / "home" / ".local" / "bin"
|
||||
|
||||
result = wrap_command(
|
||||
"bwrap",
|
||||
"uv --version",
|
||||
str(ws),
|
||||
str(ws),
|
||||
sandbox_ro_binds=[str(tool_bin)],
|
||||
)
|
||||
tokens = _parse(result)
|
||||
|
||||
try_indices = [i for i, t in enumerate(tokens) if t == "--ro-bind-try"]
|
||||
try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in try_indices}
|
||||
assert (str(tool_bin.resolve(strict=False)), str(tool_bin.resolve(strict=False))) in try_pairs
|
||||
|
||||
def test_custom_read_write_binds_use_bind_try(self, tmp_path):
|
||||
ws = tmp_path / "project"
|
||||
cache_dir = tmp_path / "cache"
|
||||
|
||||
result = wrap_command(
|
||||
"bwrap",
|
||||
"touch cache/file",
|
||||
str(ws),
|
||||
str(ws),
|
||||
sandbox_rw_binds=[str(cache_dir)],
|
||||
)
|
||||
tokens = _parse(result)
|
||||
|
||||
bind_try_indices = [i for i, t in enumerate(tokens) if t == "--bind-try"]
|
||||
bind_try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in bind_try_indices}
|
||||
resolved = str(cache_dir.resolve(strict=False))
|
||||
assert (resolved, resolved) in bind_try_pairs
|
||||
|
||||
def test_custom_relative_bind_paths_are_ignored(self, tmp_path):
|
||||
ws = tmp_path / "project"
|
||||
|
||||
result = wrap_command(
|
||||
"bwrap",
|
||||
"ls",
|
||||
str(ws),
|
||||
str(ws),
|
||||
sandbox_ro_binds=["relative/bin"],
|
||||
sandbox_rw_binds=["relative/cache"],
|
||||
)
|
||||
tokens = _parse(result)
|
||||
|
||||
assert "relative/bin" not in tokens
|
||||
assert "relative/cache" not in tokens
|
||||
|
||||
def test_custom_workspace_parent_binds_are_ignored(self, tmp_path):
|
||||
ws = tmp_path / "private" / "project"
|
||||
parent = ws.parent.resolve(strict=False)
|
||||
|
||||
result = wrap_command(
|
||||
"bwrap",
|
||||
"cat ../config.json",
|
||||
str(ws),
|
||||
str(ws),
|
||||
sandbox_ro_binds=[str(parent)],
|
||||
sandbox_rw_binds=[str(parent)],
|
||||
)
|
||||
tokens = _parse(result)
|
||||
|
||||
ro_try_indices = [i for i, token in enumerate(tokens) if token == "--ro-bind-try"]
|
||||
ro_try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in ro_try_indices}
|
||||
bind_try_indices = [i for i, token in enumerate(tokens) if token == "--bind-try"]
|
||||
bind_try_pairs = {(tokens[i + 1], tokens[i + 2]) for i in bind_try_indices}
|
||||
assert (str(parent), str(parent)) not in ro_try_pairs
|
||||
assert (str(parent), str(parent)) not in bind_try_pairs
|
||||
|
||||
|
||||
class TestUnknownBackend:
|
||||
def test_raises_value_error(self, tmp_path):
|
||||
|
||||
@ -714,6 +714,22 @@ def test_exec_config_timeout_uncapped_and_zero() -> None:
|
||||
ExecToolConfig(timeout=-1)
|
||||
|
||||
|
||||
def test_exec_config_accepts_bwrap_bind_aliases() -> None:
|
||||
cfg = ExecToolConfig.model_validate(
|
||||
{
|
||||
"sandboxRoBinds": ["/home/user/.local/bin"],
|
||||
"sandboxRwBinds": ["/home/user/.cache/uv"],
|
||||
}
|
||||
)
|
||||
|
||||
dumped = cfg.model_dump(by_alias=True)
|
||||
|
||||
assert cfg.sandbox_ro_binds == ["/home/user/.local/bin"]
|
||||
assert cfg.sandbox_rw_binds == ["/home/user/.cache/uv"]
|
||||
assert dumped["sandboxRoBinds"] == ["/home/user/.local/bin"]
|
||||
assert dumped["sandboxRwBinds"] == ["/home/user/.cache/uv"]
|
||||
|
||||
|
||||
def test_resolve_timeout_config_uncapped_and_unlimited() -> None:
|
||||
"""Config timeout drives the hard timeout uncapped; 0 means no limit (#3595)."""
|
||||
assert ExecTool(timeout=3600)._resolve_timeout(None) == 3600
|
||||
|
||||
@ -585,3 +585,53 @@ def test_local_trigger_from_dict_accepts_null_run_at_ms() -> None:
|
||||
)
|
||||
assert delivery.created_at_ms == 0
|
||||
assert delivery.attempts == 0
|
||||
|
||||
|
||||
def test_local_trigger_from_dict_coerces_string_last_run_at_ms() -> None:
|
||||
"""String lastRunAtMs must coerce to int like cron store ms fields."""
|
||||
trigger = LocalTrigger.from_dict(
|
||||
{
|
||||
"id": "t1",
|
||||
"name": "n",
|
||||
"enabled": True,
|
||||
"channel": "websocket",
|
||||
"chatId": "c1",
|
||||
"sessionKey": "websocket:c1",
|
||||
"lastRunAtMs": "1710000000000",
|
||||
"createdAtMs": 1,
|
||||
"updatedAtMs": 1,
|
||||
}
|
||||
)
|
||||
assert trigger.last_run_at_ms == 1710000000000
|
||||
assert trigger.last_run_at_ms < 1710000000001
|
||||
|
||||
trigger_null = LocalTrigger.from_dict(
|
||||
{
|
||||
"id": "t2",
|
||||
"name": "n",
|
||||
"enabled": True,
|
||||
"sessionKey": "websocket:c1",
|
||||
"lastRunAtMs": None,
|
||||
"createdAtMs": 1,
|
||||
"updatedAtMs": 1,
|
||||
}
|
||||
)
|
||||
assert trigger_null.last_run_at_ms is None
|
||||
|
||||
|
||||
def test_local_trigger_from_dict_accepts_null_run_history() -> None:
|
||||
"""Null runHistory must load as empty, matching CronJobState.from_store_dict."""
|
||||
trigger = LocalTrigger.from_dict(
|
||||
{
|
||||
"id": "t1",
|
||||
"name": "n",
|
||||
"enabled": True,
|
||||
"channel": "websocket",
|
||||
"chatId": "c1",
|
||||
"sessionKey": "websocket:c1",
|
||||
"runHistory": None,
|
||||
"createdAtMs": 1,
|
||||
"updatedAtMs": 1,
|
||||
}
|
||||
)
|
||||
assert trigger.run_history == []
|
||||
|
||||
16
tests/utils/test_length_recovery_runtime.py
Normal file
16
tests/utils/test_length_recovery_runtime.py
Normal file
@ -0,0 +1,16 @@
|
||||
"""Tests for length-recovery prompt construction."""
|
||||
|
||||
from nanobot.utils.runtime import build_length_recovery_message
|
||||
|
||||
|
||||
def test_length_recovery_message_anchors_the_existing_tail() -> None:
|
||||
omitted_prefix = "OMITTED_PREFIX"
|
||||
tail = "x" * 64
|
||||
|
||||
message = build_length_recovery_message(omitted_prefix + tail)
|
||||
|
||||
assert message["role"] == "user"
|
||||
assert omitted_prefix not in message["content"]
|
||||
assert f"<already_delivered_tail>\n{tail}\n</already_delivered_tail>" in message["content"]
|
||||
assert "Output only new continuation text" in message["content"]
|
||||
assert "Break remaining work into smaller steps" not in message["content"]
|
||||
@ -1121,6 +1121,26 @@ def test_replay_keeps_interrupted_pre_tool_text_in_activity() -> None:
|
||||
assert msgs[2]["content"] == "Done. Open index.html to play."
|
||||
|
||||
|
||||
def test_replay_merges_length_recovery_segments_into_one_assistant_message() -> None:
|
||||
msgs = replay_transcript_to_ui_messages([
|
||||
{"event": "delta", "chat_id": "t-stream", "text": "first "},
|
||||
{
|
||||
"event": "stream_end",
|
||||
"chat_id": "t-stream",
|
||||
"text": "first ",
|
||||
"resuming": True,
|
||||
"merge_next": True,
|
||||
},
|
||||
{"event": "delta", "chat_id": "t-stream", "text": "second"},
|
||||
{"event": "stream_end", "chat_id": "t-stream"},
|
||||
{"event": "turn_end", "chat_id": "t-stream"},
|
||||
])
|
||||
|
||||
assert len(msgs) == 1
|
||||
assert msgs[0]["role"] == "assistant"
|
||||
assert msgs[0]["content"] == "first second"
|
||||
|
||||
|
||||
def test_replay_tool_events_dedupes_finish_after_start() -> None:
|
||||
msgs = replay_transcript_to_ui_messages([
|
||||
{
|
||||
|
||||
@ -638,7 +638,7 @@ export const ThreadViewport = forwardRef<ThreadViewportHandle, ThreadViewportPro
|
||||
<div
|
||||
ref={messageRegionRef}
|
||||
data-testid="thread-message-region"
|
||||
className="row-start-1 flex min-h-0 flex-col justify-start px-3 pb-4 pt-4 sm:px-4"
|
||||
className="row-start-1 flex min-h-0 min-w-0 flex-col justify-start px-3 pb-4 pt-4 sm:px-4"
|
||||
>
|
||||
<div ref={messageContentRef} className="mx-auto w-full max-w-[49.5rem]">
|
||||
<ThreadMessages
|
||||
|
||||
@ -26,7 +26,7 @@ import type {
|
||||
} from "@/lib/types";
|
||||
|
||||
interface StreamBuffer {
|
||||
/** ID of the assistant message currently receiving deltas (cleared on ``stream_end``). */
|
||||
/** ID of the assistant message currently receiving deltas (cleared when its segment closes). */
|
||||
messageId: string;
|
||||
}
|
||||
|
||||
@ -780,15 +780,20 @@ export function useNanobotStream(
|
||||
?? findStreamingAssistantIndex(next, closedAssistantStreamIdsRef.current, turn);
|
||||
if (targetIndex !== null) {
|
||||
const target = next[targetIndex];
|
||||
next = replaceMessageAt(next, targetIndex, {
|
||||
const merged = {
|
||||
...target,
|
||||
content: finalAnswerText,
|
||||
isStreaming: true,
|
||||
...turn,
|
||||
});
|
||||
};
|
||||
next = replaceMessageAt(next, targetIndex, merged);
|
||||
if (!options?.closeAnswerSegment) {
|
||||
closedAssistantStreamIdsRef.current.delete(merged.id);
|
||||
activeAssistantRef.current = { id: merged.id, index: targetIndex };
|
||||
buffer.current = { messageId: merged.id };
|
||||
}
|
||||
} else {
|
||||
const id = crypto.randomUUID();
|
||||
closedAssistantStreamIdsRef.current.add(id);
|
||||
next = [
|
||||
...next,
|
||||
{
|
||||
@ -800,6 +805,12 @@ export function useNanobotStream(
|
||||
createdAt: Date.now(),
|
||||
},
|
||||
];
|
||||
if (options?.closeAnswerSegment) {
|
||||
closedAssistantStreamIdsRef.current.add(id);
|
||||
} else {
|
||||
activeAssistantRef.current = { id, index: next.length - 1 };
|
||||
buffer.current = { messageId: id };
|
||||
}
|
||||
}
|
||||
}
|
||||
if (options?.closeAnswerSegment) closeActiveAssistantStream();
|
||||
@ -911,8 +922,9 @@ export function useNanobotStream(
|
||||
|
||||
if (ev.event === "stream_end") {
|
||||
const turn = turnFieldsFromEvent(ev, "answer");
|
||||
const mergeNext = ev.resuming === true && ev.merge_next === true;
|
||||
flushPendingStreamEvents({
|
||||
closeAnswerSegment: true,
|
||||
closeAnswerSegment: !mergeNext,
|
||||
...(typeof ev.text === "string" ? { finalAnswerText: ev.text } : {}),
|
||||
turn,
|
||||
});
|
||||
@ -920,7 +932,9 @@ export function useNanobotStream(
|
||||
if (ev.resuming) {
|
||||
cancelStreamEndTimer();
|
||||
setIsStreaming(true);
|
||||
setMessages((prev) => finalizeStreamedTurn(prev, turn));
|
||||
if (!mergeNext) {
|
||||
setMessages((prev) => finalizeStreamedTurn(prev, turn));
|
||||
}
|
||||
return;
|
||||
}
|
||||
scheduleStreamEndTimer(turn);
|
||||
|
||||
@ -1189,6 +1189,8 @@ export type InboundEvent =
|
||||
text?: string;
|
||||
/** This answer segment ended, but the active agent turn will continue. */
|
||||
resuming?: boolean;
|
||||
/** The next answer segment continues this same assistant message. */
|
||||
merge_next?: boolean;
|
||||
} & InboundTurnMetadata)
|
||||
| ({
|
||||
event: "reasoning_delta";
|
||||
|
||||
@ -223,6 +223,28 @@ describe("ThreadViewport", () => {
|
||||
expect(messageRegion.className).not.toContain("5rem");
|
||||
});
|
||||
|
||||
it("allows long messages to shrink within the shared mobile grid column", () => {
|
||||
render(
|
||||
<ThreadViewport
|
||||
messages={[
|
||||
...messages,
|
||||
{
|
||||
id: "a-long-link",
|
||||
role: "assistant",
|
||||
content:
|
||||
"https://github.com/HKUDS/nanobot/discussions/17788077"
|
||||
+ "/a-very-long-unbroken-segment-that-must-not-widen-the-thread",
|
||||
createdAt: Date.now(),
|
||||
},
|
||||
]}
|
||||
isStreaming={false}
|
||||
composer={<div>composer</div>}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByTestId("thread-message-region")).toHaveClass("min-w-0");
|
||||
});
|
||||
|
||||
it("top-aligns a short active turn while the agent is responding", () => {
|
||||
render(
|
||||
<ThreadViewport
|
||||
|
||||
@ -2009,6 +2009,79 @@ describe("useNanobotStream", () => {
|
||||
expect(result.current.messages.every((message) => !message.isStreaming)).toBe(true);
|
||||
});
|
||||
|
||||
it("keeps length-recovery segments in one assistant message", async () => {
|
||||
const fake = fakeClient();
|
||||
const { result } = renderHook(() => useNanobotStream("chat-length", EMPTY_MESSAGES), {
|
||||
wrapper: wrap(fake.client),
|
||||
});
|
||||
|
||||
act(() => {
|
||||
result.current.send("give a long answer");
|
||||
});
|
||||
const activeTurnId = fake.client.sendMessage.mock.calls.at(-1)![3]?.turnId;
|
||||
|
||||
act(() => {
|
||||
fake.emit("chat-length", {
|
||||
event: "delta",
|
||||
chat_id: "chat-length",
|
||||
text: "first ",
|
||||
turn_id: activeTurnId,
|
||||
});
|
||||
});
|
||||
await flushStreamFrame();
|
||||
const assistantId = result.current.messages[1].id;
|
||||
|
||||
act(() => {
|
||||
fake.emit("chat-length", {
|
||||
event: "stream_end",
|
||||
chat_id: "chat-length",
|
||||
text: "first ",
|
||||
resuming: true,
|
||||
merge_next: true,
|
||||
turn_id: activeTurnId,
|
||||
});
|
||||
});
|
||||
|
||||
expect(result.current.messages).toHaveLength(2);
|
||||
expect(result.current.messages[1]).toMatchObject({
|
||||
id: assistantId,
|
||||
content: "first ",
|
||||
isStreaming: true,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
fake.emit("chat-length", {
|
||||
event: "delta",
|
||||
chat_id: "chat-length",
|
||||
text: "second",
|
||||
turn_id: activeTurnId,
|
||||
});
|
||||
});
|
||||
await flushStreamFrame();
|
||||
|
||||
expect(result.current.messages).toHaveLength(2);
|
||||
expect(result.current.messages[1]).toMatchObject({
|
||||
id: assistantId,
|
||||
content: "first second",
|
||||
isStreaming: true,
|
||||
});
|
||||
|
||||
act(() => {
|
||||
fake.emit("chat-length", {
|
||||
event: "turn_end",
|
||||
chat_id: "chat-length",
|
||||
turn_id: activeTurnId,
|
||||
});
|
||||
});
|
||||
|
||||
expect(result.current.messages).toHaveLength(2);
|
||||
expect(result.current.messages[1]).toMatchObject({
|
||||
id: assistantId,
|
||||
content: "first second",
|
||||
isStreaming: false,
|
||||
});
|
||||
});
|
||||
|
||||
it("keeps streaming alive across stream_end when tool activity follows", async () => {
|
||||
const fake = fakeClient();
|
||||
const onTurnEnd = vi.fn();
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user