fix(tools): bound computer use runtime state

This commit is contained in:
Xubin Ren
2026-08-09 01:27:22 +09:00
parent 9859e02215
commit a739185740
11 changed files with 284 additions and 30 deletions
+6 -2
View File
@@ -1705,14 +1705,17 @@ controls an isolated Playwright page.
| `tools.browser.enable` | `false` | Register the DOM-based `browser` tool | | `tools.browser.enable` | `false` | Register the DOM-based `browser` tool |
| `tools.browser.allowedDomains` | `[]` | Optional top-level navigation allowlist; entries include subdomains | | `tools.browser.allowedDomains` | `[]` | Optional top-level navigation allowlist; entries include subdomains |
| `tools.browser.includeScreenshot` | `false` | Attach a screenshot after browser actions | | `tools.browser.includeScreenshot` | `false` | Attach a screenshot after browser actions |
| `tools.browser.maxSessions` | `8` | Maximum retained browser sessions; least-recently-used state is closed first |
| `tools.computerUse.enable` | `false` | Register pixel-based `computer_use` | | `tools.computerUse.enable` | `false` | Register pixel-based `computer_use` |
| `tools.computerUse.backend` | `"desktop"` | `"desktop"` or `"browser"` | | `tools.computerUse.backend` | `"desktop"` | `"desktop"` or `"browser"` |
| `tools.computerUse.allowedDomains` | `[]` | Navigation allowlist for the browser backend | | `tools.computerUse.allowedDomains` | `[]` | Navigation allowlist for the browser backend |
| `tools.computerUse.targetWidth` / `targetHeight` | `1280` / `800` | Maximum screenshot dimensions exposed to the model | | `tools.computerUse.targetWidth` / `targetHeight` | `1280` / `800` | Maximum screenshot dimensions exposed to the model |
| `tools.computerUse.maxSessions` | `8` | Maximum retained sessions for the browser backend |
Each nanobot session gets separate browser state. Browser HTTP and WebSocket traffic passes Each nanobot session gets separate browser state. Browser HTTP and WebSocket traffic passes
through the shared SSRF policy; local, private, link-local, and metadata targets are blocked through the shared SSRF policy; local, private, link-local, and metadata targets are blocked
unless explicitly permitted with `tools.ssrfWhitelist`. `file:` URLs are not accepted. unless explicitly permitted with `tools.ssrfWhitelist`. When `maxSessions` is reached, the
least-recently-used browser state is closed. `file:` URLs are not accepted.
> [!IMPORTANT] > [!IMPORTANT]
> Browser URL checks are defense in depth, not an egress sandbox: Chromium performs its own DNS > Browser URL checks are defense in depth, not an egress sandbox: Chromium performs its own DNS
@@ -1721,7 +1724,8 @@ unless explicitly permitted with `tools.ssrfWhitelist`. `file:` URLs are not acc
> [!WARNING] > [!WARNING]
> The desktop backend can click, type, and change state outside the workspace. Enabling it is an > The desktop backend can click, type, and change state outside the workspace. Enabling it is an
> explicit trust decision: use a trusted model and input source, and run nanobot in a disposable > explicit trust decision: use a trusted model and input source, and run nanobot in a disposable
> OS account or VM when unattended. The workspace restriction is not an OS sandbox. > OS account or VM when unattended. The workspace restriction is not an OS sandbox. Desktop text
> input supports ASCII key events; use the browser backend when Unicode text input is required.
## Web Tools ## Web Tools
+44
View File
@@ -33,6 +33,11 @@ COMPACTABLE_TOOLS = frozenset({
"read_file", "exec", "grep", "find_files", "read_file", "exec", "grep", "find_files",
"web_search", "web_fetch", "list_dir", "list_exec_sessions", "web_search", "web_fetch", "list_dir", "list_exec_sessions",
}) })
VISUAL_TOOLS = frozenset({"browser", "computer_use"})
STALE_SCREENSHOT_PLACEHOLDER = {
"type": "text",
"text": "[Earlier screenshot omitted; use the latest screenshot from this tool.]",
}
# read_file is the recovery path for persisted results; exempting it prevents persist->read->persist loops. # read_file is the recovery path for persisted results; exempting it prevents persist->read->persist loops.
TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"}) TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"})
BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]" BACKFILL_CONTENT = "[Tool result unavailable — call was interrupted or lost]"
@@ -41,6 +46,12 @@ PLACEHOLDER_TEXTS = frozenset({
}) })
def _is_image_block(value: object) -> bool:
if not isinstance(value, dict):
return False
return cast(dict[str, Any], value).get("type") in {"image_url", "input_image"}
def _tool_call_name_is_valid(tool_call: Any) -> bool: def _tool_call_name_is_valid(tool_call: Any) -> bool:
"""Whether a persisted OpenAI-style tool_call carries a usable name. """Whether a persisted OpenAI-style tool_call carries a usable name.
@@ -84,6 +95,10 @@ class ContextGovernor:
updated = self.drop_orphan_tool_results(updated) updated = self.drop_orphan_tool_results(updated)
updated = self.backfill_missing_tool_results(updated) updated = self.backfill_missing_tool_results(updated)
updated = self.apply_tool_result_budget(config, updated) updated = self.apply_tool_result_budget(config, updated)
updated = self.drop_stale_visual_tool_images(
updated,
start_index=config.inflight_start_index,
)
updated = self.compact_inflight_overflow(config, updated, compacted_tool_call_ids) updated = self.compact_inflight_overflow(config, updated, compacted_tool_call_ids)
updated = self.snip_history(config, updated) updated = self.snip_history(config, updated)
updated = self.drop_orphan_tool_results(updated) updated = self.drop_orphan_tool_results(updated)
@@ -326,6 +341,35 @@ class ContextGovernor:
updated[idx]["content"] = normalized updated[idx]["content"] = normalized
return updated return updated
@staticmethod
def drop_stale_visual_tool_images(
messages: list[dict[str, Any]],
*,
start_index: int,
) -> list[dict[str, Any]]:
"""Keep only the latest in-flight screenshot from each visual tool."""
seen: set[str] = set()
updated = messages
for idx in range(len(messages) - 1, start_index - 1, -1):
message = messages[idx]
name = str(message.get("name") or "")
content = message.get("content")
if message.get("role") != "tool" or name not in VISUAL_TOOLS:
continue
if not isinstance(content, list):
continue
content_blocks = cast(list[object], content)
blocks = [block for block in content_blocks if not _is_image_block(block)]
if len(blocks) == len(content_blocks):
continue
if name not in seen:
seen.add(name)
continue
if updated is messages:
updated = [dict(item) for item in messages]
updated[idx]["content"] = [dict(STALE_SCREENSHOT_PLACEHOLDER), *blocks]
return updated
def compact_inflight_overflow( def compact_inflight_overflow(
self, self,
config: ContextGovernanceConfig, config: ContextGovernanceConfig,
+9 -1
View File
@@ -4,6 +4,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
from typing import Any from typing import Any
from pydantic import Field from pydantic import Field
@@ -44,6 +45,7 @@ class BrowserToolConfig(Base):
allowed_domains: list[str] = Field(default_factory=list) allowed_domains: list[str] = Field(default_factory=list)
include_screenshot: bool = False include_screenshot: bool = False
max_elements: int = Field(default=200, ge=1, le=1000) max_elements: int = Field(default=200, ge=1, le=1000)
max_sessions: int = Field(default=8, ge=1, le=64)
def _format_elements(elements: list[dict[str, Any]]) -> str: def _format_elements(elements: list[dict[str, Any]]) -> str:
@@ -134,9 +136,11 @@ class BrowserTool(Tool):
from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserRuntime from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserRuntime
runtime = BrowserRuntime(headless=self.config.headless) runtime = BrowserRuntime(headless=self.config.headless)
self._runtime = runtime self._runtime = runtime
self._execution_lock = asyncio.Lock()
self._backends = SessionBackendPool( self._backends = SessionBackendPool(
self._make_backend, self._make_backend,
backend_impl, backend_impl,
max_backends=self.config.max_sessions,
finalizer=runtime.close if runtime is not None else None, finalizer=runtime.close if runtime is not None else None,
) )
@@ -225,12 +229,16 @@ class BrowserTool(Tool):
raise ValueError(f"unknown action '{action}'") raise ValueError(f"unknown action '{action}'")
async def execute(self, action: str | None = None, **kwargs: Any) -> Any: async def execute(self, action: str | None = None, **kwargs: Any) -> Any:
async with self._execution_lock:
return await self._execute(action, **kwargs)
async def _execute(self, action: str | None = None, **kwargs: Any) -> Any:
action = (action or "").strip() action = (action or "").strip()
if action not in _ACTIONS: if action not in _ACTIONS:
return f"Error: unknown action '{action}'. Valid actions: {', '.join(_ACTIONS)}" return f"Error: unknown action '{action}'. Valid actions: {', '.join(_ACTIONS)}"
try: try:
backend = self._backends.get() backend = await self._backends.get()
except ImportError as exc: except ImportError as exc:
return f"Error: {exc}" return f"Error: {exc}"
except Exception as exc: except Exception as exc:
+10 -2
View File
@@ -60,6 +60,7 @@ class ComputerUseToolConfig(Base):
allowed_domains: list[str] = Field(default_factory=list) allowed_domains: list[str] = Field(default_factory=list)
start_url: str = "about:blank" start_url: str = "about:blank"
headless: bool = True headless: bool = True
max_sessions: int = Field(default=8, ge=1, le=64)
def _fit_size(width: int, height: int, max_width: int, max_height: int) -> tuple[int, int]: def _fit_size(width: int, height: int, max_width: int, max_height: int) -> tuple[int, int]:
@@ -97,7 +98,8 @@ def _scale_point(
nullable=True, nullable=True,
), ),
text=StringSchema( text=StringSchema(
"Text to type (action=type), or a key/combo like 'ctrl+s' or 'Enter' (action=key).", "Text to type (action=type; desktop supports ASCII), or a key/combo like "
"'ctrl+s' or 'Enter' (action=key).",
nullable=True, nullable=True,
), ),
scroll_direction=StringSchema( scroll_direction=StringSchema(
@@ -160,9 +162,11 @@ class ComputerUseTool(Tool):
from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserRuntime from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserRuntime
runtime = BrowserRuntime(headless=self.config.headless) runtime = BrowserRuntime(headless=self.config.headless)
self._runtime = runtime self._runtime = runtime
self._execution_lock = asyncio.Lock()
self._backends = SessionBackendPool( self._backends = SessionBackendPool(
self._make_backend, self._make_backend,
backend_impl, backend_impl,
max_backends=1 if self.config.backend == "desktop" else self.config.max_sessions,
finalizer=runtime.close if runtime is not None else None, finalizer=runtime.close if runtime is not None else None,
) )
@@ -277,12 +281,16 @@ class ComputerUseTool(Tool):
raise ValueError(f"unknown action '{action}'") raise ValueError(f"unknown action '{action}'")
async def execute(self, action: str | None = None, **kwargs: Any) -> Any: async def execute(self, action: str | None = None, **kwargs: Any) -> Any:
async with self._execution_lock:
return await self._execute(action, **kwargs)
async def _execute(self, action: str | None = None, **kwargs: Any) -> Any:
action = (action or "").strip() action = (action or "").strip()
if action not in _ACTIONS: if action not in _ACTIONS:
return f"Error: unknown action '{action}'. Valid actions: {', '.join(_ACTIONS)}" return f"Error: unknown action '{action}'. Valid actions: {', '.join(_ACTIONS)}"
try: try:
backend = self._backends.get() backend = await self._backends.get()
real_w, real_h = await backend.dimensions() real_w, real_h = await backend.dimensions()
except ImportError as exc: except ImportError as exc:
return f"Error: {exc}" return f"Error: {exc}"
@@ -15,6 +15,7 @@ from __future__ import annotations
import asyncio import asyncio
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from collections import OrderedDict
from collections.abc import Awaitable, Callable from collections.abc import Awaitable, Callable
from typing import Any from typing import Any
@@ -78,37 +79,60 @@ class SessionBackendPool:
factory: Callable[[], Any], factory: Callable[[], Any],
injected: Any = None, injected: Any = None,
*, *,
max_backends: int = 8,
finalizer: Callable[[], Awaitable[None]] | None = None, finalizer: Callable[[], Awaitable[None]] | None = None,
) -> None: ) -> None:
if max_backends < 1:
raise ValueError("max_backends must be at least 1")
self._factory = factory self._factory = factory
self._injected = injected self._injected = injected
self._max_backends = max_backends
self._finalizer = finalizer self._finalizer = finalizer
self._backends: dict[str, Any] = {} self._backends: OrderedDict[str, Any] = OrderedDict()
self._lock = asyncio.Lock()
self._closed = False
def get(self) -> Any: async def get(self) -> Any:
if self._injected is not None: async with self._lock:
return self._injected if self._closed:
key = current_request_session_key() or "default" raise RuntimeError("computer-use backend pool is closed")
backend = self._backends.get(key) if self._injected is not None:
if backend is None: return self._injected
backend = self._backends[key] = self._factory() key = current_request_session_key() or "default"
return backend backend = self._backends.get(key)
if backend is not None:
self._backends.move_to_end(key)
return backend
if len(self._backends) >= self._max_backends:
_, stale = self._backends.popitem(last=False)
await stale.close()
backend = self._factory()
self._backends[key] = backend
return backend
async def close(self) -> None: async def close(self) -> None:
backends = [self._injected] if self._injected is not None else list(self._backends.values()) async with self._lock:
self._injected = None if self._closed:
self._backends.clear() return
self._closed = True
backends = (
[self._injected]
if self._injected is not None
else list(self._backends.values())
)
self._injected = None
self._backends.clear()
finalizer, self._finalizer = self._finalizer, None
results = await asyncio.gather( results = await asyncio.gather(
*(backend.close() for backend in backends if backend is not None), *(backend.close() for backend in backends if backend is not None),
return_exceptions=True, return_exceptions=True,
) )
errors = [result for result in results if isinstance(result, BaseException)] errors = [result for result in results if isinstance(result, BaseException)]
if self._finalizer is not None: if finalizer is not None:
try: try:
await self._finalizer() await finalizer()
except BaseException as exc: except BaseException as exc:
errors.append(exc) errors.append(exc)
self._finalizer = None
if len(errors) == 1: if len(errors) == 1:
raise errors[0] raise errors[0]
if errors: if errors:
@@ -111,6 +111,11 @@ class DesktopBackend(ComputerBackend):
await asyncio.to_thread(pg.hscroll, clicks if direction == "right" else -clicks) await asyncio.to_thread(pg.hscroll, clicks if direction == "right" else -clicks)
async def type_text(self, text: str) -> None: async def type_text(self, text: str) -> None:
if not text.isascii():
raise ValueError(
"desktop text input supports ASCII key events only; "
"use the browser backend for Unicode text"
)
pg = self._ensure() pg = self._ensure()
await asyncio.to_thread(pg.typewrite, text, 0.01) await asyncio.to_thread(pg.typewrite, text, 0.01)
+12 -5
View File
@@ -356,6 +356,7 @@ _TOOL_RESULT_PREVIEW_CHARS = 1200
_TOOL_RESULTS_DIR = ".nanobot/tool-results" _TOOL_RESULTS_DIR = ".nanobot/tool-results"
_TOOL_RESULT_RETENTION_SECS = 7 * 24 * 60 * 60 _TOOL_RESULT_RETENTION_SECS = 7 * 24 * 60 * 60
_TOOL_RESULT_MAX_BUCKETS = 32 _TOOL_RESULT_MAX_BUCKETS = 32
_IMAGE_TOKEN_ESTIMATE = 2048
_TRUNCATED_SUFFIX = "\n... (truncated)" _TRUNCATED_SUFFIX = "\n... (truncated)"
@@ -676,6 +677,7 @@ def _estimate_prompt_tokens_with_source(
reasoning_content, tool_call_id, name, plus per-message framing overhead. reasoning_content, tool_call_id, name, plus per-message framing overhead.
""" """
parts: list[str] = [] parts: list[str] = []
image_tokens = 0
for msg in messages: for msg in messages:
content = msg.get("content") content = msg.get("content")
if isinstance(content, str): if isinstance(content, str):
@@ -687,6 +689,8 @@ def _estimate_prompt_tokens_with_source(
text = part.get("text", "") text = part.get("text", "")
if isinstance(text, str) and text: if isinstance(text, str) and text:
parts.append(text) parts.append(text)
elif part is not None and part.get("type") in {"image_url", "input_image"}:
image_tokens += _IMAGE_TOKEN_ESTIMATE
tc = msg.get("tool_calls") tc = msg.get("tool_calls")
if tc: if tc:
@@ -709,7 +713,7 @@ def _estimate_prompt_tokens_with_source(
_estimate_tools_tokens(enc, tools, leading_separator=bool(parts)) if tools else 0 _estimate_tools_tokens(enc, tools, leading_separator=bool(parts)) if tools else 0
) )
message_tokens = len(enc.encode(message_payload)) if message_payload else 0 message_tokens = len(enc.encode(message_payload)) if message_payload else 0
return message_tokens + tool_tokens + per_message_overhead, "tiktoken" return message_tokens + image_tokens + tool_tokens + per_message_overhead, "tiktoken"
except Exception: except Exception:
tool_payload = ( tool_payload = (
("\n" if message_payload else "") + json.dumps(tools, ensure_ascii=False) ("\n" if message_payload else "") + json.dumps(tools, ensure_ascii=False)
@@ -718,7 +722,7 @@ def _estimate_prompt_tokens_with_source(
) )
payload = message_payload + tool_payload payload = message_payload + tool_payload
estimated = len(payload.encode("utf-8")) estimated = len(payload.encode("utf-8"))
return estimated + per_message_overhead, "heuristic" return estimated + image_tokens + per_message_overhead, "heuristic"
def estimate_prompt_tokens( def estimate_prompt_tokens(
@@ -734,6 +738,7 @@ def estimate_message_tokens(message: dict[str, Any]) -> int:
"""Estimate prompt tokens contributed by one persisted message.""" """Estimate prompt tokens contributed by one persisted message."""
content = message.get("content") content = message.get("content")
parts: list[str] = [] parts: list[str] = []
image_tokens = 0
if isinstance(content, str): if isinstance(content, str):
parts.append(content) parts.append(content)
elif isinstance(content, list): elif isinstance(content, list):
@@ -743,6 +748,8 @@ def estimate_message_tokens(message: dict[str, Any]) -> int:
text = part.get("text", "") text = part.get("text", "")
if isinstance(text, str) and text: if isinstance(text, str) and text:
parts.append(text) parts.append(text)
elif part is not None and part.get("type") in {"image_url", "input_image"}:
image_tokens += _IMAGE_TOKEN_ESTIMATE
else: else:
parts.append(json.dumps(raw_part, ensure_ascii=False)) parts.append(json.dumps(raw_part, ensure_ascii=False))
elif content is not None: elif content is not None:
@@ -760,13 +767,13 @@ def estimate_message_tokens(message: dict[str, Any]) -> int:
parts.append(rc) parts.append(rc)
payload = "\n".join(parts) payload = "\n".join(parts)
if not payload: if not payload and not image_tokens:
return 4 return 4
try: try:
enc = _get_token_encoding() enc = _get_token_encoding()
return max(4, len(enc.encode(payload)) + 4) return max(4, len(enc.encode(payload)) + image_tokens + 4)
except Exception: except Exception:
return max(4, len(payload.encode("utf-8")) + 4) return max(4, len(payload.encode("utf-8")) + image_tokens + 4)
def estimate_prompt_tokens_chain( def estimate_prompt_tokens_chain(
+25
View File
@@ -1,6 +1,13 @@
from nanobot.agent.context_governance import ContextGovernor from nanobot.agent.context_governance import ContextGovernor
def _image_result(label: str) -> list[dict]:
return [
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{label}"}},
{"type": "text", "text": label},
]
def _assistant_tool_call(call_id: str) -> dict: def _assistant_tool_call(call_id: str) -> dict:
return { return {
"role": "assistant", "role": "assistant",
@@ -37,3 +44,21 @@ def test_drop_orphan_tool_results_drops_duplicate_tool_result() -> None:
tool_results = [m for m in result if m.get("role") == "tool"] tool_results = [m for m in result if m.get("role") == "tool"]
assert len(tool_results) == 1 assert len(tool_results) == 1
assert tool_results[0]["content"] == "first" assert tool_results[0]["content"] == "first"
def test_drop_stale_visual_tool_images_keeps_latest_per_tool() -> None:
messages = [
{"role": "tool", "name": "computer_use", "content": _image_result("history")},
{"role": "tool", "name": "computer_use", "content": _image_result("old")},
{"role": "tool", "name": "browser", "content": _image_result("browser")},
{"role": "tool", "name": "computer_use", "content": _image_result("latest")},
]
result = ContextGovernor.drop_stale_visual_tool_images(messages, start_index=1)
assert result is not messages
assert result[0]["content"] == messages[0]["content"]
assert [block["type"] for block in result[1]["content"]] == ["text", "text"]
assert result[2]["content"] == messages[2]["content"]
assert result[3]["content"] == messages[3]["content"]
assert messages[1]["content"][0]["type"] == "image_url"
+29
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import io import io
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock from unittest.mock import AsyncMock, MagicMock
@@ -11,6 +12,7 @@ import pytest
from nanobot.agent.tools.browser_tool import BrowserTool, BrowserToolConfig from nanobot.agent.tools.browser_tool import BrowserTool, BrowserToolConfig
from nanobot.agent.tools.computer_use_backends import browser_playwright from nanobot.agent.tools.computer_use_backends import browser_playwright
from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserBackend from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserBackend
from nanobot.agent.tools.context import RequestContext, request_context
class _FakeDomBackend: class _FakeDomBackend:
@@ -86,6 +88,7 @@ class TestConfigAndMetadata:
assert cfg.headless is True assert cfg.headless is True
assert cfg.include_screenshot is False assert cfg.include_screenshot is False
assert cfg.max_elements == 200 assert cfg.max_elements == 200
assert cfg.max_sessions == 8
def test_enabled_reads_config(self): def test_enabled_reads_config(self):
ctx = MagicMock() ctx = MagicMock()
@@ -171,6 +174,32 @@ class TestDispatch:
assert imgs and texts assert imgs and texts
assert "Clicked element [1]" in texts[-1]["text"] assert "Clicked element [1]" in texts[-1]["text"]
@pytest.mark.asyncio
async def test_calls_are_serialized_across_sessions(self):
class SlowBackend(_FakeDomBackend):
active = 0
max_active = 0
async def dom_snapshot(self, max_elements=200):
self.active += 1
self.max_active = max(self.max_active, self.active)
await asyncio.sleep(0.01)
self.active -= 1
return await super().dom_snapshot(max_elements)
backend = SlowBackend()
tool = BrowserTool(backend_impl=backend)
async def snapshot(session: str):
with request_context(
RequestContext(channel="test", chat_id=session, session_key=session)
):
return await tool.execute(action="snapshot")
await asyncio.gather(snapshot("a"), snapshot("b"))
assert backend.max_active == 1
class TestErrorsAndPolicy: class TestErrorsAndPolicy:
@pytest.mark.parametrize( @pytest.mark.parametrize(
+75 -5
View File
@@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import asyncio
import io import io
import sys import sys
from types import SimpleNamespace from types import SimpleNamespace
@@ -81,19 +82,23 @@ class TestConfigAndMetadata:
assert cfg.enable is False assert cfg.enable is False
assert cfg.backend == "desktop" assert cfg.backend == "desktop"
assert (cfg.target_width, cfg.target_height) == (1280, 800) assert (cfg.target_width, cfg.target_height) == (1280, 800)
assert cfg.max_sessions == 8
assert "require_approval" not in type(cfg).model_fields assert "require_approval" not in type(cfg).model_fields
def test_tools_config_accepts_camel_case(self): def test_tools_config_accepts_camel_case(self):
cfg = ToolsConfig.model_validate({ cfg = ToolsConfig.model_validate({
"browser": {"enable": True}, "browser": {"enable": True, "maxSessions": 4},
"computerUse": {"enable": True, "backend": "browser"}, "computerUse": {"enable": True, "backend": "browser", "maxSessions": 6},
}) })
assert cfg.browser.enable is True assert cfg.browser.enable is True
assert cfg.browser.max_sessions == 4
assert cfg.computer_use.enable is True assert cfg.computer_use.enable is True
assert cfg.computer_use.backend == "browser" assert cfg.computer_use.backend == "browser"
assert cfg.computer_use.max_sessions == 6
dumped = cfg.model_dump(by_alias=True) dumped = cfg.model_dump(by_alias=True)
assert "computerUse" in dumped assert "computerUse" in dumped
assert dumped["computerUse"]["maxSessions"] == 6
def test_enabled_reads_config(self): def test_enabled_reads_config(self):
ctx = MagicMock() ctx = MagicMock()
@@ -218,10 +223,10 @@ async def test_backend_pool_isolates_sessions_and_closes_all():
pool = SessionBackendPool(factory, finalizer=finalize) pool = SessionBackendPool(factory, finalizer=finalize)
with request_context(RequestContext(channel="test", chat_id="a", session_key="test:a")): with request_context(RequestContext(channel="test", chat_id="a", session_key="test:a")):
first = pool.get() first = await pool.get()
assert pool.get() is first assert await pool.get() is first
with request_context(RequestContext(channel="test", chat_id="b", session_key="test:b")): with request_context(RequestContext(channel="test", chat_id="b", session_key="test:b")):
second = pool.get() second = await pool.get()
assert first is not second assert first is not second
await pool.close() await pool.close()
@@ -233,6 +238,59 @@ async def test_backend_pool_isolates_sessions_and_closes_all():
assert finalized == [True] assert finalized == [True]
@pytest.mark.asyncio
async def test_backend_pool_evicts_least_recently_used_session():
created: list[_FakeBackend] = []
def factory():
backend = _FakeBackend()
created.append(backend)
return backend
pool = SessionBackendPool(factory, max_backends=2)
contexts = [
RequestContext(channel="test", chat_id=key, session_key=f"test:{key}")
for key in ("a", "b", "c")
]
with request_context(contexts[0]):
first = await pool.get()
with request_context(contexts[1]):
second = await pool.get()
with request_context(contexts[0]):
assert await pool.get() is first
with request_context(contexts[2]):
await pool.get()
assert first.closed is False
assert second.closed is True
await pool.close()
@pytest.mark.asyncio
async def test_desktop_tool_serializes_calls_across_sessions():
class SlowBackend(_FakeBackend):
active = 0
max_active = 0
async def dimensions(self):
self.active += 1
self.max_active = max(self.max_active, self.active)
await asyncio.sleep(0.01)
self.active -= 1
return await super().dimensions()
backend = SlowBackend(width=1280, height=800)
tool = ComputerUseTool(backend_impl=backend)
async def screenshot(session: str):
with request_context(RequestContext(channel="test", chat_id=session, session_key=session)):
return await tool.execute(action="screenshot")
await asyncio.gather(screenshot("a"), screenshot("b"))
assert backend.max_active == 1
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_desktop_backend_uses_safe_pyautogui_calls(): async def test_desktop_backend_uses_safe_pyautogui_calls():
pg = MagicMock() pg = MagicMock()
@@ -251,6 +309,18 @@ async def test_desktop_backend_uses_safe_pyautogui_calls():
pg.scroll.assert_called_once_with(-3) pg.scroll.assert_called_once_with(-3)
@pytest.mark.asyncio
async def test_desktop_backend_rejects_unicode_instead_of_typing_incorrect_keys():
pg = MagicMock()
backend = DesktopBackend()
backend._pg = pg
with pytest.raises(ValueError, match="ASCII"):
await backend.type_text("你好")
pg.typewrite.assert_not_called()
def test_desktop_backend_preserves_pyautogui_failsafe(monkeypatch): def test_desktop_backend_preserves_pyautogui_failsafe(monkeypatch):
pg = SimpleNamespace(FAILSAFE=True) pg = SimpleNamespace(FAILSAFE=True)
monkeypatch.setitem(sys.modules, "pyautogui", pg) monkeypatch.setitem(sys.modules, "pyautogui", pg)
+30
View File
@@ -29,6 +29,36 @@ def test_estimate_prompt_tokens_chain_falls_back_without_provider_counter() -> N
assert source == "tiktoken" assert source == "tiktoken"
def test_image_blocks_have_bounded_token_cost() -> None:
text = [{"role": "tool", "content": [{"type": "text", "text": "screen"}]}]
small_image = [{
"role": "tool",
"content": [
{"type": "image_url", "image_url": {"url": "data:image/png;base64,A"}},
{"type": "text", "text": "screen"},
],
}]
large_image = [{
"role": "tool",
"content": [
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64," + "A" * 100_000},
},
{"type": "text", "text": "screen"},
],
}]
text_tokens = estimate_prompt_tokens(text)
small_tokens = estimate_prompt_tokens(small_image)
large_tokens = estimate_prompt_tokens(large_image)
assert small_tokens >= text_tokens + 2_000
assert large_tokens == small_tokens
assert estimate_message_tokens(large_image[0]) >= text_tokens + 2_000
assert estimate_message_tokens({"role": "user", "content": small_image[0]["content"][:1]}) > 2_000
def test_estimate_prompt_tokens_chain_falls_back_when_provider_counter_fails() -> None: def test_estimate_prompt_tokens_chain_falls_back_when_provider_counter_fails() -> None:
tokens, source = estimate_prompt_tokens_chain( tokens, source = estimate_prompt_tokens_chain(
_BrokenCounterProvider(), _BrokenCounterProvider(),