From a7391857409aabc2c5bf8d8bfc4902f6fe7c19a7 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Sun, 9 Aug 2026 01:24:25 +0900 Subject: [PATCH] fix(tools): bound computer use runtime state --- docs/configuration.md | 8 +- nanobot/agent/context_governance.py | 44 ++++++++++ nanobot/agent/tools/browser_tool.py | 10 ++- nanobot/agent/tools/computer_use.py | 12 ++- .../agent/tools/computer_use_backends/base.py | 54 +++++++++---- .../desktop_pyautogui.py | 5 ++ nanobot/utils/helpers.py | 17 ++-- tests/agent/test_context_governance.py | 25 ++++++ tests/tools/test_browser_tool.py | 29 +++++++ tests/tools/test_computer_use_tool.py | 80 +++++++++++++++++-- tests/utils/test_token_estimation.py | 30 +++++++ 11 files changed, 284 insertions(+), 30 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index d97255b5c..5a10e6c5f 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -1705,14 +1705,17 @@ controls an isolated Playwright page. | `tools.browser.enable` | `false` | Register the DOM-based `browser` tool | | `tools.browser.allowedDomains` | `[]` | Optional top-level navigation allowlist; entries include subdomains | | `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.backend` | `"desktop"` | `"desktop"` or `"browser"` | | `tools.computerUse.allowedDomains` | `[]` | Navigation allowlist for the browser backend | | `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 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] > 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] > 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 -> 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 diff --git a/nanobot/agent/context_governance.py b/nanobot/agent/context_governance.py index 98b1291dc..415767de7 100644 --- a/nanobot/agent/context_governance.py +++ b/nanobot/agent/context_governance.py @@ -33,6 +33,11 @@ COMPACTABLE_TOOLS = frozenset({ "read_file", "exec", "grep", "find_files", "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. TOOL_RESULT_OFFLOAD_EXEMPT_TOOLS = frozenset({"read_file"}) 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: """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.backfill_missing_tool_results(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.snip_history(config, updated) updated = self.drop_orphan_tool_results(updated) @@ -326,6 +341,35 @@ class ContextGovernor: updated[idx]["content"] = normalized 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( self, config: ContextGovernanceConfig, diff --git a/nanobot/agent/tools/browser_tool.py b/nanobot/agent/tools/browser_tool.py index fa182d226..cf2f18b59 100644 --- a/nanobot/agent/tools/browser_tool.py +++ b/nanobot/agent/tools/browser_tool.py @@ -4,6 +4,7 @@ from __future__ import annotations +import asyncio from typing import Any from pydantic import Field @@ -44,6 +45,7 @@ class BrowserToolConfig(Base): allowed_domains: list[str] = Field(default_factory=list) include_screenshot: bool = False 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: @@ -134,9 +136,11 @@ class BrowserTool(Tool): from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserRuntime runtime = BrowserRuntime(headless=self.config.headless) self._runtime = runtime + self._execution_lock = asyncio.Lock() self._backends = SessionBackendPool( self._make_backend, backend_impl, + max_backends=self.config.max_sessions, finalizer=runtime.close if runtime is not None else None, ) @@ -225,12 +229,16 @@ class BrowserTool(Tool): raise ValueError(f"unknown action '{action}'") 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() if action not in _ACTIONS: return f"Error: unknown action '{action}'. Valid actions: {', '.join(_ACTIONS)}" try: - backend = self._backends.get() + backend = await self._backends.get() except ImportError as exc: return f"Error: {exc}" except Exception as exc: diff --git a/nanobot/agent/tools/computer_use.py b/nanobot/agent/tools/computer_use.py index a0daf3a62..a60d5ffa2 100644 --- a/nanobot/agent/tools/computer_use.py +++ b/nanobot/agent/tools/computer_use.py @@ -60,6 +60,7 @@ class ComputerUseToolConfig(Base): allowed_domains: list[str] = Field(default_factory=list) start_url: str = "about:blank" 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]: @@ -97,7 +98,8 @@ def _scale_point( nullable=True, ), 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, ), scroll_direction=StringSchema( @@ -160,9 +162,11 @@ class ComputerUseTool(Tool): from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserRuntime runtime = BrowserRuntime(headless=self.config.headless) self._runtime = runtime + self._execution_lock = asyncio.Lock() self._backends = SessionBackendPool( self._make_backend, 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, ) @@ -277,12 +281,16 @@ class ComputerUseTool(Tool): raise ValueError(f"unknown action '{action}'") 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() if action not in _ACTIONS: return f"Error: unknown action '{action}'. Valid actions: {', '.join(_ACTIONS)}" try: - backend = self._backends.get() + backend = await self._backends.get() real_w, real_h = await backend.dimensions() except ImportError as exc: return f"Error: {exc}" diff --git a/nanobot/agent/tools/computer_use_backends/base.py b/nanobot/agent/tools/computer_use_backends/base.py index 12209196b..6ed75c896 100644 --- a/nanobot/agent/tools/computer_use_backends/base.py +++ b/nanobot/agent/tools/computer_use_backends/base.py @@ -15,6 +15,7 @@ from __future__ import annotations import asyncio from abc import ABC, abstractmethod +from collections import OrderedDict from collections.abc import Awaitable, Callable from typing import Any @@ -78,37 +79,60 @@ class SessionBackendPool: factory: Callable[[], Any], injected: Any = None, *, + max_backends: int = 8, finalizer: Callable[[], Awaitable[None]] | None = None, ) -> None: + if max_backends < 1: + raise ValueError("max_backends must be at least 1") self._factory = factory self._injected = injected + self._max_backends = max_backends 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: - if self._injected is not None: - return self._injected - key = current_request_session_key() or "default" - backend = self._backends.get(key) - if backend is None: - backend = self._backends[key] = self._factory() - return backend + async def get(self) -> Any: + async with self._lock: + if self._closed: + raise RuntimeError("computer-use backend pool is closed") + if self._injected is not None: + return self._injected + key = current_request_session_key() or "default" + 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: - backends = [self._injected] if self._injected is not None else list(self._backends.values()) - self._injected = None - self._backends.clear() + async with self._lock: + if self._closed: + 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( *(backend.close() for backend in backends if backend is not None), return_exceptions=True, ) errors = [result for result in results if isinstance(result, BaseException)] - if self._finalizer is not None: + if finalizer is not None: try: - await self._finalizer() + await finalizer() except BaseException as exc: errors.append(exc) - self._finalizer = None if len(errors) == 1: raise errors[0] if errors: diff --git a/nanobot/agent/tools/computer_use_backends/desktop_pyautogui.py b/nanobot/agent/tools/computer_use_backends/desktop_pyautogui.py index 15d265355..c349de495 100644 --- a/nanobot/agent/tools/computer_use_backends/desktop_pyautogui.py +++ b/nanobot/agent/tools/computer_use_backends/desktop_pyautogui.py @@ -111,6 +111,11 @@ class DesktopBackend(ComputerBackend): await asyncio.to_thread(pg.hscroll, clicks if direction == "right" else -clicks) 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() await asyncio.to_thread(pg.typewrite, text, 0.01) diff --git a/nanobot/utils/helpers.py b/nanobot/utils/helpers.py index 1f1720dd3..9d86c0a5b 100644 --- a/nanobot/utils/helpers.py +++ b/nanobot/utils/helpers.py @@ -356,6 +356,7 @@ _TOOL_RESULT_PREVIEW_CHARS = 1200 _TOOL_RESULTS_DIR = ".nanobot/tool-results" _TOOL_RESULT_RETENTION_SECS = 7 * 24 * 60 * 60 _TOOL_RESULT_MAX_BUCKETS = 32 +_IMAGE_TOKEN_ESTIMATE = 2048 _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. """ parts: list[str] = [] + image_tokens = 0 for msg in messages: content = msg.get("content") if isinstance(content, str): @@ -687,6 +689,8 @@ def _estimate_prompt_tokens_with_source( text = part.get("text", "") if isinstance(text, str) and 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") 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 ) 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: tool_payload = ( ("\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 estimated = len(payload.encode("utf-8")) - return estimated + per_message_overhead, "heuristic" + return estimated + image_tokens + per_message_overhead, "heuristic" 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.""" content = message.get("content") parts: list[str] = [] + image_tokens = 0 if isinstance(content, str): parts.append(content) elif isinstance(content, list): @@ -743,6 +748,8 @@ def estimate_message_tokens(message: dict[str, Any]) -> int: text = part.get("text", "") if isinstance(text, str) and text: parts.append(text) + elif part is not None and part.get("type") in {"image_url", "input_image"}: + image_tokens += _IMAGE_TOKEN_ESTIMATE else: parts.append(json.dumps(raw_part, ensure_ascii=False)) elif content is not None: @@ -760,13 +767,13 @@ def estimate_message_tokens(message: dict[str, Any]) -> int: parts.append(rc) payload = "\n".join(parts) - if not payload: + if not payload and not image_tokens: return 4 try: enc = _get_token_encoding() - return max(4, len(enc.encode(payload)) + 4) + return max(4, len(enc.encode(payload)) + image_tokens + 4) 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( diff --git a/tests/agent/test_context_governance.py b/tests/agent/test_context_governance.py index 14daef174..471e31bd0 100644 --- a/tests/agent/test_context_governance.py +++ b/tests/agent/test_context_governance.py @@ -1,6 +1,13 @@ 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: return { "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"] assert len(tool_results) == 1 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" diff --git a/tests/tools/test_browser_tool.py b/tests/tools/test_browser_tool.py index f551f6f58..a5bc92f10 100644 --- a/tests/tools/test_browser_tool.py +++ b/tests/tools/test_browser_tool.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import io from types import SimpleNamespace 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.computer_use_backends import browser_playwright from nanobot.agent.tools.computer_use_backends.browser_playwright import BrowserBackend +from nanobot.agent.tools.context import RequestContext, request_context class _FakeDomBackend: @@ -86,6 +88,7 @@ class TestConfigAndMetadata: assert cfg.headless is True assert cfg.include_screenshot is False assert cfg.max_elements == 200 + assert cfg.max_sessions == 8 def test_enabled_reads_config(self): ctx = MagicMock() @@ -171,6 +174,32 @@ class TestDispatch: assert imgs and texts 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: @pytest.mark.parametrize( diff --git a/tests/tools/test_computer_use_tool.py b/tests/tools/test_computer_use_tool.py index 9f53f6028..fe2c49de4 100644 --- a/tests/tools/test_computer_use_tool.py +++ b/tests/tools/test_computer_use_tool.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import io import sys from types import SimpleNamespace @@ -81,19 +82,23 @@ class TestConfigAndMetadata: assert cfg.enable is False assert cfg.backend == "desktop" assert (cfg.target_width, cfg.target_height) == (1280, 800) + assert cfg.max_sessions == 8 assert "require_approval" not in type(cfg).model_fields def test_tools_config_accepts_camel_case(self): cfg = ToolsConfig.model_validate({ - "browser": {"enable": True}, - "computerUse": {"enable": True, "backend": "browser"}, + "browser": {"enable": True, "maxSessions": 4}, + "computerUse": {"enable": True, "backend": "browser", "maxSessions": 6}, }) assert cfg.browser.enable is True + assert cfg.browser.max_sessions == 4 assert cfg.computer_use.enable is True assert cfg.computer_use.backend == "browser" + assert cfg.computer_use.max_sessions == 6 dumped = cfg.model_dump(by_alias=True) assert "computerUse" in dumped + assert dumped["computerUse"]["maxSessions"] == 6 def test_enabled_reads_config(self): ctx = MagicMock() @@ -218,10 +223,10 @@ async def test_backend_pool_isolates_sessions_and_closes_all(): pool = SessionBackendPool(factory, finalizer=finalize) with request_context(RequestContext(channel="test", chat_id="a", session_key="test:a")): - first = pool.get() - assert pool.get() is first + first = await pool.get() + assert await pool.get() is first 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 await pool.close() @@ -233,6 +238,59 @@ async def test_backend_pool_isolates_sessions_and_closes_all(): 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 async def test_desktop_backend_uses_safe_pyautogui_calls(): pg = MagicMock() @@ -251,6 +309,18 @@ async def test_desktop_backend_uses_safe_pyautogui_calls(): 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): pg = SimpleNamespace(FAILSAFE=True) monkeypatch.setitem(sys.modules, "pyautogui", pg) diff --git a/tests/utils/test_token_estimation.py b/tests/utils/test_token_estimation.py index f3949ad49..b4a8f4e7d 100644 --- a/tests/utils/test_token_estimation.py +++ b/tests/utils/test_token_estimation.py @@ -29,6 +29,36 @@ def test_estimate_prompt_tokens_chain_falls_back_without_provider_counter() -> N 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: tokens, source = estimate_prompt_tokens_chain( _BrokenCounterProvider(),