From b2e56a892abc3650a1bca94fb08231a42c284d3a Mon Sep 17 00:00:00 2001 From: yanghailan Date: Fri, 13 Nov 2026 11:54:55 +0800 Subject: [PATCH] feat: add AgentKit CUA computer-use toolset --- tests/tools/cua_mcp/conftest.py | 22 + tests/tools/cua_mcp/fakes.py | 65 ++ tests/tools/cua_mcp/requirements.txt | 8 + tests/tools/cua_mcp/stdio_fixture.py | 44 ++ tests/tools/cua_mcp/test_cua_isolation.py | 194 ++++++ tests/tools/cua_mcp/test_distribution.py | 72 +++ tests/tools/cua_mcp/test_responses.py | 367 +++++++++++ tests/tools/cua_mcp/test_serial_guidance.py | 203 +++++++ tests/tools/cua_mcp/test_toolset.py | 642 ++++++++++++++++++++ veadk/tools/builtin_tools/computer_use.py | 445 ++++++++++++++ veadk/tools/computer_use/__init__.py | 1 + veadk/tools/computer_use/media.py | 36 ++ 12 files changed, 2099 insertions(+) create mode 100644 tests/tools/cua_mcp/conftest.py create mode 100644 tests/tools/cua_mcp/fakes.py create mode 100644 tests/tools/cua_mcp/requirements.txt create mode 100644 tests/tools/cua_mcp/stdio_fixture.py create mode 100644 tests/tools/cua_mcp/test_cua_isolation.py create mode 100644 tests/tools/cua_mcp/test_distribution.py create mode 100644 tests/tools/cua_mcp/test_responses.py create mode 100644 tests/tools/cua_mcp/test_serial_guidance.py create mode 100644 tests/tools/cua_mcp/test_toolset.py create mode 100644 veadk/tools/builtin_tools/computer_use.py create mode 100644 veadk/tools/computer_use/__init__.py create mode 100644 veadk/tools/computer_use/media.py diff --git a/tests/tools/cua_mcp/conftest.py b/tests/tools/cua_mcp/conftest.py new file mode 100644 index 000000000..8bc2930fb --- /dev/null +++ b/tests/tools/cua_mcp/conftest.py @@ -0,0 +1,22 @@ +"""CUA integration tests run when the independent MCP wheel is installed. + +Install tests/tools/cua_mcp/requirements.txt to execute this suite. Without the +wheel these tests are reported as skipped; no sibling repository or private MCP +source checkout is loaded into sys.path. +""" + +import importlib.util + +import pytest + +_CUA_MCP_INSTALLED = importlib.util.find_spec("agentkit_cua_mcp") is not None + + +def pytest_collection_modifyitems(config, items): + if _CUA_MCP_INSTALLED: + return + skip = pytest.mark.skip( + reason="install tests/tools/cua_mcp/requirements.txt to run CUA integration tests" + ) + for item in items: + item.add_marker(skip) diff --git a/tests/tools/cua_mcp/fakes.py b/tests/tools/cua_mcp/fakes.py new file mode 100644 index 000000000..f95fb6cb9 --- /dev/null +++ b/tests/tools/cua_mcp/fakes.py @@ -0,0 +1,65 @@ +"""Cloud-free platform contract fixtures for the installed MCP distribution.""" + +from datetime import datetime, timedelta, timezone + +from agentkit_cua_mcp.service import CuaError + + +class Platform: + def __init__(self): + self.sessions = {} + self.created = 0 + self.fail = False + self.envs = [{"Key": "ENABLE_COMPUTER", "Value": "true"}] + + async def request(self, action, body): + if self.fail: + raise CuaError("PLATFORM_UNAVAILABLE") + if action == "GetTool": + return {"ToolId": "t-example", "Status": "Ready", "Envs": self.envs} + if action == "ListSessions": + assert body["Filters"] == [ + {"Name": "UserSessionId", "Values": [body["Filters"][0]["Values"][0]]} + ] + return {"SessionInfos": list(self.sessions.values())} + if action == "CreateSession": + self.created += 1 + sid = f"s-{self.created}" + self.sessions[sid] = { + "SessionId": sid, + "UserSessionId": body["UserSessionId"], + "Status": "Ready", + "Endpoint": "https://sandbox.example.test", + "ExpireAt": ( + datetime.now(timezone.utc) + timedelta(minutes=20) + ).isoformat(), + } + return self.sessions[sid] + if action == "GetSession": + if body["SessionId"] not in self.sessions: + raise CuaError("SESSION_NOT_FOUND") + return self.sessions[body["SessionId"]] + raise AssertionError(action) + + +class Desktop: + enabled = True + ready = True + called = [] + + def __init__(self, endpoint): + self.endpoint = endpoint + + async def computer_capability(self): + return { + "enabled": self.enabled, + "ready": self.ready, + "state": "disabled" + if not self.enabled + else ("ready" if self.ready else "starting"), + "features": [], + } + + async def call(self, name, arguments): + self.called.append((name, arguments)) + return {"success": True} diff --git a/tests/tools/cua_mcp/requirements.txt b/tests/tools/cua_mcp/requirements.txt new file mode 100644 index 000000000..ed9680052 --- /dev/null +++ b/tests/tools/cua_mcp/requirements.txt @@ -0,0 +1,8 @@ +# Published baseline, pinned identically to the client default. +# Current schema-description tests require the newer local candidate until it is +# published. From the VeADK repository root, with a sibling cua-sandbox checkout: +# uv build --wheel --out-dir /tmp/cua-mcp-candidate ../cua-sandbox/packages/agentkit-cua-mcp +# uv pip install --python .venv/bin/python --reinstall /tmp/cua-mcp-candidate/agentkit_cua_mcp-1.0.0-py3-none-any.whl +# uv run --no-sync pytest -q tests/tools/cua_mcp +# Demo and bilingual integration guide: cua-sandbox repository, examples/agentkit-cua/. +agentkit-cua-mcp @ https://browser-use-test.tos-cn-beijing.volces.com/agentkit-cua-mcp/1.0.0/agentkit_cua_mcp-1.0.0-py3-none-any.whl?revision=bdf9d949ae148fb3c9b6953fa67ec65bf136681502c0737971a3e0751f128840#sha256=bdf9d949ae148fb3c9b6953fa67ec65bf136681502c0737971a3e0751f128840 diff --git a/tests/tools/cua_mcp/stdio_fixture.py b/tests/tools/cua_mcp/stdio_fixture.py new file mode 100644 index 000000000..56a1ac22c --- /dev/null +++ b/tests/tools/cua_mcp/stdio_fixture.py @@ -0,0 +1,44 @@ +"""Cloud-free subprocess used only by integration tests; no production fake mode.""" + +import asyncio +import base64 +import sys +from fakes import Platform, Desktop +from agentkit_cua_mcp.service import SandboxService +from agentkit_cua_mcp.server import run + + +class ImageDesktop(Desktop): + async def call(self, name, arguments): + if name == "cua_screenshot": + return b"\x89PNG\r\n\x1a\nTEST" + if name == "cua_click" and arguments.get("include_screenshot"): + return { + "success": True, + "data": {"action_performed": "CLICK"}, + "screenshot": base64.b64encode(b"\x89PNG\r\n\x1a\nTEST").decode(), + } + if name == "cua_actions_batch": + return { + "success": True, + "data": { + "performed": [a["action_type"] for a in arguments["actions"]], + "failed_index": None, + "error": None, + }, + } + return await super().call(name, arguments) + + +async def main(): + platform = Platform() + if "--disabled" in sys.argv: + platform.envs = [] + ImageDesktop.enabled = False + ImageDesktop.ready = False + if "--existing" in sys.argv: + await platform.request("CreateSession", {"UserSessionId": "agent_user_one"}) + await run(SandboxService(platform, "t-example", desktop_factory=ImageDesktop)) + + +asyncio.run(main()) diff --git a/tests/tools/cua_mcp/test_cua_isolation.py b/tests/tools/cua_mcp/test_cua_isolation.py new file mode 100644 index 000000000..028fc5dcb --- /dev/null +++ b/tests/tools/cua_mcp/test_cua_isolation.py @@ -0,0 +1,194 @@ +import ast +from pathlib import Path + +import pytest +from google.adk.models import LlmRequest +from google.genai import types + +from veadk.models.ark_llm import _get_responses_inputs +from test_toolset import context, toolset + + +@pytest.mark.parametrize( + "name", + [ + "cua_screenshot", + "cua_actions_batch", + "cua_click", + "cua_typing", + "cua_node_invoke", + ], +) +def test_every_cua_image_reaches_unmodified_ark_without_editing_event(name): + from veadk.tools.computer_use.media import lift_screenshots + + original = types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + name=name, + id="call", + response={"done": True}, + parts=[ + types.FunctionResponsePart( + inline_data=types.FunctionResponseBlob( + mime_type="image/png", data=b"png-fixture" + ) + ) + ], + ) + ) + ], + ) + request = LlmRequest(contents=[original]) + lift_screenshots(request) + lift_screenshots(request) + inputs = _get_responses_inputs(request)[1] + assert ( + sum( + p.get("type") == "input_image" for i in inputs for p in i.get("content", []) + ) + == 1 + ) + assert original.parts[0].function_response.parts + assert not request.contents[0].parts[0].function_response.parts + + +@pytest.mark.asyncio +@pytest.mark.parametrize("previous_id", [None, "resp-existing"]) +async def test_toolset_preserves_request_context_without_rebuilding_history( + monkeypatch, previous_id +): + from google.adk.flows.llm_flows.contents import request_processor + from veadk.models.ark_llm import ArkLlm + + def reject_history_rebuild(*args, **kwargs): + raise AssertionError("CUA must not rerun ADK's history processor") + + monkeypatch.setattr(request_processor, "run_async", reject_history_rebuild) + cua = toolset() + ctx, _ = await context("request-only") + ctx._invocation_context.agent.model = ArkLlm(model="openai/offline") + request = LlmRequest( + previous_interaction_id=previous_id, + contents=[ + types.Content( + role="user", parts=[types.Part(text="request-only instruction")] + ) + ], + ) + before = request.model_dump() + try: + await cua.process_llm_request(tool_context=ctx, llm_request=request) + assert request.model_dump() == before + finally: + await cua.close() + + +def test_generic_model_source_has_no_cua_branch(): + path = Path(__file__).resolve().parents[3] / "veadk/models/ark_llm.py" + source = path.read_text() + assert "_veadk_cua" not in source + assert "_cua_responses" not in source + ast.parse(source) + + +@pytest.mark.asyncio +async def test_toolset_does_not_replace_model_or_require_binding(): + cua = toolset() + ctx, _ = await context("one") + agent = ctx._invocation_context.agent + model = agent.model + request = LlmRequest(contents=[]) + try: + await cua.process_llm_request(tool_context=ctx, llm_request=request) + assert agent.model is model + assert not hasattr(request, "_veadk_cua_context") + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_fixed_catalog_rejects_filter_change_without_model_call(): + cua = toolset() + ctx, readonly = await context("one") + try: + await cua.get_tools(readonly) + cua.tool_filter = ["cua_screenshot"] + with pytest.raises(RuntimeError, match="catalog.*changed"): + await cua.get_tools(readonly) + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_legacy_mcp_discovery_contract_fails_before_model(): + from unittest.mock import AsyncMock + from mcp import types as mcp_types + + cua = toolset() + _, readonly = await context("legacy") + session = AsyncMock() + session.list_tools.return_value = mcp_types.ListToolsResult( + tools=[], _meta={"veadk/cua": {"discovery_mode": "lazy-v1"}} + ) + cua._manager.create_session = AsyncMock(return_value=session) + try: + with pytest.raises(RuntimeError, match="discovery failed"): + await cua.get_tools(readonly) + finally: + await cua.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["none", "template", "schema"]) +async def test_catalog_is_checked_after_child_restart_and_new_invocation(change): + from unittest.mock import AsyncMock + from google.adk.agents.invocation_context import InvocationContext + from google.adk.agents.readonly_context import ReadonlyContext + + cua = toolset() + ctx, readonly = await context("restart_catalog") + invocation = ctx._invocation_context + try: + original = await cua.get_tools(readonly) + saved = await invocation.session_service.get_session( + app_name="app", user_id="user", session_id="restart_catalog" + ) + assert any(key.startswith("veadk_cua_catalog_") for key in saved.state) + assert ctx.state.get(cua.binding_state_key(ctx)) is None + finally: + await cua.close() + + restored = toolset() + new_context = ReadonlyContext( + InvocationContext( + session_service=invocation.session_service, + invocation_id="next-turn", + agent=invocation.agent, + session=saved, + ) + ) + try: + if change != "none": + session = await restored._manager.create_session() + from mcp import types as mcp_types + + result = await session.list_tools( + params=mcp_types.PaginatedRequestParams( + _meta=restored._metadata(new_context) + ) + ) + if change == "template": + result.tools = [t for t in result.tools if t.name == "ensure_sandbox"] + else: + result.tools[1].inputSchema["description"] = "changed schema" + session.list_tools = AsyncMock(return_value=result) + with pytest.raises(RuntimeError, match="catalog.*changed"): + await restored.get_tools(new_context) + else: + tools = await restored.get_tools(new_context) + assert [t.name for t in tools] == [t.name for t in original] + finally: + await restored.close() diff --git a/tests/tools/cua_mcp/test_distribution.py b/tests/tools/cua_mcp/test_distribution.py new file mode 100644 index 000000000..815e3b711 --- /dev/null +++ b/tests/tools/cua_mcp/test_distribution.py @@ -0,0 +1,72 @@ +"""Distribution changes must not alter the CUA request/session contract.""" + +import asyncio + +import pytest + +from mcp import StdioServerParameters + +from veadk.tools.builtin_tools import computer_use + + +def capture_params(monkeypatch): + captured = [] + + class Manager: + def __init__(self, params): + captured.append(params.server_params) + + async def close(self): + pass + + monkeypatch.setattr(computer_use, "MCPSessionManager", Manager) + return captured + + +def test_missing_wheel_url_fails_with_configuration_guidance(monkeypatch): + capture_params(monkeypatch) + monkeypatch.delenv("CUA_MCP_WHEEL_URL", raising=False) + with pytest.raises(ValueError, match="CUA_MCP_WHEEL_URL"): + computer_use.create_computer_use_toolset(tool_id="t-example") + + +def test_empty_wheel_url_fails_with_configuration_guidance(monkeypatch): + capture_params(monkeypatch) + monkeypatch.setenv("CUA_MCP_WHEEL_URL", " ") + with pytest.raises(ValueError, match="CUA_MCP_WHEEL_URL"): + computer_use.create_computer_use_toolset(tool_id="t-example") + + +def test_explicit_wheel_and_uv_cache_configuration(monkeypatch): + captured = capture_params(monkeypatch) + wheel = "https://downloads.example.test/agentkit_cua_mcp-1.0.0-py3-none-any.whl" + monkeypatch.setenv("CUA_MCP_WHEEL_URL", wheel) + monkeypatch.setenv("UV_CACHE_DIR", "/opt/uv-cache") + monkeypatch.setenv("UV_OFFLINE", "true") + monkeypatch.setenv("UV_INDEX_URL", "https://pypi.org/simple") + monkeypatch.setenv("UNRELATED_SECRET", "must-not-forward") + cua = computer_use.create_computer_use_toolset(tool_id="t-example") + params = captured[0] + assert params.command == "uvx" + assert params.args[1] == wheel + assert params.env["UV_CACHE_DIR"] == "/opt/uv-cache" + assert params.env["UV_OFFLINE"] == "true" + assert params.env["UV_INDEX_URL"] == "https://pypi.org/simple" + assert "UNRELATED_SECRET" not in params.env + assert "CUA_MCP_WHEEL_URL" not in params.env + asyncio.run(cua.close()) + + +def test_explicit_connection_still_supports_preinstalled_server(monkeypatch): + captured = capture_params(monkeypatch) + monkeypatch.setenv("CUA_MCP_WHEEL_URL", "not-used") + params = StdioServerParameters( + command="agentkit-cua-mcp", args=["--transport", "stdio"], env={} + ) + cua = computer_use.create_computer_use_toolset( + tool_id="t-example", connection_params=params + ) + assert captured[0].command == "agentkit-cua-mcp" + assert captured[0].args == params.args + assert params.env == {} # Do not mutate caller configuration. + asyncio.run(cua.close()) diff --git a/tests/tools/cua_mcp/test_responses.py b/tests/tools/cua_mcp/test_responses.py new file mode 100644 index 000000000..7a6057497 --- /dev/null +++ b/tests/tools/cua_mcp/test_responses.py @@ -0,0 +1,367 @@ +"""Real ADK + stdio + Ark adapter; only cloud and model HTTP are simulated.""" + +import copy + +import pytest +from google.adk.models import LlmResponse +from google.adk.runners import InMemoryRunner +from google.genai import types + +from test_toolset import toolset +from veadk.models.ark_llm import ArkLlm, request_reorganization_by_ark + + +@pytest.mark.asyncio +async def test_expired_model_response_uses_original_ark_retry_without_cua_history( + monkeypatch, +): + import httpx + from google.adk.agents import Agent + from volcenginesdkarkruntime._exceptions import ArkBadRequestError + + async def run_case(use_cua): + requests = [] + + async def fake_generate(self, args, stream=False): + requests.append(copy.deepcopy(args)) + if len(requests) == 2: + request = httpx.Request("POST", "https://model.example.test/responses") + raise ArkBadRequestError( + "expired", + response=httpx.Response(400, request=request), + body={"code": "InvalidParameter.PreviousResponseNotFound"}, + request_id="test", + ) + yield LlmResponse( + model_version="offline", + content=types.Content(role="model", parts=[types.Part(text="done")]), + interaction_id=f"resp-{len(requests)}", + ) + + monkeypatch.setattr(ArkLlm, "generate_content_via_responses", fake_generate) + cua = toolset() if use_cua else None + runner = InMemoryRunner( + agent=Agent( + name="agent", + model=ArkLlm(model="openai/offline"), + tools=[cua] if cua else [], + ), + app_name="app", + ) + await runner.session_service.create_session( + app_name="app", user_id="user", session_id="one" + ) + try: + for text in ["remember-me", "continue"]: + _ = [ + e + async for e in runner.run_async( + user_id="user", + session_id="one", + new_message=types.Content( + role="user", parts=[types.Part(text=text)] + ), + ) + ] + assert len(requests) == 3 + assert requests[1]["previous_response_id"] == "resp-1" + assert not requests[2].get("previous_response_id") + if cua: + saved = await runner.session_service.get_session( + app_name="app", user_id="user", session_id="one" + ) + assert not any(k.startswith(cua.state_key) for k in saved.state) + return requests + finally: + await runner.close() + + baseline = await run_case(use_cua=False) + with_cua = await run_case(use_cua=True) + # CUA must not augment either the normal request or the provider's retry. + assert [r["input"] for r in with_cua] == [r["input"] for r in baseline] + + +@pytest.mark.asyncio +async def test_real_ark_serialization_does_not_leak_internal_markers(monkeypatch): + from google.adk.agents import Agent + from veadk.models.ark_llm import ArkLlmClient, ArkTypeResponse + + wire_requests = [] + + async def aresponses(self, **kwargs): + wire_requests.append(kwargs) + return ArkTypeResponse.model_validate( + { + "id": f"resp-{len(wire_requests)}", + "object": "response", + "created_at": 1, + "model": "provider-resolved-model-version", + "status": "completed", + "incomplete_details": None, + "tools": [], + "output": [ + { + "type": "message", + "id": "msg", + "role": "assistant", + "status": "completed", + "content": [ + {"type": "output_text", "text": "done", "annotations": []} + ], + } + ], + } + ) + + monkeypatch.setattr(ArkLlmClient, "aresponses", aresponses) + cua = toolset() + runner = InMemoryRunner( + agent=Agent( + name="agent", + model=ArkLlm( + model="openai/offline", extra_body={"parallel_tool_calls": False} + ), + tools=[cua], + ), + app_name="app", + ) + await runner.session_service.create_session( + app_name="app", user_id="user", session_id="one" + ) + try: + for text in ["first", "second", "third"]: + _ = [ + e + async for e in runner.run_async( + user_id="user", + session_id="one", + new_message=types.Content( + role="user", parts=[types.Part(text=text)] + ), + ) + ] + assert "tools" in wire_requests[0] + assert wire_requests[1]["previous_response_id"] == "resp-1" + assert "tools" not in wire_requests[1] + assert wire_requests[2]["previous_response_id"] == "resp-2" + assert not any(k.startswith("_veadk") for w in wire_requests for k in w) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_same_model_and_toolset_do_not_share_response_chain_between_sessions( + monkeypatch, +): + from google.adk.agents import Agent + + calls = [] + + async def fake_generate(self, args, stream=False): + calls.append(copy.deepcopy(args)) + yield LlmResponse( + model_version="offline", + interaction_id=f"resp-{len(calls)}", + content=types.Content(role="model", parts=[types.Part(text="done")]), + ) + + monkeypatch.setattr(ArkLlm, "generate_content_via_responses", fake_generate) + cua = toolset() + runner = InMemoryRunner( + agent=Agent(name="agent", model=ArkLlm(model="openai/offline"), tools=[cua]), + app_name="app", + ) + for session_id in ["one", "two"]: + await runner.session_service.create_session( + app_name="app", user_id="user", session_id=session_id + ) + try: + for session_id in ["one", "two", "one", "two"]: + _ = [ + e + async for e in runner.run_async( + user_id="user", + session_id=session_id, + new_message=types.Content( + role="user", parts=[types.Part(text=session_id)] + ), + ) + ] + assert [c.get("previous_response_id") for c in calls] == [ + None, + None, + "resp-1", + "resp-2", + ] + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_explicit_cache_disable_and_disabled_cua_preserve_original_context( + monkeypatch, +): + import sys + from pathlib import Path + from mcp import StdioServerParameters + from google.adk.agents import Agent + from veadk.tools.builtin_tools.computer_use import create_computer_use_toolset + + def run_code(code: str) -> str: + """Stub existing AIO tool for model-visible catalog checks.""" + return code + + async def run_case(use_cua): + requests = [] + + async def fake_generate(self, args, stream=False): + requests.append(copy.deepcopy(args)) + yield LlmResponse( + model_version="offline", + interaction_id="response", + content=types.Content(role="model", parts=[types.Part(text="done")]), + ) + + monkeypatch.setattr(ArkLlm, "generate_content_via_responses", fake_generate) + tools = [run_code] + if use_cua: + tools.append( + create_computer_use_toolset( + tool_id="t-example", + connection_params=StdioServerParameters( + command=sys.executable, + args=[ + str(Path(__file__).with_name("stdio_fixture.py")), + "--disabled", + ], + ), + ) + ) + runner = InMemoryRunner( + agent=Agent( + name="agent", + model=ArkLlm(model="openai/offline", enable_responses_cache=False), + tools=tools, + ), + app_name="app", + ) + await runner.session_service.create_session( + app_name="app", user_id="user", session_id="one" + ) + try: + for text in ["original-detail", "second-turn"]: + _ = [ + e + async for e in runner.run_async( + user_id="user", + session_id="one", + new_message=types.Content( + role="user", parts=[types.Part(text=text)] + ), + ) + ] + assert {t["name"] for t in requests[0]["tools"]} == ( + {"run_code", "ensure_sandbox"} if use_cua else {"run_code"} + ) + assert all(not r.get("previous_response_id") for r in requests) + return requests + finally: + await runner.close() + + baseline = await run_case(use_cua=False) + with_cua = await run_case(use_cua=True) + assert [r["input"] for r in with_cua] == [r["input"] for r in baseline] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "tool_name,args", + [("cua_screenshot", {}), ("cua_click", {"include_screenshot": True})], +) +async def test_fixed_tools_and_images_with_original_ark_cache( + monkeypatch, tool_name, args +): + from google.adk.agents import Agent + + requests = [] + + async def fake_generate(self, response_args, stream=False): + wire = request_reorganization_by_ark(copy.deepcopy(response_args)) + assert "parallel_tool_calls" not in wire + assert "parallel_tool_calls" not in wire.get("extra_body", {}) + requests.append(wire) + n = len(requests) + yield LlmResponse( + interaction_id=f"resp-{n}", + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + id="action", name=tool_name, args=args + ) + ) + if n == 1 + else types.Part(text="done") + ], + ), + ) + + monkeypatch.setattr(ArkLlm, "generate_content_via_responses", fake_generate) + cua = toolset() + model = ArkLlm(model="openai/offline", extra_body={"caching": {"type": "enabled"}}) + runner = InMemoryRunner( + agent=Agent(name="agent", model=model, tools=[cua]), app_name="app" + ) + await runner.session_service.create_session( + app_name="app", user_id="user", session_id="one" + ) + try: + for text in ["first", "next"]: + _ = [ + e + async for e in runner.run_async( + user_id="user", + session_id="one", + new_message=types.Content( + role="user", parts=[types.Part(text=text)] + ), + ) + ] + assert "cua_accessibility_nodes" in {t["name"] for t in requests[0]["tools"]} + assert requests[1]["previous_response_id"] == "resp-1" + assert "tools" not in requests[1] + assert ( + sum( + p.get("type") == "input_image" + for i in requests[1]["input"] + for p in i.get("content", []) + ) + == 1 + ) + assert requests[2]["previous_response_id"] == "resp-2" + assert runner.agent.model is model + saved = await runner.session_service.get_session( + app_name="app", user_id="user", session_id="one" + ) + assert ( + saved.state[ + cua.binding_state_key(type("Ctx", (), {"agent_name": "agent"})()) + ]["session_id"] + == "s-1" + ) + finally: + await runner.close() + + +def test_original_ark_parallel_parameter_behavior_is_not_silently_patched(): + # This provider limitation is not fixed by static discovery. Keep visible. + wire = request_reorganization_by_ark( + { + "model": "openai/offline", + "input": [], + "previous_response_id": "old", + "extra_body": {"parallel_tool_calls": False}, + } + ) + assert wire["extra_body"]["parallel_tool_calls"] is False diff --git a/tests/tools/cua_mcp/test_serial_guidance.py b/tests/tools/cua_mcp/test_serial_guidance.py new file mode 100644 index 000000000..630ef923a --- /dev/null +++ b/tests/tools/cua_mcp/test_serial_guidance.py @@ -0,0 +1,203 @@ +"""Multi-call turns are rejected as recoverable CUA results, never dispatched.""" + +from pathlib import Path +from unittest.mock import AsyncMock + +import pytest +from google.adk.events import Event +from google.genai import types +from test_toolset import context, toolset + + +@pytest.mark.asyncio +@pytest.mark.parametrize("batch_visible", [True, False]) +@pytest.mark.parametrize("source", ["counter", "event", "mixed_run_code"]) +async def test_multi_call_rejection_is_actionable_without_dispatch( + batch_visible, source +): + cua = toolset() + if not batch_visible: + cua.tool_filter = ["cua_click"] + ctx, readonly = await context("reject") + try: + tools = await cua.get_tools(readonly) + session = await cua._manager.create_session() + session.call_tool = AsyncMock(wraps=session.call_tool) + if source == "counter": + ctx._veadk_parallel_tool_call_count = 2 + else: + ctx.session.events.append( + Event( + author="agent", + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + name="run_code" + if source == "mixed_run_code" and i == 1 + else "cua_click", + args={}, + id=str(i), + ) + ) + for i in range(2) + ], + ), + ) + ) + click = next(t for t in tools if t.name == "cua_click") + result = await click.run_async(args={}, tool_context=ctx) + assert result["isError"] is True + assert result["structuredContent"]["code"] == "CUA_MULTIPLE_TOOL_CALLS_REJECTED" + assert result["structuredContent"]["executed"] is False + message = result["content"][0]["text"] + assert ("cua_actions_batch" in message) is batch_visible + assert "one tool call" in message + if batch_visible: + assert "action_type" in message and "run_code" in message + assert session.call_tool.await_count == 0 + assert ctx.state.get(cua.binding_state_key(ctx)) is None + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_real_runner_recovers_from_multi_call_to_one_batch(): + from google.adk.agents import Agent + from google.adk.models import LlmResponse + from google.adk.models.base_llm import BaseLlm + from google.adk.runners import InMemoryRunner + + class Model(BaseLlm): + step: int = 0 + + async def generate_content_async(self, request, stream=False): + if self.step == 0: + parts = [ + types.Part( + function_call=types.FunctionCall( + id=f"bad-{i}", name="cua_click", args={} + ) + ) + for i in range(2) + ] + elif self.step == 1: + responses = [ + p.function_response + for c in request.contents + for p in c.parts or [] + if p.function_response + ] + assert len(responses) == 2 + assert all(r.response["isError"] for r in responses) + assert all("cua_actions_batch" in str(r.response) for r in responses) + parts = [ + types.Part( + function_call=types.FunctionCall( + id="batch", + name="cua_actions_batch", + args={ + "actions": [ + {"action_type": "CLICK"}, + {"action_type": "TYPING", "text": "fixture"}, + ] + }, + ) + ) + ] + else: + response = [ + p.function_response + for c in request.contents + for p in c.parts or [] + if p.function_response + ][-1] + assert ( + response.name == "cua_actions_batch" + and not response.response.get("isError") + ) + assert response.response["structuredContent"]["data"]["performed"] == [ + "CLICK", + "TYPING", + ] + parts = [types.Part(text="done")] + self.step += 1 + yield LlmResponse(content=types.Content(role="model", parts=parts)) + + cua = toolset() + model = Model(model="local-scripted") + runner = InMemoryRunner( + agent=Agent(name="agent", model=model, tools=[cua]), app_name="app" + ) + await runner.session_service.create_session( + app_name="app", user_id="user", session_id="recovery" + ) + try: + session = await cua._manager.create_session() + session.call_tool = AsyncMock(wraps=session.call_tool) + events = [ + event + async for event in runner.run_async( + user_id="user", + session_id="recovery", + new_message=types.Content(role="user", parts=[types.Part(text="test")]), + ) + ] + assert model.step == 3 + assert any(e.is_final_response() for e in events) + assert session.call_tool.await_count == 1 + assert session.call_tool.call_args.args[0] == "cua_actions_batch" + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_disabled_template_never_recommends_batch(): + import sys + from mcp import StdioServerParameters + from veadk.tools.builtin_tools.computer_use import create_computer_use_toolset + + cua = create_computer_use_toolset( + tool_id="t-example", + connection_params=StdioServerParameters( + command=sys.executable, + args=[str(Path(__file__).with_name("stdio_fixture.py")), "--disabled"], + ), + ) + ctx, readonly = await context("off") + try: + tools = await cua.get_tools(readonly) + assert [t.name for t in tools] == ["ensure_sandbox"] + ctx._veadk_parallel_tool_call_count = 2 + result = await tools[0].run_async(args={}, tool_context=ctx) + assert result["isError"] and "cua_actions_batch" not in str(result) + assert ctx.state.get(cua.binding_state_key(ctx)) is None + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_mcp_descriptions_reach_model_declarations(): + cua = toolset() + _, readonly = await context("schema") + try: + tools = await cua.get_tools(readonly) + assert len(tools) == 55 + for tool in tools: + declaration = tool._get_declaration().model_dump( + mode="json", exclude_none=True + ) + assert len(declaration["description"]) >= 80 + schema = declaration.get("parameters_json_schema") or declaration.get( + "parameters" + ) + assert schema and schema.get("description") + for name, field in schema.get("properties", {}).items(): + assert field.get("description"), (tool.name, name) + batch = next(t for t in tools if t.name == "cua_actions_batch") + assert ( + "one tool call" in batch.description and "no rollback" in batch.description + ) + finally: + await cua.close() diff --git a/tests/tools/cua_mcp/test_toolset.py b/tests/tools/cua_mcp/test_toolset.py new file mode 100644 index 000000000..b37622ebc --- /dev/null +++ b/tests/tools/cua_mcp/test_toolset.py @@ -0,0 +1,642 @@ +import sys +from pathlib import Path + +import pytest +from google.adk.agents import Agent +from google.adk.agents.invocation_context import InvocationContext +from google.adk.agents.readonly_context import ReadonlyContext +from google.adk.sessions import InMemorySessionService +from google.adk.tools import ToolContext +from mcp import StdioServerParameters + +from veadk.tools.builtin_tools.computer_use import create_computer_use_toolset +from veadk.tools.builtin_tools.run_code import _tool_user_session_id + + +async def context(session_id): + service = InMemorySessionService() + session = await service.create_session( + app_name="app", user_id="user", session_id=session_id + ) + ctx = InvocationContext( + session_service=service, + invocation_id="turn", + agent=Agent(name="agent"), + session=session, + ) + return ToolContext(ctx), ReadonlyContext(ctx) + + +def toolset(): + return create_computer_use_toolset( + connection_params=StdioServerParameters( + command=sys.executable, + args=[str(Path(__file__).with_name("stdio_fixture.py"))], + ), + tool_id="t-example", + ) + + +@pytest.mark.asyncio +async def test_real_context_binding_refresh_same_invocation_and_second_session(): + cua = toolset() + ctx, readonly = await context("one") + other, other_readonly = await context("two") + try: + tools = await cua.get_tools_with_prefix(readonly) + assert "cua_screenshot" in [t.name for t in tools] + result = await tools[0].run_async(args={}, tool_context=ctx) + assert result["structuredContent"]["user_session_id"] == _tool_user_session_id( + ctx + ) + assert ctx.state[cua.binding_state_key(ctx)]["session_id"] == "s-1" + # ADK normally caches toolsets per invocation; CUA must opt out. + assert "cua_screenshot" in [ + t.name for t in await cua.get_tools_with_prefix(readonly) + ] + assert "cua_screenshot" in [ + t.name for t in await cua.get_tools_with_prefix(other_readonly) + ] + assert other.state.get(cua.binding_state_key(other)) is None + # Simulate a new Agent turn after the state delta has been persisted. + ctx.session.state[cua.binding_state_key(ctx)] = ctx.state[ + cua.binding_state_key(ctx) + ] + await cua.close() + restored = toolset() + try: + # The restarted fixture has no platform instances: do not replace the saved binding. + with pytest.raises(RuntimeError, match="discovery failed"): + await restored.get_tools(readonly) + assert ctx.state[cua.binding_state_key(ctx)]["session_id"] == "s-1" + finally: + await restored.close() + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_parallel_calls_rejected_and_model_cannot_override_context(): + cua = toolset() + ctx, readonly = await context("one") + try: + tool = (await cua.get_tools(readonly))[0] + ctx._veadk_parallel_tool_call_count = 2 + rejected = await tool.run_async(args={}, tool_context=ctx) + assert rejected["isError"] + assert rejected["structuredContent"]["executed"] is False + ctx._veadk_parallel_tool_call_count = 1 + result = await tool.run_async( + args={"user_session_id": "spoofed"}, tool_context=ctx + ) + assert result["isError"] + assert ctx.state.get(cua.binding_state_key(ctx)) is None + finally: + await cua.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("runner_kind", ["google_adk", "veadk"]) +async def test_runner_receives_actions_before_first_model_call_and_image(runner_kind): + from google.adk.models.base_llm import BaseLlm + from google.adk.models.llm_response import LlmResponse + from google.adk.runners import InMemoryRunner + from google.genai import types + + class ScriptedModel(BaseLlm): + step: int = 0 + + async def generate_content_async(self, llm_request, stream=False): + names = set(llm_request.tools_dict) + if self.step == 0: + assert "cua_screenshot" in names + call = types.FunctionCall(name="cua_screenshot", args={}) + else: + images = [ + part + for content in llm_request.contents + for part in content.parts or [] + if part.function_response and part.function_response.parts + ] + assert images, ( + "Screenshot must be media, not a base64 string for the model" + ) + yield LlmResponse( + content=types.Content(role="model", parts=[types.Part(text="done")]) + ) + return + self.step += 1 + yield LlmResponse( + content=types.Content( + role="model", parts=[types.Part(function_call=call)] + ) + ) + + cua = toolset() + model = ScriptedModel(model="local-test") + if runner_kind == "veadk": + from veadk import Agent as VeAgent, Runner + + agent = VeAgent( + name="agent", model=model, model_api_key="offline-fixture", tools=[cua] + ) + runner = Runner( + agent=agent, app_name="app", session_service=InMemorySessionService() + ) + else: + agent = Agent(name="agent", model=model, tools=[cua]) + runner = InMemoryRunner(agent=agent, app_name="app") + await runner.session_service.create_session( + app_name="app", user_id="user", session_id="one" + ) + try: + events = [ + event + async for event in runner.run_async( + user_id="user", + session_id="one", + new_message=types.Content(role="user", parts=[types.Part(text="test")]), + ) + ] + assert any(event.is_final_response() for event in events) + saved = await runner.session_service.get_session( + app_name="app", user_id="user", session_id="one" + ) + assert ( + saved.state[ + next(key for key in saved.state if key.startswith(cua.state_key)) + ]["session_id"] + == "s-1" + ) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_subagents_in_same_session_keep_distinct_bindings(): + cua = toolset() + first, readonly = await context("shared") + try: + await (await cua.get_tools(readonly))[0].run_async(args={}, tool_context=first) + first.session.state.update(first.state.to_dict()) + second = ToolContext( + InvocationContext( + session_service=first._invocation_context.session_service, + invocation_id="second", + agent=Agent(name="other_agent"), + session=first.session, + ) + ) + result = await (await cua.get_tools(second))[0].run_async( + args={}, tool_context=second + ) + assert result["structuredContent"]["session_id"] == "s-2" + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_history_parallel_run_code_is_not_silently_attached(): + from google.adk.events import Event + from google.genai import types + + cua = toolset() + ctx, readonly = await context("one") + ctx.session.events.append( + Event( + author="agent", + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall(name="run_code", args={}) + ), + types.Part( + function_call=types.FunctionCall(name="run_code", args={}) + ), + ], + ), + ) + ) + try: + with pytest.raises(ValueError, match="Parallel run_code history"): + await cua.get_tools(readonly) + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_allowlist_retains_bootstrap_and_closes_child(): + cua = toolset() + cua.tool_filter = ["cua_screenshot"] + ctx, readonly = await context("one") + tool = (await cua.get_tools(readonly))[0] + await tool.run_async(args={}, tool_context=ctx) + assert {t.name for t in await cua.get_tools(readonly)} == { + "ensure_sandbox", + "cua_screenshot", + } + await cua.close() + assert not cua._manager._sessions + + +@pytest.mark.asyncio +async def test_tool_cancellation_propagates_and_mutations_not_retried(): + import asyncio + from unittest.mock import AsyncMock + from mcp import types + from veadk.tools.builtin_tools.computer_use import _ComputerTool + + cua = toolset() + ctx, _ = await context("one") + session = AsyncMock() + session.call_tool.side_effect = ConnectionError("SENSITIVE_URL") + cua._manager.create_session = AsyncMock(return_value=session) + tool = _ComputerTool( + types.Tool(name="cua_click", inputSchema={"type": "object"}), cua + ) + result = await tool.run_async(args={}, tool_context=ctx) + assert result["isError"] + assert "SENSITIVE" not in str(result) + assert session.call_tool.await_count == 1 + session.call_tool.side_effect = asyncio.CancelledError() + with pytest.raises(asyncio.CancelledError): + await tool.run_async(args={}, tool_context=ctx) + assert session.call_tool.await_count == 2 + await cua.close() + + +@pytest.mark.asyncio +async def test_saved_binding_reconnects_after_mcp_process_restart(): + ctx, readonly = await context("one") + first = toolset() + try: + await (await first.get_tools(readonly))[0].run_async(args={}, tool_context=ctx) + ctx.session.state.update(ctx.state.to_dict()) + finally: + await first.close() + second = create_computer_use_toolset( + connection_params=StdioServerParameters( + command=sys.executable, + args=[str(Path(__file__).with_name("stdio_fixture.py")), "--existing"], + ), + tool_id="t-example", + ) + try: + assert "cua_screenshot" in { + tool.name for tool in await second.get_tools(readonly) + } + result = await (await second.get_tools(readonly))[0].run_async( + args={}, tool_context=ctx + ) + assert result["structuredContent"]["session_id"] == "s-1" + finally: + await second.close() + + +@pytest.mark.asyncio +async def test_unchanged_run_code_result_not_used_for_binding(monkeypatch): + import importlib + + module = importlib.import_module("veadk.tools.builtin_tools.run_code") + ctx, readonly = await context("one") + invoked = [] + + def invoke(**kwargs): + invoked.append(kwargs) + return {"Result": {"Result": "only-program-output"}} + + monkeypatch.setattr(module, "resolve_agentkit_tool_id", lambda *args: "t-example") + monkeypatch.setattr( + module, + "get_agentkit_endpoint_config", + lambda: ("agentkit", "cn-beijing", "unused.example", "https"), + ) + monkeypatch.setattr(module, "invoke_agentkit_run_code", invoke) + assert module.run_code("print(1)", "python3", ctx) == "only-program-output" + cua = create_computer_use_toolset( + connection_params=StdioServerParameters( + command=sys.executable, + args=[str(Path(__file__).with_name("stdio_fixture.py")), "--existing"], + ), + tool_id="t-example", + ) + try: + result = await (await cua.get_tools(readonly))[0].run_async( + args={}, tool_context=ctx + ) + assert ( + result["structuredContent"]["user_session_id"] + == invoked[0]["tool_user_session_id"] + ) + assert result["structuredContent"]["session_id"] == "s-1" + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_runner_stops_before_model_when_mcp_discovery_fails(caplog): + from unittest.mock import AsyncMock + from google.adk.models.base_llm import BaseLlm + from google.adk.models.llm_response import LlmResponse + from google.adk.runners import InMemoryRunner + from google.genai import types + + calls = [] + + class NoCallModel(BaseLlm): + async def generate_content_async(self, llm_request, stream=False): + calls.append(True) + yield LlmResponse( + content=types.Content( + role="model", parts=[types.Part(text="must not run")] + ) + ) + + cua = toolset() + cua._manager.create_session = AsyncMock( + side_effect=ConnectionError("private endpoint") + ) + runner = InMemoryRunner( + agent=Agent(name="agent", model=NoCallModel(model="offline"), tools=[cua]), + app_name="app", + ) + await runner.session_service.create_session( + app_name="app", user_id="user", session_id="one" + ) + try: + with pytest.raises(RuntimeError, match="discovery failed"): + _ = [ + event + async for event in runner.run_async( + user_id="user", + session_id="one", + new_message=types.Content( + role="user", parts=[types.Part(text="test")] + ), + ) + ] + assert not calls + assert "install a compatible template-fixed-v1 MCP wheel" in caplog.text + assert "current local candidate" not in caplog.text + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_ark_accepts_default_cache_without_mutating_agent(): + from google.adk.models.llm_request import LlmRequest + from veadk.models.ark_llm import ArkLlm + + cua = toolset() + ctx, _ = await context("one") + agent = ctx._invocation_context.agent + agent.model = ArkLlm(model="openai/offline", enable_responses_cache=True) + try: + await cua.process_llm_request(tool_context=ctx, llm_request=LlmRequest()) + assert agent.model.enable_responses_cache is True + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_ark_receives_screenshot_media_from_function_response_parts(): + from google.adk.models.llm_request import LlmRequest + from google.genai import types + from veadk.models.ark_llm import ArkLlm, _get_responses_inputs + + cua = toolset() + ctx, _ = await context("one") + ctx._invocation_context.agent.model = ArkLlm( + model="openai/offline", enable_responses_cache=False + ) + req = LlmRequest( + contents=[ + types.Content( + role="user", + parts=[ + types.Part( + function_response=types.FunctionResponse( + name="cua_screenshot", + id="shot-1", + response={"isError": False}, + parts=[ + types.FunctionResponsePart( + inline_data=types.FunctionResponseBlob( + mime_type="image/png", data=b"png" + ) + ) + ], + ) + ) + ], + ) + ] + ) + try: + await cua.process_llm_request(tool_context=ctx, llm_request=req) + inputs = _get_responses_inputs(req)[1] + assert any( + item.get("type") == "message" + and any(p.get("type") == "input_image" for p in item.get("content", [])) + for item in inputs + ) + await cua.process_llm_request(tool_context=ctx, llm_request=req) + assert len(req.contents) == 2 + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_discovery_and_pure_chat_never_call_ensure_or_create(): + from unittest.mock import AsyncMock + from google.adk.models.base_llm import BaseLlm + from google.adk.models import LlmResponse + from google.adk.runners import InMemoryRunner + from google.genai import types + + class Chat(BaseLlm): + async def generate_content_async(self, llm_request, stream=False): + assert "cua_screenshot" in llm_request.tools_dict + yield LlmResponse( + content=types.Content(role="model", parts=[types.Part(text="hello")]) + ) + + cua = toolset() + session = await cua._manager.create_session() + session.call_tool = AsyncMock(wraps=session.call_tool) + runner = InMemoryRunner( + agent=Agent(name="agent", model=Chat(model="offline"), tools=[cua]), + app_name="app", + ) + await runner.session_service.create_session( + app_name="app", user_id="user", session_id="one" + ) + try: + for _ in range(2): + events = [ + e + async for e in runner.run_async( + user_id="user", + session_id="one", + new_message=types.Content( + role="user", parts=[types.Part(text="hello")] + ), + ) + ] + assert events + assert session.call_tool.await_count == 0 + saved = await runner.session_service.get_session( + app_name="app", user_id="user", session_id="one" + ) + assert not any(k.startswith(cua.state_key) for k in saved.state) + finally: + await runner.close() + + +@pytest.mark.asyncio +async def test_concurrent_readonly_discovery_does_not_ensure(): + import asyncio + from unittest.mock import AsyncMock + + cua = toolset() + ctx, ro = await context("one") + try: + session = await cua._manager.create_session() + session.call_tool = AsyncMock(wraps=session.call_tool) + results = await asyncio.gather(*(cua.get_tools(ro) for _ in range(5))) + assert all("cua_screenshot" in {t.name for t in r} for r in results) + assert session.call_tool.await_count == 0 + assert ctx.state.get(cua.binding_state_key(ctx)) is None + tool = next(t for t in results[0] if t.name == "cua_screenshot") + await tool.run_async(args={}, tool_context=ctx) + assert session.call_tool.await_count == 1 + assert ctx.state[cua.binding_state_key(ctx)]["session_id"] == "s-1" + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_disabled_template_list_leaves_session_unbound(): + cua = create_computer_use_toolset( + tool_id="t-example", + connection_params=StdioServerParameters( + command=sys.executable, + args=[str(Path(__file__).with_name("stdio_fixture.py")), "--disabled"], + ), + ) + ctx, ro = await context("one") + try: + assert [t.name for t in await cua.get_tools(ro)] == ["ensure_sandbox"] + assert ctx.state.get(cua.binding_state_key(ctx)) is None + finally: + await cua.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "failure", ["permission", "cancel", "old_server", "invalid_binding"] +) +async def test_discovery_fails_closed_without_creating(failure): + import asyncio + from unittest.mock import AsyncMock + from mcp import types + from google.adk.models import LlmRequest + + cua = toolset() + ctx, ro = await context("one") + session = AsyncMock() + session.list_tools.return_value = types.ListToolsResult( + tools=[], + _meta={"veadk/cua": {"discovery_mode": "template-fixed-v1", "binding": None}}, + ) + if failure == "permission": + session.list_tools.side_effect = RuntimeError("SENSITIVE") + if failure == "cancel": + session.list_tools.side_effect = asyncio.CancelledError() + if failure == "old_server": + session.list_tools.return_value.meta = None + if failure == "invalid_binding": + session.list_tools.return_value.meta["veadk/cua"]["binding"] = { + "session_id": "s-other", + "user_session_id": "other", + } + cua._manager.create_session = AsyncMock(return_value=session) + try: + with pytest.raises( + asyncio.CancelledError if failure == "cancel" else RuntimeError + ): + await cua.get_tools(ro) + assert session.call_tool.await_count == 0 + with pytest.raises(RuntimeError, match="model invocation aborted"): + await cua.process_llm_request(tool_context=ctx, llm_request=LlmRequest()) + finally: + await cua.close() + + +@pytest.mark.asyncio +async def test_readonly_attach_persists_existing_instance_in_database(tmp_path): + from google.adk.sessions import DatabaseSessionService + + cua = create_computer_use_toolset( + tool_id="t-example", + connection_params=StdioServerParameters( + command=sys.executable, + args=[str(Path(__file__).with_name("stdio_fixture.py")), "--existing"], + ), + ) + service = DatabaseSessionService( + db_url=f"sqlite+aiosqlite:///{tmp_path / 'sessions.db'}" + ) + session = await service.create_session( + app_name="app", user_id="user", session_id="one" + ) + ctx = ReadonlyContext( + InvocationContext( + session_service=service, + invocation_id="turn", + agent=Agent(name="agent"), + session=session, + ) + ) + try: + await cua.get_tools(ctx) + saved = await service.get_session( + app_name="app", user_id="user", session_id="one" + ) + assert saved.state[cua.binding_state_key(ctx)]["session_id"] == "s-1" + finally: + await cua.close() + await service.close() + + +@pytest.mark.asyncio +async def test_action_failure_preserves_binding_and_hides_routing_from_model(): + from unittest.mock import AsyncMock + from mcp import types + from veadk.tools.builtin_tools.computer_use import _ComputerTool + + cua = toolset() + ctx, _ = await context("one") + session = AsyncMock() + session.call_tool.return_value = types.CallToolResult( + isError=True, + content=[types.TextContent(type="text", text="COMPUTER_DISABLED")], + _meta={ + "veadk/cua": { + "binding": {"user_session_id": "agent_user_one", "session_id": "s-1"} + } + }, + ) + cua._manager.create_session = AsyncMock(return_value=session) + try: + tool = _ComputerTool( + types.Tool(name="cua_screenshot", inputSchema={"type": "object"}), cua + ) + result = await tool.run_async(args={}, tool_context=ctx) + assert result["isError"] + assert ctx.state[cua.binding_state_key(ctx)]["session_id"] == "s-1" + assert "meta" not in result and "_meta" not in result + assert session.call_tool.await_count == 1 + finally: + await cua.close() diff --git a/veadk/tools/builtin_tools/computer_use.py b/veadk/tools/builtin_tools/computer_use.py new file mode 100644 index 000000000..9999e905b --- /dev/null +++ b/veadk/tools/builtin_tools/computer_use.py @@ -0,0 +1,445 @@ +# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# SPDX-License-Identifier: Apache-2.0 +"""AgentKit CUA stdio Toolset. The independently installed MCP package owns provisioning.""" + +from __future__ import annotations + +import asyncio +import base64 +import hashlib +import json +import os +import re +import weakref +from datetime import timedelta + +from google.adk.tools.base_toolset import BaseToolset +from google.adk.events import Event, EventActions +from google.adk.tools.mcp_tool.mcp_session_manager import ( + MCPSessionManager, + StdioConnectionParams, +) +from google.adk.tools.mcp_tool.mcp_tool import McpTool +from mcp import StdioServerParameters, types +from google.genai import types as genai_types + +_META_KEY = "veadk/cua" +_FORWARD_ENV = ( + "VOLCENGINE_ACCESS_KEY", + "VOLCENGINE_SECRET_KEY", + "VOLCENGINE_SESSION_TOKEN", + "AGENTKIT_TOOL_TTL", + "CUA_ALLOW_CREATE", + "UV_CACHE_DIR", + "UV_INDEX_URL", + "UV_OFFLINE", +) + + +class _ComputerTool(McpTool): + def __init__(self, definition, owner): + super().__init__(mcp_tool=definition, mcp_session_manager=owner._manager) + self._owner = owner + self._batch_available = False + + async def run_async(self, *, args, tool_context): + # Deliberately bypass McpTool's transport retry: a desktop mutation may + # have completed before a broken connection/timeout becomes visible. + if getattr(tool_context, "_veadk_parallel_tool_call_count", 1) > 1: + return self._serial_rejection() + for event in reversed(tool_context.session.events): + content = event.content + if content and content.role == "model": + calls = [ + p.function_call for p in content.parts or [] if p.function_call + ] + if len(calls) > 1: + return self._serial_rejection() + break + meta = self._owner._metadata(tool_context) + try: + session = await asyncio.wait_for(self._owner._manager.create_session(), 30) + response = await session.call_tool( + self.name, + args, + meta=meta, + read_timeout_seconds=timedelta(seconds=self._owner.timeout), + ) + except Exception: + return { + "isError": True, + "content": [ + { + "type": "text", + "text": "CUA MCP transport failed; outcome may be unknown. Operation was not retried.", + } + ], + } + binding = self._owner._response_binding(response, meta) + if binding is not None: + # Save even on action errors: the instance may already have been + # created. State deltas are persisted by the normal tool event. + tool_context.state[self._owner.binding_state_key(tool_context)] = binding + result = response.model_dump(exclude_none=True, mode="json") + result.pop("meta", None) + result.pop("_meta", None) + # ADK extracts typed media parts; raw MCP base64 dictionaries remain text. + for index, block in enumerate(response.content): + if isinstance(block, types.ImageContent): + result["content"][index] = genai_types.Part.from_bytes( + data=base64.b64decode(block.data, validate=True), + mime_type=block.mimeType, + ) + return result + + def _serial_rejection(self): + message = ( + "This CUA call was not executed: multiple tool calls in one model response " + "are rejected. Submit one tool call at a time and wait for its result. " + ) + if self._batch_available: + message += ( + "For a known ordered sequence of supported desktop input actions, submit " + "one cua_actions_batch call with actions containing action_type and the " + "matching parameters (not MCP tool names). Do not put run_code, browser " + "lifecycle, app/process, recording or script tools inside the batch. " + "If the next step depends on a screenshot/result, use separate turns. " + ) + message += ( + "This rejection covers only this CUA call, not other tools in the same " + "response; they may have executed. Inspect their results before proceeding." + ) + return { + "isError": True, + "content": [{"type": "text", "text": message}], + "structuredContent": { + "code": "CUA_MULTIPLE_TOOL_CALLS_REJECTED", + "executed": False, + "scope": "this_cua_call", + "message": message, + }, + } + + +class ComputerUseToolset(BaseToolset): + """Per-request context, no global current session and no model routing arguments. + + The caller owns the returned Toolset and must await ``close()`` on Runtime + shutdown. Use one Toolset per fixed ToolId/region/credential namespace. + Successful bindings are persisted through ADK ToolContext state deltas. + """ + + def __init__( + self, + *, + connection_params: StdioServerParameters, + tool_id: str, + region: str = "cn-beijing", + tool_filter: list[str] | None = None, + timeout: float = 120, + ): + super().__init__(tool_filter=tool_filter) + if not re.fullmatch(r"t-[A-Za-z0-9_-]+", tool_id): + raise ValueError("A valid AgentKit ToolId is required") + if not 1 <= timeout <= 300: + raise ValueError("CUA request timeout must be between 1 and 300 seconds") + self._use_invocation_cache = False + self._discovery_failures = {} + self._prepare_locks = weakref.WeakValueDictionary() + self._catalogs = {} + self.tool_id, self.region, self.timeout = tool_id, region, timeout + suffix = hashlib.sha256(f"{region}:{tool_id}".encode()).hexdigest()[:16] + self.state_key = f"veadk_cua_binding_{suffix}" + env = dict(connection_params.env or {}) + configured = env.get("AGENTKIT_TOOL_ID_CUA") or env.get("AGENTKIT_TOOL_ID") + if configured and configured != tool_id: + raise ValueError("MCP subprocess ToolId must match Toolset ToolId") + if env.get("AGENTKIT_TOOL_REGION", region) != region: + raise ValueError("MCP subprocess region must match Toolset region") + env.update(AGENTKIT_TOOL_ID_CUA=tool_id, AGENTKIT_TOOL_REGION=region) + params = connection_params.model_copy(update={"env": env}) + self._manager = MCPSessionManager( + StdioConnectionParams(server_params=params, timeout=30) + ) + + def binding_state_key(self, ctx): + """Use a distinct state entry for each agent in a shared ADK session.""" + suffix = hashlib.sha256(ctx.agent_name.encode()).hexdigest()[:16] + return f"{self.state_key}_{suffix}" + + def _metadata(self, ctx): + if ctx is None: + raise ValueError( + "ComputerUseToolset requires a real ADK invocation context" + ) + for event in ctx.session.events: + content = event.content + calls = ( + [p.function_call for p in content.parts or [] if p.function_call] + if content + else [] + ) + if sum(call.name == "run_code" for call in calls) > 1: + raise ValueError( + "Parallel run_code history cannot identify one desktop; start a serial session" + ) + key = f"{ctx.agent_name}_{ctx.user_id}_{ctx.session.id}" + if not re.fullmatch(r"[A-Za-z0-9_-]{1,200}", key): + raise ValueError( + "UserSessionId must match AgentKit limits; no automatic key rewriting" + ) + binding = ctx.state.get(self.binding_state_key(ctx)) + route = {"user_session_id": key} + if binding is not None: + if not isinstance(binding, dict) or binding.get("user_session_id") != key: + raise ValueError( + "Stored CUA binding does not match the actual ADK context" + ) + sid = binding.get("session_id") + if not isinstance(sid, str) or not re.fullmatch( + r"[A-Za-z0-9_-]{1,200}", sid + ): + raise ValueError("Invalid stored CUA SessionId") + route["session_id"] = sid + return {_META_KEY: route} + + async def get_tools(self, readonly_context=None): + invocation = getattr(readonly_context, "_invocation_context", None) + try: + meta = self._metadata(readonly_context) + except Exception: + if invocation is not None: + self._remember_discovery_failure(invocation) + raise + try: + lock_key = ( + readonly_context.session.app_name, + meta[_META_KEY]["user_session_id"], + ) + lock = self._prepare_locks.setdefault(lock_key, asyncio.Lock()) + async with lock: + meta = self._metadata(readonly_context) + cached = self._catalogs.get(id(invocation)) + if cached and cached[0]() is invocation: + if ( + cached[1][_META_KEY]["user_session_id"] + != meta[_META_KEY]["user_session_id"] + ): + raise ValueError("CUA session changed during invocation") + # A successful action may have added the first binding. + result = cached[2] + else: + if id(invocation) in self._discovery_failures: + raise RuntimeError( + "CUA preparation previously failed in this invocation" + ) + session = await asyncio.wait_for(self._manager.create_session(), 30) + # Only tools/list: never ensure or create in discovery. + result = await asyncio.wait_for( + session.list_tools( + params=types.PaginatedRequestParams(_meta=meta) + ), + self.timeout, + ) + contract = (result.meta or {}).get(_META_KEY, {}) + if contract.get("discovery_mode") != "template-fixed-v1": + raise RuntimeError( + "CUA MCP must support template-fixed-v1 discovery; refresh the wheel" + ) + binding = self._response_binding(result, meta) + if ( + binding is not None + and readonly_context.state.get( + self.binding_state_key(readonly_context) + ) + != binding + ): + await invocation.session_service.append_event( + session=invocation.session, + event=Event( + invocation_id=invocation.invocation_id, + author=readonly_context.agent_name, + branch=invocation.branch, + actions=EventActions( + state_delta={ + self.binding_state_key( + readonly_context + ): binding + } + ), + ), + ) + if ( + readonly_context.state.get( + self.binding_state_key(readonly_context) + ) + != binding + ): + raise RuntimeError("CUA binding was not persisted") + meta = self._metadata(readonly_context) + if result.nextCursor: + raise RuntimeError("Unexpected pagination in CUA catalog") + identity = id(invocation) + self._catalogs[identity] = ( + weakref.ref( + invocation, lambda ref: self._catalogs.pop(identity, None) + ), + meta, + result, + ) + except asyncio.CancelledError: + if invocation is not None: + self._remember_discovery_failure(invocation) + raise + except Exception: + if invocation is not None: + self._remember_discovery_failure(invocation) + raise RuntimeError( + "CUA tool discovery failed; discovery does not create instances. " + "Check template permissions and install a compatible template-fixed-v1 " + "MCP wheel." + ) from None + if result.nextCursor: + if invocation is not None: + self._remember_discovery_failure(invocation) + raise RuntimeError("Unexpected pagination in the CUA tool catalog") + if invocation is not None: + self._discovery_failures.pop(id(invocation), None) + tools = [_ComputerTool(definition, self) for definition in result.tools] + # Always keep the bootstrap entry even when a desktop allowlist is supplied. + selected = [ + tool + for tool in tools + if tool.name == "ensure_sandbox" + or self._is_tool_selected(tool, readonly_context) + ] + batch_available = any(t.name == "cua_actions_batch" for t in selected) + for tool in selected: + tool._batch_available = batch_available + # Fixed means fixed across turns and process restarts, not just inside + # one invocation. Persist a digest, never credentials or model state. + key = "veadk_cua_catalog_" + self.binding_state_key(readonly_context) + signature = hashlib.sha256( + json.dumps( + [ + t._get_declaration().model_dump(mode="json", exclude_none=True) + for t in selected + ], + sort_keys=True, + separators=(",", ":"), + ).encode() + ).hexdigest() + old = readonly_context.state.get(key) + if old is not None and old != signature: + self._remember_discovery_failure(invocation) + raise RuntimeError("CUA catalog changed; start a new conversation/session") + if old is None: + await invocation.session_service.append_event( + session=invocation.session, + event=Event( + invocation_id=invocation.invocation_id, + author=readonly_context.agent_name, + branch=invocation.branch, + actions=EventActions(state_delta={key: signature}), + ), + ) + return selected + + @staticmethod + def _response_binding(response, meta): + contract = (response.meta or {}).get(_META_KEY, {}) + binding = contract.get("binding") + if binding is None: + return None + expected = meta[_META_KEY] + if ( + not isinstance(binding, dict) + or set(binding) != {"user_session_id", "session_id"} + or binding.get("user_session_id") != expected["user_session_id"] + or not isinstance(binding.get("session_id"), str) + or not re.fullmatch(r"[A-Za-z0-9_-]{1,200}", binding["session_id"]) + or ( + expected.get("session_id") + and binding["session_id"] != expected["session_id"] + ) + ): + raise ValueError("CUA MCP returned an inconsistent sandbox binding") + return dict(binding) + + def _remember_discovery_failure(self, invocation): + identity = id(invocation) + self._discovery_failures[identity] = weakref.ref( + invocation, lambda ref: self._discovery_failures.pop(identity, None) + ) + + async def process_llm_request(self, *, tool_context, llm_request): + # ADK catches get_tools failures and substitutes []; stop the model rather + # than silently continuing without the requested CUA capability. + if id(tool_context._invocation_context) in self._discovery_failures: + raise RuntimeError("CUA tool discovery failed; model invocation aborted") + # Adapt only Ark requests; other models retain their native tool media. + from veadk.models.ark_llm import ArkLlm + + if isinstance(tool_context._invocation_context.agent.canonical_model, ArkLlm): + from veadk.tools.computer_use.media import lift_screenshots + + # Adapt only images already in the request; history and cache + # continuation remain owned by ADK and the original Ark adapter. + lift_screenshots(llm_request) + + async def close(self): + await self._manager.close() + self._prepare_locks.clear() + self._catalogs.clear() + self._discovery_failures.clear() + + +def create_computer_use_toolset( + *, + connection_params: StdioServerParameters | None = None, + tool_id: str | None = None, + region: str | None = None, + tool_filter: list[str] | None = None, +) -> ComputerUseToolset: + """Launch the configured CUA MCP wheel with uvx on first use. + + The independent server is maintained in the cua-sandbox repository. When + using the default uvx launcher, set ``CUA_MCP_WHEEL_URL`` to a trusted, + explicitly versioned wheel URL in deployment configuration. Alternatively, + pass ``connection_params`` for a preinstalled/offline command. + Runtime must provide uvx; import does not download or create cloud resources. + Discovery reads template/instance capabilities without creating a sandbox. + The first valid desktop action (or explicit ensure) creates/reuses an instance; + pure chat does not create a billable sandbox through this Toolset. + """ + tool_id = ( + tool_id + or os.getenv("AGENTKIT_TOOL_ID_CUA") + or os.getenv("AGENTKIT_TOOL_ID", "") + ) + region = region or os.getenv("AGENTKIT_TOOL_REGION", "cn-beijing") + if connection_params is None: + wheel_url = os.getenv("CUA_MCP_WHEEL_URL", "").strip() + if not wheel_url: + raise ValueError( + "CUA_MCP_WHEEL_URL must be set to an explicitly versioned CUA MCP " + "wheel URL when using the default uvx launcher; alternatively, " + "pass connection_params for a preinstalled or local server." + ) + connection_params = StdioServerParameters( + command="uvx", + args=[ + "--from", + wheel_url, + "agentkit-cua-mcp", + "--transport", + "stdio", + ], + env={name: os.environ[name] for name in _FORWARD_ENV if name in os.environ}, + ) + return ComputerUseToolset( + connection_params=connection_params, + tool_id=tool_id, + region=region, + tool_filter=tool_filter, + ) diff --git a/veadk/tools/computer_use/__init__.py b/veadk/tools/computer_use/__init__.py new file mode 100644 index 000000000..62c3946b6 --- /dev/null +++ b/veadk/tools/computer_use/__init__.py @@ -0,0 +1 @@ +"""CUA-only request adapters; no generic model patches or sandbox server code.""" diff --git a/veadk/tools/computer_use/media.py b/veadk/tools/computer_use/media.py new file mode 100644 index 000000000..2b823588a --- /dev/null +++ b/veadk/tools/computer_use/media.py @@ -0,0 +1,36 @@ +"""Adapt CUA tool-result images without changing provider model classes.""" + +from google.genai import types + + +def lift_screenshots(request): + """Expose nested screenshot media to Ark without editing stored events.""" + expanded = [] + for content in request.contents: + images, parts = [], [] + for part in content.parts or []: + response = part.function_response + if ( + response + and response.name + and response.name.startswith("cua_") + and response.parts + ): + for media in response.parts: + blob = media.inline_data + if blob and blob.mime_type and blob.mime_type.startswith("image/"): + images.append( + types.Part.from_bytes( + data=blob.data, mime_type=blob.mime_type + ) + ) + part = part.model_copy( + update={ + "function_response": response.model_copy(update={"parts": None}) + } + ) + parts.append(part) + expanded.append(content.model_copy(update={"parts": parts})) + if images: + expanded.append(types.Content(role="user", parts=images)) + request.contents = expanded