diff --git a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py index d5458586d..4f5f0f70a 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/agent_env_environment.py @@ -309,7 +309,8 @@ def __init__(self, environment_card: EnvironmentCard, handler: Any) -> None: self._app: Any = None ext_handlers = _served_extensions(handler, _discover_extensions(handler)) environment_card = _merge_extensions(environment_card, ext_handlers + _declared_extensions(handler, ext_handlers)) - tool_handlers = _served_tools(handler, _with_forwarded(_discover_tools(handler, environment_card.name), handler)) + tool_handlers = _named_tools(handler, _served_tools( + handler, _with_forwarded(_discover_tools(handler, environment_card.name), handler))) if tool_handlers and not self.supports_tools: raise ValueError(f"@tool methods require a FastMCP-backed application; {type(self).__name__} cannot serve MCP tools") self._resources = _discover_resources(handler) @@ -1107,11 +1108,15 @@ def read() -> Any: def _instructions(handler: Any) -> str | None: - """What the handler's @server_instructions method returns; None without one.""" + """What the handler's @server_instructions method returns, as its ``served_instructions`` serves it; None without + one.""" marked = [method for method, _marker in _marked_members(handler, _INSTRUCTIONS_ATTR)] if len(marked) > 1: raise ValueError(f"an environment has one @server_instructions method, not {sorted(m.__name__ for m in marked)}") - return marked[0]() if marked else None + if not marked: + return None + hook = getattr(handler, "served_instructions", None) + return hook(marked[0]()) if hook is not None else marked[0]() def _discover_methods(handler: Any) -> dict: @@ -1320,6 +1325,20 @@ def _served_tools(handler: Any, tool_handlers: list) -> list: return [(descriptor, by_name[descriptor.name]) for descriptor in chosen] +def _named_tools(handler: Any, tool_handlers: list) -> list: + """The tools as the handler's ``serve_tools_as`` names and describes them, in their order; as they are without one.""" + hook = getattr(handler, "serve_tools_as", None) + if hook is None: + return tool_handlers + named = list(hook([descriptor for descriptor, _fn in tool_handlers])) + if len(named) != len(tool_handlers): + raise ValueError(f"serve_tools_as returned {len(named)} tools for {len(tool_handlers)}") + names = [descriptor.name for descriptor in named] + if len(set(names)) != len(names): + raise ValueError(f"serve_tools_as named two tools alike: {sorted({n for n in names if names.count(n) > 1})}") + return [(descriptor, fn) for descriptor, (_old, fn) in zip(named, tool_handlers)] + + def _merge_tools(card: EnvironmentCard, tool_handlers: list) -> EnvironmentCard: """Advertise decorator-declared tools under card.capabilities; card-declared names win.""" if not tool_handlers: diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/__init__.py index 12d3350a2..678080857 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/__init__.py @@ -82,7 +82,7 @@ from .keys import PRINTABLE_KEYS from .manifest import ManifestSegment, RecordingManifest from .platform import Platform -from .settings import SETTINGS_ENV, ComputerSettings +from .settings import SETTINGS_ENV, TOOL_PREFIX_ENV, ComputerSettings, tool_prefix_problem if TYPE_CHECKING: from .card import ( @@ -127,7 +127,7 @@ def __getattr__(name: str) -> Any: "Launch", "Ready", "Packages", "AptSource", "Download", "RunScript", "AppDeclaration", "AppOs", "APP_SETTINGS_ENV", "part_bytes", - "Display", "Image", "Effect", "ExecResult", "ReadyCheck", "Platform", "ComputerSettings", "SETTINGS_ENV", + "Display", "Image", "Effect", "ExecResult", "ReadyCheck", "Platform", "ComputerSettings", "SETTINGS_ENV", "TOOL_PREFIX_ENV", "tool_prefix_problem", "Tab", "PageState", "PageDialog", "FileChooser", "PRINTABLE_KEYS", "display", "set_display", "capture", "click", "drag", "scroll", "type_text", "press", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/driver.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/driver.py index 6dfb0a756..bd5dc6856 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/driver.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/driver.py @@ -5,6 +5,7 @@ import asyncio import logging +import os from pathlib import Path from typing import Any, Callable, ClassVar, Literal, Optional @@ -13,9 +14,9 @@ from starlette.responses import Response from starlette.websockets import WebSocket -from ..agent_env_environment import (_CARD_CONFIG_ATTR, AgentEnvEnvironment, AgentEnvFastMCPApplication, - AgentEnvError, _error_response, add_websocket_route, control_check, - environment_card) +from ..agent_env_environment import (_CARD_CONFIG_ATTR, _TOOL_ATTR, AgentEnvEnvironment, AgentEnvFastMCPApplication, + AgentEnvError, _error_response, _marked_members, add_websocket_route, + control_check, environment_card) from ..tool_context import ToolContext from ..types import EnvironmentCard, EnvironmentExtension, EnvironmentTool from . import _linux, apps, card, svc @@ -27,8 +28,9 @@ from .platform import Platform, detect from .recording import Recording from .runtime import ComputerRuntime -from .settings import ComputerSettings, load_settings -from .tools import Tools +from .prefix import ToolPrefix +from .settings import TOOL_PREFIX_ENV, ComputerSettings, load_settings, tool_prefix_problem +from .tools import COMPUTER_TOOLS, Tools #: The caller every call comes from until agents carry an identity of their own. DEFAULT_IDENTITY = "default" @@ -99,9 +101,15 @@ def __init_subclass__(cls, **kwargs: Any) -> None: check_class(cls) check_features(cls) - def __init__(self, settings: ComputerSettings | None = None, platform: Platform | None = None) -> None: + def __init__(self, settings: ComputerSettings | None = None, platform: Platform | None = None, + tool_prefix: str | None = None) -> None: self.settings = settings if settings is not None else load_settings(model=type(self).Settings) self.platform = platform if platform is not None else detect() + prefix = os.environ.get(TOOL_PREFIX_ENV, "") if tool_prefix is None else tool_prefix + if problem := tool_prefix_problem(prefix): + raise ValueError(problem) + #: The start of every name it serves a tool under: the deploy's ``AGENTENV_TOOL_PREFIX``, none by default. + self._computer_prefix = ToolPrefix(prefix) self._computer = ComputerRuntime(self) #: The routes each app's card opens to agents, by app, read once it answers. self._computer_agent_routes: dict[str, list[str]] = {} @@ -113,7 +121,7 @@ def declaration(cls, platform: Platform) -> EnvironmentCard: running its ``__init__``: what a ``computer_driver`` env records for each platform it supports. The card hooks may read the class and the base class's state, not what a subclass's ``__init__`` sets.""" driver = cls.__new__(cls) - ComputerDriverEnvironment.__init__(driver, settings=cls.Settings(), platform=platform) + ComputerDriverEnvironment.__init__(driver, settings=cls.Settings(), platform=platform, tool_prefix="") config = dict(getattr(cls, _CARD_CONFIG_ATTR, None) or {}) card = EnvironmentCard(**{**config, "name": config.get("name") or cls.__name__}) return AgentEnvFastMCPApplication(card, driver).environment_card @@ -203,6 +211,17 @@ async def _computer_relay(self, websocket: WebSocket) -> None: def served_tools(self, tools: list[EnvironmentTool]) -> list[EnvironmentTool]: return card.served_tools(self, tools) + def serve_tools_as(self, tools: list[EnvironmentTool]) -> list[EnvironmentTool]: + return self._computer_prefix.tools(tools) + + def served_instructions(self, text: str) -> str: + return self._computer_prefix.instructions(text, self._computer_other_tools()) + + def _computer_other_tools(self) -> list[str]: + """Its driver's own and forwarded tools, whose text stays as written, by their own names.""" + return [marker.name for _member, marker in _marked_members(self, _TOOL_ATTR) if marker.name not in COMPUTER_TOOLS] + [ + forwarded.tool.name for forwarded in self.forwarded_tools()] + def served_extensions(self, extensions: list[EnvironmentExtension]) -> list[EnvironmentExtension]: return card.served_extensions(self, extensions) @@ -213,7 +232,8 @@ def extension_params(self, uri: str) -> dict: return card.extension_params(self, uri) def declared_extensions(self) -> list[EnvironmentExtension]: - return card.declared_extensions(self) + return [extension.model_copy(update={"description": self._computer_prefix.text(extension.description)}) + for extension in card.declared_extensions(self)] def _computer_require_init(self) -> None: if self._computer is None: diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/extensions.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/extensions.py index a569d094f..2732f7190 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/extensions.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/extensions.py @@ -79,8 +79,8 @@ async def _computer_spaces_set(self, default_space: str, default_norm_scale: Opt self.settings.default_norm_scale) self.settings = self.settings.model_copy(update={"default_space": space, "default_norm_scale": scale}) for tool in self.mounted_card().capabilities.tools or []: - if tool.name in card.POINT_TOOLS: - await self.update_tool(tool.name, description=card.point_description(self, tool.name)) + if (own := self._computer_prefix.own(tool.name)) in card.POINT_TOOLS: + await self.update_tool(tool.name, description=self._computer_prefix.text(card.point_description(self, own))) return await self._computer_spaces_get() @extension(COMPUTER_KEYBOARD, op="get", method="GET", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/prefix.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/prefix.py new file mode 100644 index 000000000..efaba1d73 --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/prefix.py @@ -0,0 +1,51 @@ +"""A computer's tool prefix: the names it serves its tools under, so several computers serve theirs side by side on +one gateway. The SDK's ``computer_*`` names in its text are rewritten to them; its driver's own and forwarded tools +keep their text as written, and a sentence tells an agent how those are named.""" +from __future__ import annotations + +import re +from collections.abc import Iterable + +from ..types import EnvironmentTool +from .tools import COMPUTER_TOOLS + +#: The SDK's tool names, and a ``computer_*`` family of them, each as a whole word. +_SDK_NAMES = re.compile(r"(? None: + self.prefix = prefix + + def served(self, name: str) -> str: + """The name the tool ``name`` is served under.""" + return f"{self.prefix}{name}" + + def own(self, served: str) -> str: + """A served tool's own name.""" + return served.removeprefix(self.prefix) + + def text(self, text: str | None) -> str | None: + """``text`` naming the SDK's tools as they are served; any other name stays as written.""" + return _SDK_NAMES.sub(lambda m: self.prefix + m.group(0), text) if self.prefix and text else text + + def tools(self, tools: list[EnvironmentTool]) -> list[EnvironmentTool]: + """``tools`` under their served names, the SDK's own with their descriptions' names served too.""" + return [tool.model_copy(update={"name": self.served(tool.name), **( + {"description": self.text(tool.description)} if tool.name in COMPUTER_TOOLS else {})}) for tool in tools] + + def rule(self, others: Iterable[str]) -> str | None: + """The sentence telling an agent how ``others``, the tools whose text stays as written, are served; None + without a prefix or any of them.""" + example = next(iter(others), None) + if not self.prefix or example is None: + return None + return (f"This computer serves its tools under `{self.prefix}`: a tool this text names without `{self.prefix}` is " + f"served with it, so `{example}` is `{self.prefix}{example}`. Names inside a tool's arguments stay as written.") + + def instructions(self, text: str, others: Iterable[str]) -> str: + """An MCP server's instructions as the computer serves them: after the rule, if there is one, its names served.""" + return "\n\n".join(part for part in (self.rule(others), self.text(text)) if part) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/runtime.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/runtime.py index 11d537c1e..8723045fb 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/runtime.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/runtime.py @@ -116,7 +116,7 @@ async def act(self, ctx: ToolContext, operation: str, args: tuple, *, frame: Fra async with self.queue.turn(self.driver.settings.queue_wait_s): if operation in DIALOG_BLOCKS and self.has("page") and (dialog := (await self.page_now()).dialog): raise DialogOpen(f'a {dialog.type} dialog is open on the page: "{dialog.message}"; answer it with ' - "computer_dialog first") + f"{self.driver._computer_prefix.served('computer_dialog')} first") if mapped and frame is not None: self.frames.check_current(frame, await self.display_now()) if at_pointer: diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/settings.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/settings.py index 99d295adb..5e4a629f4 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/settings.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/settings.py @@ -3,12 +3,29 @@ import json import os +import re from typing import Literal from pydantic import BaseModel, ConfigDict, Field, ValidationError SETTINGS_ENV = "AGENTENV_COMPUTER_SETTINGS" +#: The prefix a deploy gives the names of the computer's tools, so several computers serve them side by side. +TOOL_PREFIX_ENV = "AGENTENV_TOOL_PREFIX" +_TOOL_PREFIX = re.compile(r"[A-Za-z0-9_-]*") +#: The longest a tool prefix may be, its separator included, so a prefixed name stays well within the 64 characters +#: model APIs take for a tool's name. +MAX_TOOL_PREFIX = 16 + + +def tool_prefix_problem(prefix: str) -> str | None: + """What's wrong with ``prefix`` as a tool prefix, or None: only what a tool's name may hold, at most + ``MAX_TOOL_PREFIX`` characters.""" + if not _TOOL_PREFIX.fullmatch(prefix): + return f"tool prefix {prefix!r} holds characters a tool's name can't" + if len(prefix) > MAX_TOOL_PREFIX: + return f"tool prefix {prefix!r} is {len(prefix)} characters, and a tool prefix is at most {MAX_TOOL_PREFIX}" + return None Space = Literal["frame", "norm"] diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/skill.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/skill.py index 0756845a9..63b28b417 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/skill.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/skill.py @@ -61,11 +61,14 @@ async def skill(driver: ComputerDriverEnvironment) -> tuple[str, str, dict[str, if notes := driver.notes.strip(): sections += ["## Notes", notes] body = "\n\n".join(sections) + "\n" + prefix = driver._computer_prefix + lead = f"{rule}\n\n" if (rule := prefix.rule(driver._computer_other_tools())) else "" if bundled is None: - return description, body, {} + return prefix.text(description), prefix.text(lead + body), {} own_description, own_body, files = bundled - return (f"{own_description} This computer: {system}, {display.width}x{display.height}, and its computer_* tools.", - own_body.rstrip() + "\n\n" + body, files) + return (prefix.text(f"{own_description} This computer: {system}, {display.width}x{display.height}, and its " + "computer_* tools."), + prefix.text(lead + own_body.rstrip() + "\n\n" + body), files) def _input(driver: ComputerDriverEnvironment) -> str: diff --git a/packages/agentenv-protocol/src/agentenv_protocol/computer/tools.py b/packages/agentenv-protocol/src/agentenv_protocol/computer/tools.py index 17ad4e31c..d067cd95c 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/computer/tools.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/computer/tools.py @@ -9,7 +9,7 @@ from mcp.types import CallToolResult from pydantic import Field -from ..agent_env_environment import tool +from ..agent_env_environment import _TOOL_ATTR, tool from ..tool_context import ToolContext from .errors import InvalidArgument, Unsupported from .frames import Frame @@ -201,7 +201,8 @@ async def computer_type( raise InvalidArgument(f"text is at most {self.max_text_length} characters on this computer, got {len(text)}") control = sorted({f"U+{ord(c):04X}" for c in text if c not in "\n\t" and unicodedata.category(c) == "Cc"}) if control: - raise InvalidArgument(f"text has control characters {control}; press keys with computer_key") + raise InvalidArgument(f"text has control characters {control}; press keys with " + f"{self._computer_prefix.served('computer_key')}") return await computer.act(ctx, "type_text", (text,), frame=computer.frames.latest(self.caller_identity(ctx)), screenshot=screenshot, summary=f"Typed {len(text)} characters") @@ -509,3 +510,7 @@ def _check_range(name: str, value: int, low: int, high: int) -> None: def _display(frame: Frame) -> str: d = frame.display return f"{d.width}×{d.height}" + (f" at scale {d.scale:g}" if d.scale != 1 else "") + + +#: The tools the SDK serves on every computer, by their own names. +COMPUTER_TOOLS: tuple[str, ...] = tuple(sorted(name for name, member in vars(Tools).items() if hasattr(member, _TOOL_ATTR))) diff --git a/packages/agentenv-protocol/tests/test_computer_tool_prefix.py b/packages/agentenv-protocol/tests/test_computer_tool_prefix.py new file mode 100644 index 000000000..64b47a05e --- /dev/null +++ b/packages/agentenv-protocol/tests/test_computer_tool_prefix.py @@ -0,0 +1,153 @@ +"""A computer deployed with a tool prefix serves every tool under it, and names its tools that way wherever it names +them: the card, tools/list, its skill, its extensions' descriptions and its errors. Without one, nothing changes.""" +from __future__ import annotations + +import re + +import pytest + +pytest.importorskip("mcp.server.fastmcp") +pytest.importorskip("PIL") + +from computer_fake import LINUX, FakeDriver, http_client, mcp_session # noqa: E402 +from mcp.types import CallToolResult, TextContent # noqa: E402 + +from agentenv_protocol import EnvironmentTool, ForwardedTool, server_instructions # noqa: E402 +from agentenv_protocol.computer import TOOL_PREFIX_ENV # noqa: E402 + +PREFIX = "desktop_" +UNPREFIXED = re.compile(r"(? tuple[set, set]: + async with http_client(driver) as http, mcp_session(driver) as session: + card = (await http.get("/.well-known/agent-env.json")).json() + listed = {tool.name for tool in (await session.list_tools()).tools} + return {tool["name"] for tool in card["capabilities"]["tools"]}, listed + + +@pytest.mark.asyncio +async def test_a_deploys_tool_prefix_names_every_tool_on_the_card_and_in_tools_list(monkeypatch): + own_card, own_listed = await _names(FakeDriver()) + monkeypatch.setenv(TOOL_PREFIX_ENV, PREFIX) + driver = FakeDriver() + card, listed = await _names(driver) + assert card == listed == {PREFIX + name for name in own_card} and own_card == own_listed + async with mcp_session(driver) as session: + shot = await session.call_tool(f"{PREFIX}computer_screenshot", {}) + unknown = await session.call_tool("computer_screenshot", {}) + assert not shot.isError and any(part.type == "image" for part in shot.content) + assert unknown.isError + + +def test_the_declaration_a_driver_env_stores_names_the_tools_without_the_deploys_prefix(monkeypatch): + monkeypatch.setenv(TOOL_PREFIX_ENV, PREFIX) + names = {tool.name for tool in FakeDriver.declaration(LINUX).capabilities.tools} + assert "computer_screenshot" in names and not any(name.startswith(PREFIX) for name in names) + + +@pytest.mark.asyncio +async def test_its_skill_its_extensions_and_its_notes_name_the_tools_as_it_serves_them(monkeypatch): + monkeypatch.setenv(TOOL_PREFIX_ENV, PREFIX) + monkeypatch.setattr(FakeDriver, "notes", "Take a computer_screenshot first; desktop_computer_click is already right.") + async with http_client(FakeDriver()) as http: + skill = (await http.get("/agentenv/ext/skill/get", params={"name": "desktop"})).json() + card = (await http.get("/.well-known/agent-env.json")).json() + text = skill["description"] + skill["skill_md"] + assert f"`{PREFIX}computer_screenshot`" in text and f"`{PREFIX}computer_*`" in text and f"`{PREFIX}computer_click`" in text + assert "Take a desktop_computer_screenshot first; desktop_computer_click is already right." in text + assert not UNPREFIXED.search(text) and "desktop_desktop_" not in text and "serves its tools under" not in text + descriptions = " ".join(e["description"] for e in card["capabilities"]["extensions"] if e.get("description")) + assert f"The seven {PREFIX}computer_* tools" in descriptions and not UNPREFIXED.search(descriptions) + tools = " ".join(t.get("description") or "" for t in card["capabilities"]["tools"]) + assert not UNPREFIXED.search(tools) + + +@pytest.mark.asyncio +async def test_an_error_names_the_tool_to_use_instead_as_it_is_served(monkeypatch): + monkeypatch.setenv(TOOL_PREFIX_ENV, PREFIX) + async with mcp_session(FakeDriver()) as session: + refused = await session.call_tool(f"{PREFIX}computer_type", {"text": "a\x07b"}) + assert refused.isError and f"press keys with {PREFIX}computer_key" in refused.content[0].text + + +@pytest.mark.asyncio +async def test_tool_access_rules_name_the_tools_as_it_serves_them(monkeypatch): + monkeypatch.setenv(TOOL_PREFIX_ENV, PREFIX) + driver = FakeDriver() + async with http_client(driver) as http, mcp_session(driver) as session: + disabled = await http.post("/tools/disable", json={"role": "default", "tools": [f"{PREFIX}computer_click"]}) + listed = {tool.name for tool in (await session.list_tools()).tools} + refused = await session.call_tool(f"{PREFIX}computer_click", {"x": 1, "y": 1}) + assert disabled.status_code == 200 + assert f"{PREFIX}computer_click" not in listed and f"{PREFIX}computer_screenshot" in listed + assert refused.isError and "is disabled for role 'default'" in refused.content[0].text + + +@pytest.mark.asyncio +async def test_setting_the_default_space_rewrites_the_prefixed_point_tools(monkeypatch): + monkeypatch.setenv(TOOL_PREFIX_ENV, PREFIX) + driver = FakeDriver() + async with http_client(driver) as http, mcp_session(driver) as session: + await http.post("/agentenv/ext/computer-spaces/set", json={"default_space": "norm", "default_norm_scale": 1000}) + after = {t.name: t.description for t in (await session.list_tools()).tools} + assert "norm space by default: 0 to 1000" in after[f"{PREFIX}computer_click"] + + +@pytest.mark.parametrize("prefix, problem", [ + ("desk top_", "holds characters a tool's name can't"), + ("desktop/", "holds characters a tool's name can't"), + ("décor_", "holds characters a tool's name can't"), + ("accounting-desk1_", "is 17 characters, and a tool prefix is at most 16"), +]) +def test_a_prefix_a_tool_name_cant_hold_or_that_is_too_long_is_refused(monkeypatch, prefix, problem): + monkeypatch.setenv(TOOL_PREFIX_ENV, prefix) + with pytest.raises(ValueError, match=problem): + FakeDriver() + + +@pytest.mark.asyncio +async def test_a_prefix_of_sixteen_characters_is_taken(monkeypatch): + monkeypatch.setenv(TOOL_PREFIX_ENV, "accounting-desk_") + card, _listed = await _names(FakeDriver()) + assert "accounting-desk_computer_pointer_position" in card + + +class Forwarding(FakeDriver): + """A driver forwarding another program's batch tool, whose steps name that program's own actions.""" + + def forwarded_tools(self): + schema = {"type": "object", "properties": {"tool": {"type": "string", "enum": ["click", "double_click"]}}} + return [ForwardedTool(EnvironmentTool(name="run_steps", description="Run click or double_click in one call.", + inputSchema=schema), self.steps)] + + async def steps(self, tool): + return CallToolResult(content=[TextContent(type="text", text=f"ran {tool}")]) + + @server_instructions + def instructions(self): + return "Act through run_steps, and take a computer_screenshot first." + + +@pytest.mark.asyncio +async def test_a_forwarded_tool_keeps_its_own_text_and_the_skill_and_instructions_lead_with_the_prefix_rule(monkeypatch): + monkeypatch.setenv(TOOL_PREFIX_ENV, PREFIX) + driver = Forwarding() + instructions = driver.create_app().instructions + async with http_client(driver) as http, mcp_session(driver) as session: + tools = {tool.name: tool for tool in (await session.list_tools()).tools} + ran = await session.call_tool(f"{PREFIX}run_steps", {"tool": "click"}) + skill = (await http.get("/agentenv/ext/skill/get", params={"name": "desktop"})).json() + forwarded = tools[f"{PREFIX}run_steps"] + assert forwarded.description == "Run click or double_click in one call." + assert forwarded.inputSchema["properties"]["tool"]["enum"] == ["click", "double_click"] + assert not ran.isError and ran.content[0].text == "ran click" + rule = (f"This computer serves its tools under `{PREFIX}`: a tool this text names without `{PREFIX}` is served with it, " + f"so `run_steps` is `{PREFIX}run_steps`. Names inside a tool's arguments stay as written.") + assert skill["skill_md"].partition("---\n\n")[2].startswith(rule) + assert instructions == f"{rule}\n\nAct through run_steps, and take a {PREFIX}computer_screenshot first." + + +def test_without_a_prefix_a_forwarding_drivers_instructions_are_its_own(monkeypatch): + monkeypatch.delenv(TOOL_PREFIX_ENV, raising=False) + assert Forwarding().create_app().instructions == "Act through run_steps, and take a computer_screenshot first." diff --git a/src/agent_env/env/envs/computer/env.py b/src/agent_env/env/envs/computer/env.py index 93a6502e5..5979e9990 100644 --- a/src/agent_env/env/envs/computer/env.py +++ b/src/agent_env/env/envs/computer/env.py @@ -13,7 +13,7 @@ from agentenv_protocol import METHOD_ADD, RPC_PATH from agentenv_protocol import client as protocol_v1 -from agentenv_protocol.computer.tools import Tools +from agentenv_protocol.computer.tools import COMPUTER_TOOLS from agentenv_protocol.tool_access import (DISABLE, DISABLE_TOOL_EXTENSION_URI, ENABLE, ENABLE_TOOL_EXTENSION_URI, WILDCARD) from agentenv_protocol.tool_context import DEFAULT_ROLE @@ -61,7 +61,6 @@ #: so may a VM bundle, which a Mac probes. _DEPLOYABLE = {"container": ("ready",), "disk": ("ready", "unverified"), "vm_bundle": ("ready", "unverified")} #: Every ``computer_*`` tool a driver can serve. -COMPUTER_TOOLS = tuple(sorted(name for name in vars(Tools) if name.startswith("computer_"))) #: What a computer env's ``tools`` gives its agents, as the tool-access changes a deploy makes for them: ``standard``, #: AgentEnv's ``computer_*`` tools alone; ``driver``, the driver's own alone; ``both``, every tool it serves. TOOLS = {"standard": ((DISABLE, [WILDCARD]), (ENABLE, list(COMPUTER_TOOLS))), "driver": ((DISABLE, list(COMPUTER_TOOLS)),), @@ -276,12 +275,13 @@ async def deploy(self, ttl_seconds: int = 10800, disk_size_gb: float = 10, gateway_mode: GatewayMode = GatewayMode.PERFORMANCE, cpu: float | None = None, memory_mb: int | None = None, sandbox_type: str | None = None, env_state_type: str | None = None, env_state_instance_id: str | None = None, *, attribution: Optional[Attribution] = None, - env_params: Optional[dict] = None) -> DeployedEnv: + env_params: Optional[dict] = None, tool_prefix: str = "") -> DeployedEnv: + """With ``tool_prefix``, its driver serves its tools under it, as a multi env's computers serve theirs.""" return await deploy_through_provider( self, environment_name=self.environment_name, ttl_seconds=ttl_seconds, sandbox_type=sandbox_type, then=self._give_tools, disk_size_gb=disk_size_gb, gateway_mode=gateway_mode, cpu=cpu, memory_mb=memory_mb, env_state_type=env_state_type, env_state_instance_id=env_state_instance_id, attribution=attribution, - env_params=env_params) + env_params=env_params, tool_prefix=tool_prefix) async def run_env_params(self, env_params: Optional[dict], deployed_envs: list[DeployedEnv]) -> Optional[dict]: """With ``site_routing`` and no ``sites``, ``env_params`` naming as its sites the websites the run deployed @@ -390,9 +390,12 @@ async def reset(self, deployed: DeployedEnv) -> None: await self._give_tools(deployed) async def _give_tools(self, deployed: DeployedEnv) -> None: - """Make the tool-access changes ``tools`` makes for the agents' role, which a task's own then follow.""" + """Make the tool-access changes ``tools`` makes for the agents' role, which a task's own then follow, naming + each tool as the deploy's tool prefix has the driver serve it.""" + prefix = (deployed.metadata or {}).get("computer", {}).get("tool_prefix", "") for action, tools in TOOLS[self.tools or "both"]: - await deployed.invoke(_TOOL_ACCESS_URIS[action], action, {"role": DEFAULT_ROLE, "tools": tools}) + served = [tool if tool == WILDCARD else f"{prefix}{tool}" for tool in tools] + await deployed.invoke(_TOOL_ACCESS_URIS[action], action, {"role": DEFAULT_ROLE, "tools": served}) async def close(self) -> None: """End the computer's recordings, so what they recorded reaches the store, then tear it down.""" diff --git a/src/agent_env/env/envs/computer/utils/boot.py b/src/agent_env/env/envs/computer/utils/boot.py index d8fe7cf1e..f106949fe 100644 --- a/src/agent_env/env/envs/computer/utils/boot.py +++ b/src/agent_env/env/envs/computer/utils/boot.py @@ -14,7 +14,7 @@ from agentenv_protocol import CONTROL_TOKEN_ENV from agentenv_protocol import client as protocol_v1 -from agentenv_protocol.computer import APP_SETTINGS_ENV, COMPUTER_READY, SETTINGS_ENV +from agentenv_protocol.computer import APP_SETTINGS_ENV, COMPUTER_READY, SETTINGS_ENV, TOOL_PREFIX_ENV from agentenv_protocol.computer.app import COMPUTER_APP from agentenv_protocol.computer.apps import ANSWER_TIMEOUT_S @@ -89,11 +89,11 @@ def pulls_containers(link: SandboxProvider, store_id: str) -> bool: def settings_environment(computer: ComputerArtifact, values: dict[str, Any], *, environment_name: str, - driver_settings: Optional[dict] = None) -> dict[str, str]: + driver_settings: Optional[dict] = None, tool_prefix: str = "") -> dict[str, str]: """The environment a boot of ``computer`` gives it: each deploy-phase setting as ``AGENTENV_SETTING_``, from ``values`` over its default, minted afresh when it's generated and resolved through the secret store when it's a - secret; ``driver_settings``, when there are any, in ``AGENTENV_COMPUTER_SETTINGS``; and the name its card is served - under.""" + secret; ``driver_settings``, when there are any, in ``AGENTENV_COMPUTER_SETTINGS``; the name its card is served + under; and ``tool_prefix``, when there is one, in ``AGENTENV_TOOL_PREFIX``.""" environment = {"ENVIRONMENT_NAME": environment_name} for name, setting in computer.config.items(): if setting.phase != "deploy": @@ -104,6 +104,8 @@ def settings_environment(computer: ComputerArtifact, values: dict[str, Any], *, environment[f"AGENTENV_SETTING_{name.upper()}"] = str(value) if driver_settings: environment[SETTINGS_ENV] = json.dumps(driver_settings) + if tool_prefix: + environment[TOOL_PREFIX_ENV] = tool_prefix return environment diff --git a/src/agent_env/providers/env_providers/env_computer_provider.py b/src/agent_env/providers/env_providers/env_computer_provider.py index 7df383a65..2d8196b19 100644 --- a/src/agent_env/providers/env_providers/env_computer_provider.py +++ b/src/agent_env/providers/env_providers/env_computer_provider.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, ClassVar, Optional from agentenv_protocol import client as protocol_v1 -from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_RECORDING +from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_RECORDING, tool_prefix_problem from agentenv_protocol.computer.app import COMPUTER_APP from agentenv_protocol.computer.app_runtime import HOLD_BOOT_LAUNCH_ENV @@ -70,10 +70,11 @@ async def deploy( memory_mb: int | None = None, attribution: Optional[Attribution] = None, env_params: Optional[dict] = None, + tool_prefix: str = "", ) -> DeployedSandboxEnv: """The env's first ready build, in its base's order of forms and arches (or of ``env_params.form``), that the sandbox provider runs natively, booted, once what it boots that no compose started is composed; returns the record once - the driver is ready with the env's display.""" + the driver is ready with the env's display. With ``tool_prefix``, its driver serves its tools under it.""" from agent_env.env.envs.computer.env import ComputerEnv if not isinstance(env, ComputerEnv): @@ -81,6 +82,8 @@ async def deploy( if env.computer is None: # an app that names no computer to deploy alone on raise DeployRefused(env.deploy_refusal()) params, values, problems = computer_params.checked(env, env_params) + if problem := tool_prefix_problem(tool_prefix): + problems.append(problem) if problems: raise DeployRefused(f"{env.id}: " + "; ".join(problems)) if params.site_routing and not params.sites: @@ -94,7 +97,8 @@ async def deploy( raise DeployRefused(f"{env.id}: env_params.site_routing installs into the computer's container, and a disk " "computer runs inside a launcher container it can't reach; deploy a container form") attempt = functools.partial(self._boot, env, builds, params, values, ttl_seconds=ttl_seconds, - disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, attribution=attribution) + disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, attribution=attribution, + tool_prefix=tool_prefix) return await self._run(sandbox_provider, env.id, attempt) async def install_changelog_triggers(self, environment_name: str) -> None: @@ -103,7 +107,7 @@ async def install_changelog_triggers(self, environment_name: str) -> None: async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], params: computer_params.ComputerParams, values: dict, sandbox_provider: SandboxProvider, *, ttl_seconds: int, disk_size_gb: float, cpu: float | None, memory_mb: int | None, - attribution: Optional[Attribution]) -> DeployedSandboxEnv: + attribution: Optional[Attribution], tool_prefix: str = "") -> DeployedSandboxEnv: """One attempt on one sandbox provider: the first build it runs natively, booted and readied; its record.""" from agent_env.env.envs.computer.env import ComputerEnv @@ -131,7 +135,8 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par policy = NetworkPolicy.from_dict(env.network) if env.network else None booted = await boot( sandbox_provider, image=build["image"], port=env.driver.port, network_policy=policy, token=token, - environment=_environment(env, computer, values, params), timeout=ttl_seconds, attribution=attribution, + environment=_environment(env, computer, values, params, tool_prefix), timeout=ttl_seconds, + attribution=attribution, disk_size_gb=disk_size_gb, cpu=cpu, memory_mb=memory_mb, ready_s=ready_timeout(env.composed_apps(), env.computer.os.type), form=form, held=bool(params.site_routing), os_type=env.computer.os.type) @@ -151,6 +156,7 @@ async def _boot(self, env: ComputerEnv, builds: list[tuple[str, str, dict]], par "port": env.driver.port, "display": values.get("resolution"), "network": env.network, "computer_settings": recorded_settings(computer, values), **({"app_settings": app_settings} if app_settings else {}), + **({"tool_prefix": tool_prefix} if tool_prefix else {}), "env_params": params.model_dump(exclude_defaults=True)}}, ) if routed := await _finish_boot(env, record, booted.sandbox, params, ttl_seconds, values.get("resolution")): @@ -171,7 +177,8 @@ async def reboot(self, env: ComputerEnv, deployed: DeployedSandboxEnv) -> None: await end_recordings(deployed, "reset") sandbox_provider = build_sandbox_provider(deployed.sandbox_type) sandbox = await sandbox_provider.get_sandbox(deployed.sandbox_id) - environment = {**_environment(env, computer, values, params), **control_environment(deployed.control_token)} + environment = {**_environment(env, computer, values, params, record.get("tool_prefix", "")), + **control_environment(deployed.control_token)} try: if record.get("form") in ("disk", "vm_bundle"): await sandbox_provider.reset_vm(sandbox, image=record["image"], env=environment) @@ -249,9 +256,9 @@ def _values(computer: ComputerArtifact, values: dict) -> dict: def _environment(env: ComputerEnv, computer: ComputerArtifact, values: dict, - params: computer_params.ComputerParams) -> dict[str, str]: + params: computer_params.ComputerParams, tool_prefix: str = "") -> dict[str, str]: return {**settings_environment(computer, values, environment_name=env.computer_name, - driver_settings=params.driver_settings), + driver_settings=params.driver_settings, tool_prefix=tool_prefix), **app_settings_environment(env.composed_apps(), computer_params.app_settings(env, params.app_settings), env.computer.os.type), **({HOLD_BOOT_LAUNCH_ENV: "1"} if params.site_routing else {})} diff --git a/tst/integration/computer/tool_prefix_test.py b/tst/integration/computer/tool_prefix_test.py new file mode 100644 index 000000000..7454caae1 --- /dev/null +++ b/tst/integration/computer/tool_prefix_test.py @@ -0,0 +1,60 @@ +"""The desktop example's computer deployed twice side by side with tool prefixes, as a multi env's computers will be: +each serves every tool under its own prefix and takes screenshots by those names, its skill names its tools that way, +and one whose ``tools`` is ``standard`` gives its agents only its prefixed ``computer_*`` tools, at deploy and again +after a reset. Needs Docker and the local registry on localhost:5000.""" + +import re + +import pytest +from mcp import ClientSession +from mcp.client.streamable_http import streamablehttp_client + +from agentenv_protocol.types import SKILL_EXTENSION_URI + +from agent_env.env import Env +from agent_env.env.envs.computer.env import COMPUTER_TOOLS, ComputerEnv +from agent_env.store.routing import namespace_routing +from agent_env.task.teardown import teardown_run +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.reset_env import ResetEnvTaskStep + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + +UNPREFIXED = re.compile(r"(? tuple[set[str], bool]: + """The tools the computer lists to its agents, and whether a screenshot by its prefixed name answered.""" + async with streamablehttp_client(deployed.mcp_url) as (read, write, _): + async with ClientSession(read, write) as session: + await session.initialize() + names = {tool.name for tool in (await session.list_tools()).tools} + shot = await session.call_tool(f"{prefix}computer_screenshot", {}) + return names, not shot.isError and any(part.type == "image" for part in shot.content) + + +@pytest.mark.asyncio +async def test_two_computers_serve_their_tools_under_their_own_prefixes_side_by_side_and_through_a_reset(desktop): + with namespace_routing(): + standard = ComputerEnv.put(id=f"{desktop.id}-standard", computer=desktop.computer, driver=desktop.driver, + environment_name="laptop", tools="standard") + context = TaskStepContext() + try: + context.deployed_envs.append(await Env.get(desktop.id).deploy(tool_prefix="desktop_")) + context.deployed_envs.append(await Env.get(standard.id).deploy(tool_prefix="laptop_")) + left, right = context.deployed_envs + left_names, left_shot = await _served(left, "desktop_") + right_names, right_shot = await _served(right, "laptop_") + skill = await left.invoke(SKILL_EXTENSION_URI, "get", {"name": "desktop"}) + await ResetEnvTaskStep(id="reset", version=None, env_id=standard.id, fail_task_on_error=True).execute(context) + after_reset, after_shot = await _served(right, "laptop_") + finally: + report = await teardown_run(context) + + assert not report.failed + assert left_shot and right_shot and after_shot + assert left_names and all(name.startswith("desktop_") for name in left_names) + assert "desktop_computer_screenshot" in left_names and "computer_screenshot" not in left_names + assert right_names == after_reset and right_names <= {f"laptop_{tool}" for tool in COMPUTER_TOOLS} + assert "laptop_computer_click" in right_names + assert "`desktop_computer_screenshot`" in skill["skill_md"] and not UNPREFIXED.search(skill["skill_md"]) diff --git a/tst/unit/env/envs/test_computer_env.py b/tst/unit/env/envs/test_computer_env.py index a8a363820..93b93daeb 100644 --- a/tst/unit/env/envs/test_computer_env.py +++ b/tst/unit/env/envs/test_computer_env.py @@ -22,7 +22,7 @@ import httpx import pytest from agentenv_protocol import CONTROL_TOKEN_ENV -from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_READY, COMPUTER_RECORDING +from agentenv_protocol.computer import COMPUTER_DISPLAY, COMPUTER_READY, COMPUTER_RECORDING, TOOL_PREFIX_ENV from agentenv_protocol.computer.app import COMPUTER_APP from agentenv_protocol.computer.app_runtime import HOLD_BOOT_LAUNCH_ENV from agentenv_protocol.tool_access import DISABLE_TOOL_EXTENSION_URI, ENABLE_TOOL_EXTENSION_URI @@ -783,6 +783,44 @@ async def test_a_computer_envs_tools_are_given_its_agents_at_deploy_and_again_af assert [call for call in boots.calls if call[0] in (DISABLE_TOOL_EXTENSION_URI, ENABLE_TOOL_EXTENSION_URI)] == changes +@pytest.mark.asyncio +async def test_a_deploy_with_a_tool_prefix_has_its_driver_serve_its_tools_under_it_and_gives_them_by_those_names( + base, driver, docker, boots, monkeypatch): + env = _put(base, driver, display={"width": 1600, "height": 900}, tools="standard") + deployed = await Env.get(env.id).deploy(tool_prefix="desktop_") + (booted,), (sandbox,) = boots.boots, boots.sandboxes + given = [call for call in boots.calls if call[0] in (DISABLE_TOOL_EXTENSION_URI, ENABLE_TOOL_EXTENSION_URI)] + reset = AsyncMock() + monkeypatch.setattr(local_sandbox.LocalSandboxProvider, "get_sandbox", AsyncMock(return_value=sandbox)) + monkeypatch.setattr(local_sandbox.LocalSandboxProvider, "reset_container", reset) + monkeypatch.setattr(env_computer_provider, "ready", AsyncMock(return_value=DRIVER_CARD)) + boots.calls.clear() + + await ResetEnvTaskStep(id="reset", version=None, env_id=env.id, fail_task_on_error=True).execute( + TaskStepContext(deployed_envs=[deployed])) + + changes = [(DISABLE_TOOL_EXTENSION_URI, "disable", {"role": "default", "tools": ["*"]}), + (ENABLE_TOOL_EXTENSION_URI, "enable", {"role": "default", + "tools": [f"desktop_{tool}" for tool in COMPUTER_TOOLS]})] + assert booted["environment"][TOOL_PREFIX_ENV] == "desktop_" == deployed.metadata["computer"]["tool_prefix"] + assert reset.call_args.kwargs["env"][TOOL_PREFIX_ENV] == "desktop_" + assert given == changes + assert [call for call in boots.calls if call[0] in (DISABLE_TOOL_EXTENSION_URI, ENABLE_TOOL_EXTENSION_URI)] == changes + + +@pytest.mark.asyncio +async def test_a_deploy_without_a_tool_prefix_gives_its_driver_none(base, driver, docker, boots): + deployed = await Env.get(_put(base, driver, display={"width": 1600, "height": 900}).id).deploy() + assert TOOL_PREFIX_ENV not in boots.boots[0]["environment"] and "tool_prefix" not in deployed.metadata["computer"] + + +@pytest.mark.asyncio +async def test_a_tool_prefix_a_tool_name_cant_hold_is_refused_before_anything_boots(base, driver, docker, boots): + with pytest.raises(Exception, match="tool prefix 'desk top_' holds characters a tool's name can't"): + await Env.get(_put(base, driver, display={"width": 1600, "height": 900}).id).deploy(tool_prefix="desk top_") + assert boots.boots == [] + + @pytest.mark.asyncio async def test_a_deploy_whose_tools_the_driver_cant_take_fails_and_tears_down_its_computer_alone( base, driver, docker, boots, monkeypatch):