From 203d9ec8534889722349da3b053b4d707767e8e0 Mon Sep 17 00:00:00 2001 From: "hanzhi.421" Date: Sun, 6 Sep 2026 19:57:30 +0800 Subject: [PATCH 1/3] feat(agents): add unified AgentKit remote sandbox agent --- tests/agents/test_remote_sandbox_agent.py | 615 ++++++++++++++++++ tests/agents/test_sandbox_session.py | 116 ++++ veadk/__init__.py | 8 +- veadk/agent.py | 9 +- veadk/agents/REMOTE_SANDBOX.md | 71 ++ veadk/agents/_sandbox_code.py | 254 ++++++++ veadk/agents/_sandbox_session.py | 415 ++++++++++++ veadk/agents/_sandbox_skill.py | 278 ++++++++ veadk/agents/agentkit_remote_sandbox_agent.py | 248 +++++++ veadk/tools/sandbox/codex_worker_client.py | 198 ++++++ 10 files changed, 2209 insertions(+), 3 deletions(-) create mode 100644 tests/agents/test_remote_sandbox_agent.py create mode 100644 tests/agents/test_sandbox_session.py create mode 100644 veadk/agents/REMOTE_SANDBOX.md create mode 100644 veadk/agents/_sandbox_code.py create mode 100644 veadk/agents/_sandbox_session.py create mode 100644 veadk/agents/_sandbox_skill.py create mode 100644 veadk/agents/agentkit_remote_sandbox_agent.py create mode 100644 veadk/tools/sandbox/codex_worker_client.py diff --git a/tests/agents/test_remote_sandbox_agent.py b/tests/agents/test_remote_sandbox_agent.py new file mode 100644 index 000000000..1353e0f4b --- /dev/null +++ b/tests/agents/test_remote_sandbox_agent.py @@ -0,0 +1,615 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Offline tests of native transfer, both wire protocols and type discovery.""" + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from aiohttp import web +from aiohttp.test_utils import TestServer +from google.adk.models.base_llm import BaseLlm +from google.adk.models.llm_response import LlmResponse +from google.adk.runners import Runner +from google.adk.sessions import InMemorySessionService +from google.genai import types +from pydantic import PrivateAttr + +from veadk import Agent +from veadk.agents.agentkit_remote_sandbox_agent import ( + AgentkitRemoteSandboxAgent, + SandboxAgentError, +) + + +class ParentModel(BaseLlm): + model: str = "offline-parent" + _calls: int = PrivateAttr(default=0) + + async def generate_content_async(self, llm_request, stream=False): + self._calls += 1 + yield LlmResponse( + content=types.Content( + role="model", + parts=[ + types.Part( + function_call=types.FunctionCall( + name="transfer_to_agent", args={"agent_name": "sandbox"} + ) + ) + ], + ) + ) + + +async def make_runner(child): + model = ParentModel() + root = Agent( + name="coordinator", model=model, model_api_key="fixture", sub_agents=[child] + ) + sessions = InMemorySessionService() + await sessions.create_session( + app_name="probe", user_id="user", session_id="session" + ) + return Runner(agent=root, app_name="probe", session_service=sessions), model + + +async def collect(runner, on_event=None, *, user_id="user"): + result = [] + async for event in runner.run_async( + user_id=user_id, + session_id="session", + new_message=types.Content( + role="user", + parts=[types.Part(text="Compute two plus three in the sandbox")], + ), + ): + result.append(event) + if on_event: + await on_event(event) + return result + + +class CodeFixture: + def __init__(self): + self.release = asyncio.Event() + self.starts = 0 + self.session_keys = [] + self.cancelled = False + self.fail = False + + async def handle(self, request): + path = request.path + assert request.query.get("route") == "fixture" + if path.endswith("/readyz"): + return web.json_response( + {"schemaVersion": 1, "capabilities": ["tool_events"]} + ) + if path.endswith("/sessions"): + self.session_keys.append(request.headers["Idempotency-Key"]) + return web.json_response({"status": "ready", "sessionId": "session-code"}) + if path.endswith("/cancel"): + self.cancelled = True + return web.json_response({"status": "interrupted"}) + if request.method == "POST": + self.starts += 1 + return web.json_response({"turnId": "turn-code"}) + if not path.endswith("/events"): + return web.json_response({"status": "running", "codexTurnId": "codex-1"}) + response = web.StreamResponse(headers={"Content-Type": "text/event-stream"}) + await response.prepare(request) + + async def send(seq, kind, payload): + data = { + "schemaVersion": 1, + "eventId": seq, + "sessionId": "session-code", + "turnId": "turn-code", + "type": kind, + "payload": payload, + } + await response.write(("data: " + json.dumps(data) + "\n\n").encode()) + + await send( + 1, + "tool.started", + { + "item": { + "id": "exec-1", + "type": "commandExecution", + "command": "python -c 'print(2+3)'", + } + }, + ) + await asyncio.wait_for(self.release.wait(), 5) + await send(2, "tool.delta", {"itemId": "exec-1", "delta": "5\n"}) + await send( + 3, + "tool.completed", + { + "item": { + "id": "exec-1", + "type": "commandExecution", + "aggregatedOutput": "5\n", + "exitCode": 0, + } + }, + ) + await send(4, "message.delta", {"itemId": "msg", "delta": "Answer: 5"}) + await send( + 5, + "turn.completed", + { + "status": "failed" if self.fail else "completed", + "finalText": "Answer: 5", + "error": {"message": "fixture failure"}, + }, + ) + return response + + +@pytest.mark.asyncio +async def test_code_child_streams_tools_before_final_and_reuses_session(): + fixture = CodeFixture() + app = web.Application() + app.router.add_route("*", "/{path:.*}", fixture.handle) + async with TestServer(app) as server: + child = AgentkitRemoteSandboxAgent( + name="sandbox", + tool_type="CodeEnv", + endpoint=str(server.make_url("/?route=fixture")), + ) + runner, model = await make_runner(child) + + async def observe(event): + if event.author == "sandbox" and event.get_function_calls(): + assert not any(p.text == "Answer: 5" for p in event.content.parts) + fixture.release.set() + + first = await collect(runner, observe) + second = await collect(runner, observe) + for events in (first, second): + assert not [e.error_message for e in events if e.error_message] + fc = next( + e.get_function_calls()[0] + for e in events + if e.author == "sandbox" and e.get_function_calls() + ) + fr = next( + e.get_function_responses()[0] + for e in events + if e.author == "sandbox" and e.get_function_responses() + ) + assert fc.id == fr.id and fr.response["exitCode"] == 0 + assert any( + e.partial and e.content.parts[0].thought + for e in events + if e.author == "sandbox" + ) + assert ( + sum( + e.author == "sandbox" + and not e.partial + and e.content.parts[0].text == "Answer: 5" + for e in events + ) + == 1 + ) + assert fixture.starts == 2 + assert fixture.session_keys[0] == fixture.session_keys[1] + assert model._calls <= 2 # no post-delegation summary + + +@pytest.mark.asyncio +async def test_code_failure_is_not_a_successful_final(): + fixture = CodeFixture() + fixture.fail = True + fixture.release.set() + app = web.Application() + app.router.add_route("*", "/{path:.*}", fixture.handle) + async with TestServer(app) as server: + runner, _ = await make_runner( + AgentkitRemoteSandboxAgent( + name="sandbox", + tool_type="CodeEnv", + endpoint=str(server.make_url("/?route=fixture")), + ) + ) + events = await collect(runner) + assert any("fixture failure" in (e.error_message or "") for e in events) + + +@pytest.mark.asyncio +async def test_explicit_type_skips_get_tool_and_is_passed_to_lease(): + agent = AgentkitRemoteSandboxAgent( + name="sandbox", tool_id="private", tool_type="CodeEnv" + ) + ctx = SimpleNamespace( + app_name="app", user_id="u", session=SimpleNamespace(state={}) + ) + lease = SimpleNamespace( + session_id="physical", select_endpoint=lambda **kw: "https://example.test" + ) + with ( + patch.object( + AgentkitRemoteSandboxAgent, + "_discover_type", + side_effect=AssertionError("must not discover"), + ), + patch( + "veadk.agents.agentkit_remote_sandbox_agent.ensure_agentkit_session_lease", + return_value=lease, + ), + ): + assert (await agent._resolve(ctx, "logical"))[0] == "CodeEnv" + + +def test_private_discovery_errors_and_constructor_is_lazy(): + agent = AgentkitRemoteSandboxAgent(name="sandbox", tool_id="private") + with ( + patch( + "veadk.agents.agentkit_remote_sandbox_agent.get_agentkit_credentials", + return_value=("fake-ak", "fake-sk", {}), + ), + patch("agentkit.sdk.tools.client.AgentkitToolsClient") as client, + ): + client.return_value.get_tool.return_value = SimpleNamespace(tool_type="Private") + with pytest.raises(SandboxAgentError, match="specify tool_type"): + agent._discover_type("private", {}) + + +@pytest.mark.parametrize("kind", ["", "Private", "All-in-one"]) +def test_invalid_explicit_type_rejected(kind): + with pytest.raises(ValueError): + AgentkitRemoteSandboxAgent(name="sandbox", tool_type=kind) + + +@pytest.mark.asyncio +async def test_discovery_cache_is_user_scoped(): + agent = AgentkitRemoteSandboxAgent(name="sandbox", tool_id="tool") + lease = SimpleNamespace( + session_id="physical", select_endpoint=lambda **kw: "https://example.test" + ) + with ( + patch.object( + AgentkitRemoteSandboxAgent, "_discover_type", return_value="Skill" + ) as discover, + patch( + "veadk.agents.agentkit_remote_sandbox_agent.ensure_agentkit_session_lease", + return_value=lease, + ), + ): + for user in ["u1", "u1", "u2"]: + await agent._resolve( + SimpleNamespace( + app_name="app", user_id=user, session=SimpleNamespace(state={}) + ), + user, + ) + assert discover.call_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [True, False]) +async def test_skill_a2a_preserves_tool_parts_and_context(streaming): + from google.adk.a2a.converters.part_converter import convert_genai_part_to_a2a_part + + release = asyncio.Event() + contexts = [] + methods = [] + + def wire(part): + return convert_genai_part_to_a2a_part(part).model_dump( + mode="json", by_alias=True, exclude_none=True + ) + + calls = [ + wire( + types.Part( + function_call=types.FunctionCall( + id="skill-exec", name="add_numbers", args={"a": 2, "b": 3} + ) + ) + ) + ] + results = [ + wire( + types.Part( + function_response=types.FunctionResponse( + id="skill-exec", name="add_numbers", response={"sum": 5} + ) + ) + ) + ] + + def task(state="working"): + result = { + "kind": "task", + "id": "a2a-task", + "contextId": "remote-context", + "status": {"state": state}, + } + if state == "completed": + result["artifacts"] = [ + {"artifactId": "tools", "parts": calls + results}, + { + "artifactId": "answer", + "parts": [{"kind": "text", "text": "Skill computed 5"}], + }, + ] + return result + + async def handle(request): + assert request.query.get("route") == "fixture" + if request.method == "GET": + return web.json_response( + { + "name": "skill", + "description": "fixture", + "version": "1", + "url": "http://127.0.0.1:1/a2a", + "capabilities": {"streaming": streaming}, + "defaultInputModes": ["text/plain"], + "defaultOutputModes": ["text/plain"], + "skills": [], + } + ) + body = await request.json() + methods.append(body["method"]) + if body["method"].startswith("message/"): + contexts.append(body["params"]["message"].get("contextId")) + + def reply(result): + return {"jsonrpc": "2.0", "id": body["id"], "result": result} + + if body["method"] == "tasks/get": + return web.json_response(reply(task("completed"))) + if not streaming: + return web.json_response(reply(task())) + response = web.StreamResponse(headers={"Content-Type": "text/event-stream"}) + await response.prepare(request) + + async def send(result): + await response.write( + ("data: " + json.dumps(reply(result)) + "\n\n").encode() + ) + + await send(task()) + + def artifact(identifier, parts): + return { + "kind": "artifact-update", + "taskId": "a2a-task", + "contextId": "remote-context", + "artifact": {"artifactId": identifier, "parts": parts}, + "append": False, + "lastChunk": True, + } + + await send(artifact("calls", calls)) + await asyncio.wait_for(release.wait(), 5) + await send(artifact("results", results)) + # Exercise partial artifact updates that default RemoteA2aAgent can drop. + await send( + { + "kind": "artifact-update", + "taskId": "a2a-task", + "contextId": "remote-context", + "artifact": { + "artifactId": "answer", + "parts": [{"kind": "text", "text": "Skill "}], + }, + "append": False, + "lastChunk": False, + } + ) + await send( + { + "kind": "artifact-update", + "taskId": "a2a-task", + "contextId": "remote-context", + "artifact": { + "artifactId": "answer", + "parts": [{"kind": "text", "text": "computed 5"}], + }, + "append": True, + "lastChunk": True, + } + ) + await send( + { + "kind": "status-update", + "taskId": "a2a-task", + "contextId": "remote-context", + "status": {"state": "completed"}, + "final": True, + } + ) + return response + + app = web.Application() + app.router.add_route("*", "/{path:.*}", handle) + async with TestServer(app) as server: + runner, model = await make_runner( + AgentkitRemoteSandboxAgent( + name="sandbox", + tool_type="Skill", + endpoint=str(server.make_url("/?route=fixture")), + ) + ) + + async def observe(event): + if event.author == "sandbox" and event.get_function_calls(): + release.set() + + for _ in range(2): + events = await collect(runner, observe) + assert not [e.error_message for e in events if e.error_message] + assert any(e.author == "sandbox" and e.get_function_calls() for e in events) + assert any( + e.author == "sandbox" and e.get_function_responses() for e in events + ) + assert any( + e.author == "sandbox" + and not e.partial + and e.content.parts[0].text == "Skill computed 5" + for e in events + ) + assert contexts == [None, "remote-context"] + assert ("message/stream" if streaming else "message/send") in methods + assert model._calls <= 2 + + +@pytest.mark.asyncio +async def test_cancelling_code_invocation_interrupts_remote_turn(): + fixture = CodeFixture() + app = web.Application() + app.router.add_route("*", "/{path:.*}", fixture.handle) + async with TestServer(app) as server: + runner, _ = await make_runner( + AgentkitRemoteSandboxAgent( + name="sandbox", + tool_type="CodeEnv", + endpoint=str(server.make_url("/?route=fixture")), + ) + ) + seen = asyncio.Event() + + async def observe(event): + if event.author == "sandbox" and event.get_function_calls(): + seen.set() + + running = asyncio.create_task(collect(runner, observe)) + await asyncio.wait_for(seen.wait(), 5) + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running + assert fixture.cancelled + fixture.release.set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [True, False]) +async def test_skill_against_native_adk_a2a_server(streaming): + """Real ADK A2A serialization must expose the tool before it completes.""" + import socket + import uvicorn + from google.adk.agents import LlmAgent + from google.adk.a2a.utils.agent_to_a2a import to_a2a + + release = asyncio.Event() + + async def add_numbers(a: int, b: int) -> dict: + """Add numbers after the remote client has observed the call.""" + await asyncio.wait_for(release.wait(), 10) + return {"sum": a + b} + + class RemoteModel(BaseLlm): + model: str = "offline-remote" + _step: int = PrivateAttr(default=0) + + async def generate_content_async(self, llm_request, stream=False): + self._step += 1 + part = ( + types.Part( + function_call=types.FunctionCall( + name="add_numbers", args={"a": 2, "b": 3} + ) + ) + if self._step == 1 + else types.Part(text="Native Skill computed 5") + ) + yield LlmResponse(content=types.Content(role="model", parts=[part])) + + remote = LlmAgent(name="native_skill", model=RemoteModel(), tools=[add_numbers]) + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + from a2a.types import AgentCard, AgentCapabilities + + card = AgentCard( + name="native_skill", + description="fixture", + version="1", + url=f"http://127.0.0.1:{port}/", + capabilities=AgentCapabilities(streaming=streaming), + default_input_modes=["text/plain"], + default_output_modes=["text/plain"], + skills=[], + ) + application = to_a2a(remote, host="127.0.0.1", port=port, agent_card=card) + server = uvicorn.Server( + uvicorn.Config(application, host="127.0.0.1", port=port, log_level="error") + ) + serving = asyncio.create_task(server.serve()) + try: + async with asyncio.timeout(10): + while not server.started: + if serving.done(): + await serving + await asyncio.sleep(0.05) + runner, _ = await make_runner( + AgentkitRemoteSandboxAgent( + name="sandbox", tool_type="Skill", endpoint=f"http://127.0.0.1:{port}" + ) + ) + + async def observe(event): + if event.author == "sandbox" and event.get_function_calls(): + release.set() + + events = await collect(runner, observe) + assert release.is_set() + assert not [e.error_message for e in events if e.error_message] + assert any(e.author == "sandbox" and e.get_function_responses() for e in events) + assert any( + e.author == "sandbox" + and not e.partial + and e.content.parts[0].text == "Native Skill computed 5" + for e in events + ) + finally: + release.set() + server.should_exit = True + await asyncio.wait_for(serving, 10) + + +@pytest.mark.asyncio +async def test_concurrent_users_have_distinct_worker_binding_keys(): + fixture = CodeFixture() + fixture.release.set() + app = web.Application() + app.router.add_route("*", "/{path:.*}", fixture.handle) + async with TestServer(app) as server: + child = AgentkitRemoteSandboxAgent( + name="sandbox", + tool_type="CodeEnv", + endpoint=str(server.make_url("/?route=fixture")), + ) + runner, _ = await make_runner(child) + await runner.session_service.create_session( + app_name="probe", user_id="other-user", session_id="session" + ) + first, second = await asyncio.gather( + collect(runner), collect(runner, user_id="other-user") + ) + assert len(set(fixture.session_keys)) == 2 + assert not [e.error_message for e in first + second if e.error_message] + assert {e.invocation_id for e in first}.isdisjoint( + {e.invocation_id for e in second} + ) diff --git a/tests/agents/test_sandbox_session.py b/tests/agents/test_sandbox_session.py new file mode 100644 index 000000000..4cd011f3a --- /dev/null +++ b/tests/agents/test_sandbox_session.py @@ -0,0 +1,116 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Control-plane pagination, rotation and ambiguous creation regression tests.""" + +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from veadk.agents import _sandbox_session as sessions + + +def info(sid="old", user="logical", seconds=3600, status="Ready"): + return SimpleNamespace( + session_id=sid, + user_session_id=user, + status=status, + created_at="2026-09-06T00:00:00Z", + expire_at=(datetime.now(timezone.utc) + timedelta(seconds=seconds)).isoformat(), + ) + + +def test_reuses_lease_only_with_sufficient_lifetime(): + current = info() + client = Mock() + client.list_sessions.return_value = SimpleNamespace( + session_infos=[current], next_token=None + ) + result = sessions._get_or_create_agentkit_session( + client=client, + tool_id="tool", + tool_user_session_id="logical", + ttl=1800, + min_remaining_seconds=900, + ) + assert result is current + client.create_session.assert_not_called() + + +def test_rotates_expiring_session_and_recovers_lost_create_response(): + old = info(seconds=5) + rotated = sessions._rotated_user_session_id("logical", [old]) + created = info(sid="new", user=rotated, status="Starting") + client = Mock() + client.list_sessions.side_effect = [ + SimpleNamespace(session_infos=[old], next_token=None), + SimpleNamespace(session_infos=[old, created], next_token=None), + ] + client.create_session.side_effect = TimeoutError("fixture") + result = sessions._get_or_create_agentkit_session( + client=client, + tool_id="tool", + tool_user_session_id="logical", + ttl=1800, + min_remaining_seconds=900, + ) + assert result is created + assert client.create_session.call_count == 1 + assert client.create_session.call_args.args[0].user_session_id == rotated + + +def test_list_failure_does_not_create_another_session(): + client = Mock() + client.list_sessions.side_effect = TimeoutError("fixture") + with pytest.raises(TimeoutError): + sessions._get_or_create_agentkit_session( + client=client, tool_id="tool", tool_user_session_id="logical", ttl=1800 + ) + client.create_session.assert_not_called() + + +def test_pagination_and_repeated_token_guard(): + client = Mock() + client.list_sessions.side_effect = [ + SimpleNamespace(session_infos=[], next_token="next"), + SimpleNamespace(session_infos=[info()], next_token=None), + ] + assert ( + len( + sessions._list_agentkit_sessions( + client=client, tool_id="tool", physical_user_session_id_base="logical" + ) + ) + == 1 + ) + assert client.list_sessions.call_args.args[0].next_token == "next" + client.list_sessions.side_effect = None + client.list_sessions.return_value = SimpleNamespace( + session_infos=[], next_token="repeat" + ) + with pytest.raises(RuntimeError, match="repeated NextToken"): + sessions._list_agentkit_sessions( + client=client, tool_id="tool", physical_user_session_id_base="logical" + ) + + +def test_logical_id_encoding_is_stable_and_bounded(): + encoded = sessions._safe_agentkit_user_session_id("用户/" * 500) + assert len(encoded) <= 185 + assert encoded == sessions._safe_agentkit_user_session_id("用户/" * 500) + assert sessions._safe_agentkit_user_session_id( + "a/b" + ) != sessions._safe_agentkit_user_session_id("a_b") diff --git a/veadk/__init__.py b/veadk/__init__.py index 7891d6c47..a678e6f2f 100644 --- a/veadk/__init__.py +++ b/veadk/__init__.py @@ -31,7 +31,13 @@ def __getattr__(name): from veadk.runner import Runner return Runner + if name == "AgentkitRemoteSandboxAgent": + from veadk.agents.agentkit_remote_sandbox_agent import ( + AgentkitRemoteSandboxAgent, + ) + + return AgentkitRemoteSandboxAgent raise AttributeError(f"module 'veadk' has no attribute '{name}'") -__all__ = ["Agent", "Runner", "VERSION"] +__all__ = ["Agent", "Runner", "AgentkitRemoteSandboxAgent", "VERSION"] diff --git a/veadk/agent.py b/veadk/agent.py index 17740dc42..d524a547e 100644 --- a/veadk/agent.py +++ b/veadk/agent.py @@ -16,6 +16,7 @@ import os import warnings +from contextlib import aclosing from typing import TYPE_CHECKING, AsyncGenerator, Dict, Literal, Optional, Union from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow @@ -787,8 +788,12 @@ async def _run_async_impl( stream, so the surrounding ``Runner`` is unaffected. """ if self.runtime == "adk": - async for event in super()._run_async_impl(ctx): - yield event + # A transfer can close this wrapper before the LLM stream ends. + # Close it in the same task so tracing contexts do not leak into + # async-generator finalization or the receiving sub-agent. + async with aclosing(super()._run_async_impl(ctx)) as events: + async for event in events: + yield event return from veadk.runtime import get_runtime diff --git a/veadk/agents/REMOTE_SANDBOX.md b/veadk/agents/REMOTE_SANDBOX.md new file mode 100644 index 000000000..e66fc7bf8 --- /dev/null +++ b/veadk/agents/REMOTE_SANDBOX.md @@ -0,0 +1,71 @@ +# AgentKit remote sandbox agent + +```python +from veadk import Agent, AgentkitRemoteSandboxAgent + +root_agent = Agent( + name="coordinator", + instruction="Transfer sandbox execution tasks to sandbox.", + sub_agents=[AgentkitRemoteSandboxAgent( + name="sandbox", + description="Execute tasks in the remote sandbox and report results.", + tool_id="your-tool-id", + # tool_type="CodeEnv", # optional: CodeEnv or Skill + request_timeout=900, + expiry_buffer=90, + )], +) +``` + +The entry is a BaseAgent, not an AgentTool. Native transfer hands the invocation +to the child. The child yields ADK events directly; the coordinator does not +summarize its tool trace. It works with the standard Runner and `veadk web`. + +Explicit `tool_type` wins and skips GetTool discovery. Otherwise GetTool must +return Skill or CodeEnv; Private/All-in-one/VibeSkill/unknown values fail with a +message to specify the compatible protocol explicitly. Construction/import does +not perform network requests. Omitted tool_id uses AGENTKIT_TOOL_ID. Existing +AgentKit region and credential resolution applies. GetTool permission is needed +only for discovery; Session permissions are always needed for managed sessions. + +Skill uses A2A streaming when advertised, or polls non-streaming tasks. The +adapter preserves structured ADK/A2A tool parts and partial artifact updates. +It cannot reconstruct tool traces a remote service never emits. CodeEnv uses +Sidecar HTTP/SSE protocol v1 with a tool_events capability, available in the +actb-mono code-env 1.1.2.3 image. No A2A wrapper is required for CodeEnv. + +Calls/results retain child author and matching IDs. Tool stdout progress is +shown as thought/progress text, separate from the final response. Errors are +explicit. Only text task input is currently supported. The new invocation sends +the latest user text; it does not forward the entire parent system prompt or +parent reasoning history. Follow-up context lives in the remote thread/context. + +An optional pre-provisioned `endpoint` requires explicit tool_type and bypasses +GetTool/CreateSession; this is useful for local fixtures. It has no platform TTL +management. `api_key` supplies the optional X-API-Key header. Runtime inbound_auth +credentials are resolved per invocation; clients are not shared between users. + +Session binding includes application, user, conversation and agent identity. +Physical Session rotation discards the old A2A context. Logical IDs are hashed; +creation responses lost on the network are recovered by querying the platform. +Before submission, remaining lifetime is checked again after data-plane readiness. +request_timeout bounds execution, ready_timeout bounds readiness, expiry_buffer +reserves remaining platform lifetime. A running task is never resubmitted merely +because its stream disconnected. Cancellation of the owning invocation attempts +remote cancellation; acknowledgement is not proof all external effects stopped. +A browser disconnect and explicit Runner cancellation are different operations; +the hosting application owns their relationship. + +Skill context is instance-local (up to 256 conversations and 64 invocations per +physical binding); no cross-process recovery is claimed. Code Sidecar stores +turns/events in SQLite and refuses uncertain reruns after restart. Neither +transport migrates filesystem state to a new physical sandbox. For mutually +untrusted users allocate separate physical sandboxes; threads are not filesystem +security boundaries. Shared explicit endpoints are intended for trusted callers. + +Validated with ADK 2.2.0 using deterministic models and HTTP fixtures. Native +/run_sse plus a real Code image verifies pre-completion deltas, structured tools +and a single final without parent summarization. SDK tests cover A2A streaming +and polling, two-turn context reuse, type precedence, cancellation, error results, +Session pagination/rotation and ambiguous create recovery. Cloud/provider +acceptance requires separately configured Tools and model credentials. diff --git a/veadk/agents/_sandbox_code.py b/veadk/agents/_sandbox_code.py new file mode 100644 index 000000000..0d091bd3a --- /dev/null +++ b/veadk/agents/_sandbox_code.py @@ -0,0 +1,254 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Translate Code Sidecar protocol v1 directly into ADK events.""" + +from __future__ import annotations + +import asyncio +import re +from contextlib import suppress + +from google.adk.events import Event +from google.genai import types + +from veadk.agents.agentkit_remote_sandbox_agent import SandboxAgentError, binding_key +from veadk.tools.sandbox.codex_worker_client import CodexWorkerClient, CodexWorkerError + + +def observation(value, secrets=()): + """Bound display data and redact credential fields and signed URLs.""" + if isinstance(value, dict): + return { + k: "[REDACTED]" + if re.search(r"(?i)authorization|token|api.?key|password|secret", k) + else observation(v, secrets) + for k, v in value.items() + } + if isinstance(value, list): + return [observation(v, secrets) for v in value[:100]] + if isinstance(value, str): + for secret in secrets: + if secret: + value = value.replace(secret, "[REDACTED]") + value = re.sub(r"https?://[^\s\"<>]+", "[URL]", value) + value = re.sub( + r"(?i)(bearer\s+|(?:api.?key|token|password|secret)\s*[=:]\s*)[^\s,;]+", + r"\1[REDACTED]", + value, + ) + return value[:16000] + return value + + +def check_lease(agent, lease): + if lease: + remaining = lease.remaining_seconds() + if ( + remaining is None + or remaining <= agent.request_timeout + agent.expiry_buffer + ): + raise SandboxAgentError( + "Sandbox session no longer has enough execution lifetime" + ) + + +def event_for(agent, ctx, parts, *, partial=False, metadata=None): + return Event( + author=agent.name, + invocation_id=ctx.invocation_id, + branch=ctx.branch, + partial=partial, + content=types.Content(role="model", parts=parts), + custom_metadata=metadata or {}, + ) + + +async def code_events( + *, agent, ctx, endpoint, headers, binding, turn_key, task, context, lease +): + sid = tid = None + completed = False + tools = {} + async with CodexWorkerClient(endpoint, headers=headers) as client: + try: + async with asyncio.timeout(agent.ready_timeout): + while True: + try: + ready = await client.request("GET", "/readyz") + break + except CodexWorkerError as exc: + if exc.status_code == 404: + raise SandboxAgentError( + "CodeEnv image has no Sidecar; install code-env 1.1.2.3 or newer" + ) from None + if exc.status_code and exc.status_code < 500: + raise SandboxAgentError(str(exc)) from None + await asyncio.sleep(1) + if ready.get("schemaVersion") != 1 or "tool_events" not in ready.get( + "capabilities", [] + ): + raise SandboxAgentError( + "CodeEnv Sidecar protocol is incompatible; version 1 with tool_events is required" + ) + check_lease(agent, lease) + async with asyncio.timeout(agent.request_timeout): + sid = await client.create_session(binding) + started = await client.start_turn(sid, task, turn_key) + tid = started["turnId"] + async for raw in client.events(sid, tid): + kind, payload = raw["type"], raw["payload"] + metadata = { + "remote_sandbox": { + "type": "CodeEnv", + "session_id": sid, + "turn_id": tid, + "event_id": raw["eventId"], + "event_type": kind, + } + } + if kind == "message.delta" and payload.get("delta"): + yield event_for( + agent, + ctx, + [types.Part(text=payload["delta"])], + partial=True, + metadata=metadata, + ) + elif kind in {"tool.started", "tool.completed"}: + item = payload["item"] + item_id = str(item["id"]) + call_id = "remote_" + binding_key(tid, item_id)[:24] + name = re.sub( + r"[^a-zA-Z0-9_]", "_", str(item.get("tool") or item["type"]) + )[:64] + safe = observation(item, headers.values()) + if item_id not in tools: + tools[item_id] = name + args = { + k: v + for k, v in safe.items() + if k + not in { + "id", + "type", + "status", + "aggregatedOutput", + "exitCode", + "durationMs", + "result", + "error", + } + } + yield event_for( + agent, + ctx, + [ + types.Part( + function_call=types.FunctionCall( + id=call_id, name=name, args=args + ) + ) + ], + metadata=metadata, + ) + if kind == "tool.completed": + yield event_for( + agent, + ctx, + [ + types.Part( + function_response=types.FunctionResponse( + id=call_id, + name=tools[item_id], + response=safe, + ) + ) + ], + metadata=metadata, + ) + elif kind in {"tool.delta", "reasoning.delta"} and payload.get( + "delta" + ): + prefix = "Tool output: " if kind == "tool.delta" else "" + metadata["remote_sandbox"]["item_id"] = payload.get("itemId") + yield event_for( + agent, + ctx, + [ + types.Part( + text=prefix + + observation(payload["delta"], headers.values()), + thought=True, + ) + ], + partial=True, + metadata=metadata, + ) + elif ( + kind == "message.completed" + and payload.get("phase") == "commentary" + ): + yield event_for( + agent, + ctx, + [types.Part(text=payload.get("text", ""), thought=True)], + metadata=metadata, + ) + elif kind == "turn.completed": + completed = True + if payload["status"] != "completed": + error = ( + payload.get("error", {}).get("message") + or payload.get("reason") + or payload["status"] + ) + raise SandboxAgentError( + f"CodeEnv turn {payload['status']}: {observation(error, headers.values())}" + ) + yield event_for( + agent, + ctx, + [ + types.Part( + text=payload.get("finalText") + or "Sandbox task completed." + ) + ], + metadata=metadata, + ) + return + raise SandboxAgentError( + "CodeEnv stream ended without a terminal event; do not resubmit" + ) + except CodexWorkerError as exc: + raise SandboxAgentError(str(exc)) from None + finally: + # Covers timeout, task cancellation and async-generator close. A + # browser subscriber alone must not close the owning Runner task. + if sid and not completed: + + async def cancel_accepted(): + turn = tid + if turn is None: + turn = ( + await client.request( + "GET", client.turn_path(sid) + "/by-key/" + turn_key + ) + )["turnId"] + await client.cancel(sid, turn) + + with suppress(Exception): + await asyncio.wait_for(cancel_accepted(), timeout=25) diff --git a/veadk/agents/_sandbox_session.py b/veadk/agents/_sandbox_session.py new file mode 100644 index 000000000..89f9066a8 --- /dev/null +++ b/veadk/agents/_sandbox_session.py @@ -0,0 +1,415 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Session leases for remote sandbox agents (control plane only).""" + +from __future__ import annotations +import hashlib +import re +import threading +import time +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any, Optional +from veadk.tools.builtin_tools._agentkit import ( + get_agentkit_credentials, + get_agentkit_endpoint_config, +) +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) +_SESSION_READY_TIMEOUT = 120.0 +_SESSION_POLL_INTERVAL = 1.0 +_SESSION_TERMINAL_STATUSES = frozenset({"failed", "terminating", "terminated"}) +_SESSION_REUSABLE_STATUSES = frozenset({"starting", "ready"}) +_SESSION_LIST_PAGE_SIZE = 100 +_SESSION_USER_ID_MAX_LENGTH = 200 +_SESSION_ROTATION_SUFFIX_LENGTH = 15 +_SESSION_LOCK_STRIPE_COUNT = 64 +_session_locks = tuple(threading.Lock() for _ in range(_SESSION_LOCK_STRIPE_COUNT)) + + +@dataclass(frozen=True) +class AgentKitSessionLease: + """A resolved AgentKit Session and its session-scoped data-plane endpoints.""" + + tool_id: str + logical_user_session_id: str + user_session_id: str + session_id: str + status: str + endpoint: str + internal_endpoint: str + created_at: str + expire_at: str + + def select_endpoint(self, *, prefer_internal_endpoint: bool = False) -> str: + if prefer_internal_endpoint: + return self.internal_endpoint or self.endpoint + return self.endpoint or self.internal_endpoint + + def remaining_seconds(self, *, now: datetime | None = None) -> float | None: + expires_at = _parse_agentkit_timestamp(self.expire_at) + if expires_at is None: + return None + current = now or datetime.now(timezone.utc) + if current.tzinfo is None: + current = current.replace(tzinfo=timezone.utc) + return (expires_at - current.astimezone(timezone.utc)).total_seconds() + + +def _parse_agentkit_timestamp(value: object) -> datetime | None: + if not isinstance(value, str) or not value.strip(): + return None + normalized = value.strip() + if normalized.endswith("Z"): + normalized = normalized[:-1] + "+00:00" + try: + parsed = datetime.fromisoformat(normalized) + except ValueError: + logger.warning("Invalid AgentKit Session timestamp") + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _safe_agentkit_user_session_id(logical_user_session_id: str) -> str: + """Return a contract-compliant, stable base for physical UserSessionIds.""" + if not logical_user_session_id: + raise ValueError("tool_user_session_id must not be empty") + normalized = re.sub(r"[^A-Za-z0-9_-]", "_", logical_user_session_id) + digest = hashlib.sha256(logical_user_session_id.encode("utf-8")).hexdigest()[:12] + if normalized != logical_user_session_id: + normalized = f"{normalized}_{digest}" + max_base_length = _SESSION_USER_ID_MAX_LENGTH - _SESSION_ROTATION_SUFFIX_LENGTH + if len(normalized) > max_base_length: + normalized = f"{normalized[: max_base_length - 13]}_{digest}" + return normalized + + +def _session_lock(tool_id: str, logical_user_session_id: str) -> threading.Lock: + digest = hashlib.sha256( + f"{tool_id}\0{logical_user_session_id}".encode("utf-8") + ).digest() + index = int.from_bytes(digest[:4], "big") % _SESSION_LOCK_STRIPE_COUNT + return _session_locks[index] + + +def _session_has_enough_time( + session: object, + *, + min_remaining_seconds: float, + now: datetime, +) -> bool: + expires_at = _parse_agentkit_timestamp(getattr(session, "expire_at", None)) + if expires_at is None: + return min_remaining_seconds <= 0 + return (expires_at - now).total_seconds() > min_remaining_seconds + + +def _session_is_reusable( + session: object, + *, + physical_user_session_id_base: str, + min_remaining_seconds: float, + now: datetime, +) -> bool: + user_session_id = getattr(session, "user_session_id", None) + if not isinstance(user_session_id, str) or not ( + user_session_id == physical_user_session_id_base + or user_session_id.startswith(f"{physical_user_session_id_base}_r_") + ): + return False + status = (getattr(session, "status", None) or "").strip().lower() + return status in _SESSION_REUSABLE_STATUSES and _session_has_enough_time( + session, + min_remaining_seconds=min_remaining_seconds, + now=now, + ) + + +def _rotated_user_session_id( + physical_user_session_id_base: str, + sessions: list[object], +) -> str: + existing_ids = sorted( + str(getattr(session, "session_id", "") or "") for session in sessions + ) + generation = hashlib.sha256("\n".join(existing_ids).encode("utf-8")).hexdigest()[ + :12 + ] + suffix = f"_r_{generation}" + return f"{physical_user_session_id_base[: _SESSION_USER_ID_MAX_LENGTH - len(suffix)]}{suffix}" + + +def _get_or_create_agentkit_session( + *, + client, + tool_id: str, + tool_user_session_id: str, + ttl: int, + min_remaining_seconds: float = 0, +): + """Return a reusable physical Session for a stable logical session key.""" + from agentkit.sdk.tools import types as tools_types + + physical_user_session_id_base = _safe_agentkit_user_session_id(tool_user_session_id) + candidates = _list_agentkit_sessions( + client=client, + tool_id=tool_id, + physical_user_session_id_base=physical_user_session_id_base, + ) + now = datetime.now(timezone.utc) + reusable = [ + info + for info in candidates + if _session_is_reusable( + info, + physical_user_session_id_base=physical_user_session_id_base, + min_remaining_seconds=min_remaining_seconds, + now=now, + ) + ] + if reusable: + reusable.sort( + key=lambda info: getattr(info, "created_at", "") or "", reverse=True + ) + chosen = reusable[0] + logger.debug( + f"Reusing AgentKit session {getattr(chosen, 'session_id', None)} " + f"for logical UserSessionId={tool_user_session_id}" + ) + return chosen + + physical_user_session_id = ( + physical_user_session_id_base + if not candidates + else _rotated_user_session_id( + physical_user_session_id_base, + candidates, + ) + ) + try: + return client.create_session( + tools_types.CreateSessionRequest( + ToolId=tool_id, + UserSessionId=physical_user_session_id, + Ttl=ttl, + ) + ) + except Exception: + # CreateSession may have succeeded even if its response was lost. Recover + # the physical Session before deciding whether the operation failed. + refreshed = _list_agentkit_sessions( + client=client, + tool_id=tool_id, + physical_user_session_id_base=physical_user_session_id_base, + ) + recovered = [ + info + for info in refreshed + if getattr(info, "user_session_id", None) == physical_user_session_id + and (getattr(info, "status", None) or "").strip().lower() + in _SESSION_REUSABLE_STATUSES + ] + if recovered: + recovered.sort( + key=lambda info: getattr(info, "created_at", "") or "", + reverse=True, + ) + return recovered[0] + raise + + +def _list_agentkit_sessions( + *, + client, + tool_id: str, + physical_user_session_id_base: str, +) -> list[object]: + """List every physical Session associated with one logical key.""" + if not hasattr(client, "list_sessions"): + # Compatibility for older clients and lightweight test doubles. Current + # AgentKit SDK versions expose ListSessions and use the paginated path. + return [] + + from agentkit.sdk.tools import types as tools_types + + sessions: list[object] = [] + next_token: str | None = None + seen_tokens: set[str] = set() + while True: + request_kwargs: dict[str, object] = { + "ToolId": tool_id, + "Filters": [ + tools_types.FiltersItemForListSessions( + NameContains="UserSessionId", + Values=[physical_user_session_id_base], + ) + ], + "MaxResults": _SESSION_LIST_PAGE_SIZE, + } + if next_token: + request_kwargs["NextToken"] = next_token + listing = client.list_sessions( + tools_types.ListSessionsRequest(**request_kwargs) + ) + for info in getattr(listing, "session_infos", None) or []: + user_session_id = getattr(info, "user_session_id", None) + if user_session_id == physical_user_session_id_base or ( + isinstance(user_session_id, str) + and user_session_id.startswith(f"{physical_user_session_id_base}_r_") + ): + sessions.append(info) + + next_token = getattr(listing, "next_token", None) or None + if not next_token: + return sessions + if next_token in seen_tokens: + raise RuntimeError("AgentKit ListSessions returned a repeated NextToken") + seen_tokens.add(next_token) + + +def _agentkit_session_lease( + *, + session: object, + fallback_session: object | None, + tool_id: str, + logical_user_session_id: str, +) -> AgentKitSessionLease: + def value(name: str) -> str: + current = getattr(session, name, None) + fallback = getattr(fallback_session, name, None) if fallback_session else None + return str(current or fallback or "") + + return AgentKitSessionLease( + tool_id=tool_id, + logical_user_session_id=logical_user_session_id, + user_session_id=value("user_session_id") + or _safe_agentkit_user_session_id(logical_user_session_id), + session_id=value("session_id"), + status=value("status"), + endpoint=value("endpoint"), + internal_endpoint=value("internal_endpoint"), + created_at=value("created_at"), + expire_at=value("expire_at"), + ) + + +def ensure_agentkit_session_lease( + *, + tool_id: str, + tool_user_session_id: str, + tool_state: Optional[dict[str, Any]] = None, + ttl: int = 1800, + min_remaining_seconds: float = 0, + wait_until_ready: bool = True, + ready_timeout: float = _SESSION_READY_TIMEOUT, + poll_interval: float = _SESSION_POLL_INTERVAL, +) -> AgentKitSessionLease: + """Resolve a live Session lease for a stable logical UserSessionId.""" + from agentkit.sdk.tools import types as tools_types + from agentkit.sdk.tools.client import AgentkitToolsClient + + if not 60 <= ttl <= 86400: + raise ValueError("ttl must be between 60 and 86400 seconds") + if min_remaining_seconds < 0: + raise ValueError("min_remaining_seconds must be greater than or equal to 0") + if min_remaining_seconds >= 86400: + raise ValueError("min_remaining_seconds must be less than 86400 seconds") + if ready_timeout < 0: + raise ValueError("ready_timeout must be greater than or equal to 0") + if poll_interval <= 0: + raise ValueError("poll_interval must be greater than 0") + + required_ttl = min(86400, max(ttl, int(min_remaining_seconds + ready_timeout) + 1)) + _, region, _, _ = get_agentkit_endpoint_config() + ak, sk, header = get_agentkit_credentials(tool_state) + client = AgentkitToolsClient( + access_key=ak, + secret_key=sk, + region=region, + session_token=header.get("X-Security-Token", ""), + ) + + with _session_lock(tool_id, tool_user_session_id): + session = _get_or_create_agentkit_session( + client=client, + tool_id=tool_id, + tool_user_session_id=tool_user_session_id, + ttl=required_ttl, + min_remaining_seconds=min_remaining_seconds, + ) + session_id = getattr(session, "session_id", None) + if not session_id: + raise RuntimeError("AgentKit CreateSession response is missing SessionId") + + if not wait_until_ready: + lease = _agentkit_session_lease( + session=session, + fallback_session=None, + tool_id=tool_id, + logical_user_session_id=tool_user_session_id, + ) + if lease.endpoint or lease.internal_endpoint: + return lease + + deadline = time.monotonic() + ready_timeout + last_status = "Unknown" + while True: + current_session = client.get_session( + tools_types.GetSessionRequest( + ToolId=tool_id, + SessionId=session_id, + ) + ) + status = (getattr(current_session, "status", None) or "").strip() + last_status = status or "Unknown" + logger.debug("AgentKit session %s status: %s", session_id, last_status) + normalized_status = status.lower() + if normalized_status == "ready": + lease = _agentkit_session_lease( + session=current_session, + fallback_session=session, + tool_id=tool_id, + logical_user_session_id=tool_user_session_id, + ) + if not lease.select_endpoint(): + raise RuntimeError( + f"AgentKit session {session_id} is Ready but has no endpoint" + ) + if not _session_has_enough_time( + current_session, + min_remaining_seconds=min_remaining_seconds, + now=datetime.now(timezone.utc), + ): + raise RuntimeError( + f"AgentKit session {session_id} became Ready without enough " + "remaining lifetime" + ) + return lease + if normalized_status in _SESSION_TERMINAL_STATUSES: + raise RuntimeError( + f"AgentKit session {session_id} entered terminal status {last_status}" + ) + + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + f"Timed out waiting for AgentKit session {session_id} to become " + f"Ready; last status: {last_status}" + ) + time.sleep(min(poll_interval, remaining)) diff --git a/veadk/agents/_sandbox_skill.py b/veadk/agents/_sandbox_skill.py new file mode 100644 index 000000000..839d2c416 --- /dev/null +++ b/veadk/agents/_sandbox_skill.py @@ -0,0 +1,278 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Skill A2A transport using public client APIs and structured ADK parts.""" + +from __future__ import annotations + +import asyncio +import json +from contextlib import suppress + +import httpx +from a2a.client import ClientConfig, ClientFactory +from a2a.client.errors import A2AClientHTTPError, A2AClientTimeoutError +from a2a.types import ( + AgentCard, + Message, + Part, + TaskArtifactUpdateEvent, + TaskIdParams, + TaskQueryParams, + TaskState, + TextPart, +) +from google.adk.a2a.converters.part_converter import convert_a2a_part_to_genai_part +from google.genai import types + +from veadk.agents._sandbox_code import check_lease, event_for, observation +from veadk.agents.agentkit_remote_sandbox_agent import SandboxAgentError + +_ACTIVE = {TaskState.submitted, TaskState.working} + + +def endpoint_path(endpoint, path): + url = httpx.URL(endpoint) + if url.scheme not in {"http", "https"} or url.userinfo or url.fragment: + raise SandboxAgentError("Invalid Skill endpoint") + return url.copy_with(path=url.path.rstrip("/") + "/" + path.lstrip("/")) + + +async def skill_events( + *, agent, ctx, endpoint, headers, binding, turn_key, task, context, lease +): + async with httpx.AsyncClient( + headers=headers, timeout=httpx.Timeout(30, connect=10), trust_env=False + ) as http: + async with asyncio.timeout(agent.ready_timeout): + while True: + try: + response = await http.get( + endpoint_path(endpoint, "/.well-known/agent-card.json") + ) + if response.status_code == 200: + card = AgentCard.model_validate(response.json()) + break + if response.status_code not in {502, 503, 504}: + raise SandboxAgentError( + f"Skill Agent Card returned HTTP {response.status_code}" + ) + except (httpx.TransportError, ValueError): + pass + await asyncio.sleep(1) + check_lease(agent, lease) + # Keep the platform authority/query, using only the Card's RPC path. + card = card.model_copy(deep=True) + card.url = str(endpoint_path(endpoint, httpx.URL(card.url).path)) + card.additional_interfaces = None + client = ClientFactory( + ClientConfig( + httpx_client=http, + streaming=bool(card.capabilities.streaming), + polling=True, + supported_transports=["JSONRPC"], + use_client_preference=True, + ) + ).create(card) + invocations = context.setdefault("invocations", {}) + if turn_key in invocations: + raise SandboxAgentError( + "This A2A invocation was already submitted; do not execute it again" + ) + if len(invocations) >= 64: + raise SandboxAgentError( + "Skill invocation capacity reached for this physical session" + ) + invocations[turn_key] = {"status": "submitting"} + message = Message( + message_id=turn_key, + role="user", + parts=[Part(root=TextPart(text=task))], + context_id=context.get("context_id"), + ) + remote_task = None + completed = False + emitted = set() + final_text = "" + snapshots = {} + + def convert(parts, *, source): + nonlocal final_text + result = [] + for wire in parts or []: + part = convert_a2a_part_to_genai_part(wire) + for part in part if isinstance(part, list) else [part]: + if part is None: + continue + if part.function_call or part.function_response: + safe = observation( + part.model_dump(exclude_none=True), headers.values() + ) + part = types.Part.model_validate(safe) + # Replayed Task snapshots repeat call/result parts. + identity = json.dumps(safe, sort_keys=True) + if identity in emitted: + continue + emitted.add(identity) + result.append( + event_for( + agent, + ctx, + [part], + metadata={ + "remote_sandbox": { + "type": "Skill", + "source": source, + } + }, + ) + ) + elif part.text: + result.append( + event_for( + agent, + ctx, + [part], + partial=True, + metadata={ + "remote_sandbox": { + "type": "Skill", + "source": source, + } + }, + ) + ) + return result + + def consume(response): + nonlocal remote_task, final_text, completed + if isinstance(response, Message): + completed = True + if response.context_id: + context["context_id"] = response.context_id + final_text = "\n".join( + p.root.text for p in response.parts if isinstance(p.root, TextPart) + ) + return convert(response.parts, source=response.message_id) + remote_task, update = response + invocations[turn_key] = { + "status": remote_task.status.state.value, + "task_id": remote_task.id, + } + if remote_task.context_id: + context["context_id"] = remote_task.context_id + if isinstance(update, TaskArtifactUpdateEvent): + artifact = update.artifact + text = "".join( + p.root.text for p in artifact.parts if isinstance(p.root, TextPart) + ) + snapshots[artifact.artifact_id] = ( + (snapshots.get(artifact.artifact_id, "") + text) + if update.append + else text + ) + return convert(artifact.parts, source=artifact.artifact_id) + events = [] + messages = list(remote_task.history or []) if update is None else [] + if remote_task.status.message: + messages.append(remote_task.status.message) + for msg in messages: + if msg.role == "user" or msg.message_id in emitted: + continue + emitted.add(msg.message_id) + events.extend(convert(msg.parts, source=msg.message_id)) + if remote_task.status.state not in _ACTIVE: + text = "\n".join( + p.root.text for p in msg.parts if isinstance(p.root, TextPart) + ) + if text: + final_text = text + # Final/polled snapshots provide authoritative artifacts. Streaming + # artifact deltas were already shown; only extract tools on final. + if update is None or remote_task.status.state not in _ACTIVE: + for artifact in remote_task.artifacts or []: + snapshots[artifact.artifact_id] = "".join( + p.root.text + for p in artifact.parts + if isinstance(p.root, TextPart) + ) + tool_parts = [ + p for p in artifact.parts if not isinstance(p.root, TextPart) + ] + events.extend(convert(tool_parts, source=artifact.artifact_id)) + return events + + try: + async with asyncio.timeout(agent.request_timeout): + try: + async for response in client.send_message( + message, + request_metadata={ + "user_id": ctx.user_id, + "session_id": ctx.session.id, + }, + ): + for event in consume(response): + yield event + except ( + httpx.TransportError, + A2AClientHTTPError, + A2AClientTimeoutError, + ): + # An accepted Task can be queried; never send the message twice. + if remote_task is None: + raise SandboxAgentError( + "A2A submission result is unconfirmed; do not resubmit" + ) from None + while remote_task and remote_task.status.state in _ACTIVE: + await asyncio.sleep(1) + queried = await client.get_task( + TaskQueryParams(id=remote_task.id, history_length=20) + ) + for event in consume((queried, None)): + yield event + if remote_task: + completed = remote_task.status.state not in _ACTIVE + if remote_task.status.state != TaskState.completed: + raise SandboxAgentError( + f"Skill task ended with state {remote_task.status.state.value}" + ) + if not completed: + raise SandboxAgentError( + "Skill stream ended without a terminal task" + ) + text = ( + "\n".join(text for text in snapshots.values() if text) + or final_text + or "Sandbox task completed." + ) + yield event_for( + agent, + ctx, + [types.Part(text=text)], + metadata={ + "remote_sandbox": { + "type": "Skill", + "task_id": remote_task.id if remote_task else None, + } + }, + ) + finally: + if remote_task and not completed: + with suppress(Exception): + await asyncio.wait_for( + client.cancel_task(TaskIdParams(id=remote_task.id)), timeout=10 + ) + await client.close() diff --git a/veadk/agents/agentkit_remote_sandbox_agent.py b/veadk/agents/agentkit_remote_sandbox_agent.py new file mode 100644 index 000000000..1ac62b1bc --- /dev/null +++ b/veadk/agents/agentkit_remote_sandbox_agent.py @@ -0,0 +1,248 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""A native ADK sub-agent backed by an AgentKit Skill or CodeEnv session.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +from contextlib import aclosing +from typing import Literal + +from google.adk.agents.base_agent import BaseAgent +from google.adk.events import Event +from google.genai import types +from pydantic import Field, PrivateAttr, SecretStr + +from veadk.agents._sandbox_session import ensure_agentkit_session_lease +from veadk.tools.builtin_tools._agentkit import ( + get_agentkit_credentials, + get_agentkit_endpoint_config, + resolve_agentkit_tool_id, +) + + +def binding_key(*parts: object) -> str: + return hashlib.sha256(json.dumps(parts, ensure_ascii=False).encode()).hexdigest() + + +class SandboxAgentError(RuntimeError): + """A safe diagnostic suitable for returning to the caller.""" + + +class AgentkitRemoteSandboxAgent(BaseAgent): + """Register in ``Agent(sub_agents=[...])``; no network during construction. + + Explicit tool_type wins over discovery. Private tools require an explicit + compatible type. endpoint is an optional pre-provisioned session URL for + local testing; it requires an explicit type and bypasses the control plane. + Remote calls/results are observation events, never locally executed tools. + """ + + tool_id: str | None = None + tool_type: Literal["Skill", "CodeEnv"] | None = None + request_timeout: float = Field(default=900, gt=0, lt=86000) + expiry_buffer: float = Field(default=90, ge=0, lt=86400) + ready_timeout: float = Field(default=120, gt=0) + ttl: int = Field(default=1800, ge=60, le=86400) + prefer_internal_endpoint: bool = False + endpoint: str | None = Field(default=None, repr=False) + api_key: SecretStr | None = Field(default=None, repr=False, exclude=True) + _types: dict = PrivateAttr(default_factory=dict) + _contexts: dict = PrivateAttr(default_factory=dict) + _locks: dict = PrivateAttr(default_factory=dict) + _type_lock: asyncio.Lock = PrivateAttr(default_factory=asyncio.Lock) + + def model_post_init(self, context): + super().model_post_init(context) + if self.endpoint and self.tool_type is None: + raise ValueError("endpoint requires explicit tool_type") + if self.request_timeout + self.expiry_buffer + 2 * self.ready_timeout >= 86400: + raise ValueError( + "Execution and readiness budget must fit within 86400 seconds" + ) + + def _discover_type(self, tool_id, state): + from agentkit.sdk.tools.client import AgentkitToolsClient + from agentkit.sdk.tools.types import GetToolRequest + + _, region, _, _ = get_agentkit_endpoint_config() + ak, sk, headers = get_agentkit_credentials(state) + result = AgentkitToolsClient( + access_key=ak, + secret_key=sk, + region=region, + session_token=headers.get("X-Security-Token", ""), + ).get_tool(GetToolRequest(ToolId=tool_id)) + if result.tool_type not in {"Skill", "CodeEnv"}: + raise SandboxAgentError( + "GetTool must return Skill or CodeEnv; specify tool_type explicitly for a compatible custom tool" + ) + return result.tool_type + + async def _resolve(self, ctx, logical): + if self.endpoint: + return self.tool_type, self.endpoint, binding_key(self.endpoint), None + tool_id = self.tool_id or resolve_agentkit_tool_id() + state = dict(ctx.session.state) + kind = self.tool_type + if kind is None: + # Identity-scoped cache: never share discovery across application users. + cache_key = ( + ctx.app_name, + ctx.user_id, + tool_id, + get_agentkit_endpoint_config(), + ) + async with self._type_lock: + kind = self._types.get(cache_key) + if kind is None: + kind = await asyncio.to_thread(self._discover_type, tool_id, state) + if len(self._types) >= 256: + self._types.clear() + self._types[cache_key] = kind + lease = await asyncio.to_thread( + ensure_agentkit_session_lease, + tool_id=tool_id, + tool_user_session_id=logical, + tool_state=state, + ttl=self.ttl, + min_remaining_seconds=self.request_timeout + + self.expiry_buffer + + self.ready_timeout, + ready_timeout=self.ready_timeout, + ) + return ( + kind, + lease.select_endpoint( + prefer_internal_endpoint=self.prefer_internal_endpoint + ), + lease.session_id, + lease, + ) + + async def _headers(self, ctx): + headers = {} + if self.api_key: + headers["X-API-Key"] = self.api_key.get_secret_value() + if ctx.credential_service: + from google.adk.agents.callback_context import CallbackContext + from veadk.utils.auth import build_auth_config + + credential = await ctx.credential_service.load_credential( + auth_config=build_auth_config( + credential_key="inbound_auth", + auth_method="header", + header_scheme="bearer", + ), + callback_context=CallbackContext(ctx), + ) + if credential: + token = credential.api_key + if not token and credential.http and credential.http.credentials: + token = credential.http.credentials.token + if not token and credential.oauth2: + token = credential.oauth2.access_token + if token: + headers["inbound_auth"] = token + return headers + + async def _run_async_impl(self, ctx): + from veadk.agents._sandbox_code import code_events + from veadk.agents._sandbox_skill import skill_events + + logical = binding_key( + ctx.app_name, + ctx.user_id, + ctx.session.id, + self.name, + self.tool_id, + self.endpoint, + ) + if logical not in self._locks: + if len(self._locks) >= 256: + raise SandboxAgentError( + "Agent session capacity reached; create a new agent instance" + ) + self._locks[logical] = asyncio.Lock() + lock = self._locks[logical] + if lock.locked(): + raise SandboxAgentError( + "A sandbox invocation is already active for this session" + ) + async with lock: + try: + async with asyncio.timeout(self.ready_timeout * 2 + 60): + kind, endpoint, physical, lease = await self._resolve(ctx, logical) + headers = await self._headers(ctx) + user_event = next( + ( + e + for e in reversed(ctx.session.events) + if e.author == "user" and e.content + ), + None, + ) + content = user_event.content if user_event else ctx.user_content + if not content or not content.parts: + raise SandboxAgentError("A user task is required") + task = "\n".join(p.text for p in content.parts if p.text) + if not task.strip(): + raise SandboxAgentError( + "This sandbox entry currently requires a text task" + ) + binding = binding_key(logical, physical) + turn_key = binding_key(binding, ctx.invocation_id) + # Only a physical-session-bound context can be reused. + previous = self._contexts.get(logical) + if previous is None or previous.get("physical") != physical: + previous = {"physical": physical} + self._contexts[logical] = previous + backend = code_events if kind == "CodeEnv" else skill_events + async with asyncio.timeout(self.request_timeout + self.ready_timeout): + async with aclosing( + backend( + agent=self, + ctx=ctx, + endpoint=endpoint, + headers=headers, + binding=binding, + turn_key=turn_key, + task=task, + context=previous, + lease=lease, + ) + ) as stream: + async for event in stream: + yield event + except asyncio.CancelledError: + raise + except Exception as exc: + message = ( + str(exc) + if isinstance(exc, SandboxAgentError) + else f"Sandbox invocation failed ({type(exc).__name__}); execution may be unconfirmed" + ) + yield Event( + author=self.name, + invocation_id=ctx.invocation_id, + branch=ctx.branch, + error_message=message, + content=types.Content( + role="model", parts=[types.Part(text=message)] + ), + ) diff --git a/veadk/tools/sandbox/codex_worker_client.py b/veadk/tools/sandbox/codex_worker_client.py new file mode 100644 index 000000000..83061dc1b --- /dev/null +++ b/veadk/tools/sandbox/codex_worker_client.py @@ -0,0 +1,198 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Async client for the CodeEnv Codex worker HTTP/SSE protocol v1.""" + +from __future__ import annotations + +import asyncio +import json +from urllib.parse import quote + +import httpx + + +class CodexWorkerError(RuntimeError): + """A sanitized worker failure; never includes credential URLs or wire data.""" + + def __init__(self, message: str, *, status_code: int | None = None): + super().__init__(message) + self.status_code = status_code + + +class CodexWorkerClient: + def __init__( + self, + endpoint: str, + *, + api_key: str | None = None, + headers: dict[str, str] | None = None, + ): + url = httpx.URL(endpoint) + if ( + url.scheme not in {"http", "https"} + or not url.host + or url.userinfo + or url.fragment + ): + raise ValueError( + "endpoint must be an HTTP(S) URL without userinfo or fragment" + ) + if url.path.rstrip("/") not in {"", "/v1/codex-worker"}: + raise ValueError("endpoint path must be / or /v1/codex-worker") + self._url = url + self._headers = dict(headers or {}) + if api_key: + self._headers["X-API-Key"] = api_key + self._http = None + + async def __aenter__(self): + self._http = httpx.AsyncClient( + headers=self._headers, + timeout=httpx.Timeout(60, connect=10, write=10, pool=10), + follow_redirects=False, + trust_env=False, + ) + return self + + async def __aexit__(self, *args): + await self._http.aclose() + + def _endpoint(self, path): + # copy_with preserves platform authentication query parameters. + return self._url.copy_with(path="/v1/codex-worker" + path) + + async def request(self, method, path, *, body=None, key=None): + for attempt in range(3): + try: + response = await self._http.request( + method, + self._endpoint(path), + json=body, + headers={"Idempotency-Key": key} if key else None, + ) + except httpx.TransportError: + if attempt == 2: + raise CodexWorkerError( + "Worker request unavailable; execution may have started" + ) from None + await asyncio.sleep(0.1 * (attempt + 1)) + continue + if response.status_code >= 400: + # POST retries only use a stable idempotency key. A 5xx can + # represent an accepted request, never generate a replacement key. + if response.status_code >= 500 and key and attempt < 2: + await asyncio.sleep(0.1 * (attempt + 1)) + continue + raise CodexWorkerError( + f"Worker returned HTTP {response.status_code}", + status_code=response.status_code, + ) + if response.status_code >= 300: + raise CodexWorkerError("Worker redirects are not supported") + try: + data = response.json() + if not isinstance(data, dict): + raise ValueError + return data + except ValueError: + raise CodexWorkerError("Invalid worker response") from None + raise CodexWorkerError("Worker unavailable") + + async def create_session(self, key): + result = await self.request("POST", "/sessions", body={}, key=key) + if result.get("status") != "ready": + raise CodexWorkerError( + "Session initialization is unconfirmed; do not create a replacement" + ) + return result["sessionId"] + + @staticmethod + def turn_path(sid, tid=None): + path = f"/sessions/{quote(sid, safe='')}/turns" + return path if tid is None else path + "/" + quote(tid, safe="") + + async def start_turn(self, sid, task, key): + return await self.request( + "POST", self.turn_path(sid), body={"task": task}, key=key + ) + + async def cancel(self, sid, tid): + # turn/start may still be waiting for its Codex ID. Retry cancellation, + # not execution, for up to the worker's start RPC deadline. + for _ in range(40): + state = await self.request("GET", self.turn_path(sid, tid)) + if state["status"] in {"completed", "failed", "interrupted", "unknown"}: + return state + if state.get("codexTurnId"): + return await self.request("POST", self.turn_path(sid, tid) + "/cancel") + await asyncio.sleep(0.5) + raise CodexWorkerError("Cancellation could not confirm the remote turn ID") + + async def events(self, sid, tid, *, after_event_id=0): + cursor = after_event_id + for attempt in range(4): + try: + async with self._http.stream( + "GET", + self._endpoint(self.turn_path(sid, tid) + "/events"), + headers={"Last-Event-ID": str(cursor)}, + ) as response: + if response.status_code != 200: + raise CodexWorkerError( + f"Worker event stream returned HTTP {response.status_code}" + ) + data = [] + size = 0 + async for line in response.aiter_lines(): + size += len(line) + if size > 4 * 1024 * 1024: + raise CodexWorkerError("Worker event too large") + if line.startswith("data:"): + data.append(line[5:].lstrip()) + elif not line: + size = 0 + if not data: + continue + try: + event = json.loads("\n".join(data)) + except ValueError: + raise CodexWorkerError("Invalid worker event") from None + data = [] + if event.get("schemaVersion") != 1: + raise CodexWorkerError( + "Unsupported worker event or expired stream" + ) + if ( + event.get("sessionId") != sid + or event.get("turnId") != tid + ): + raise CodexWorkerError("Worker event binding mismatch") + sequence = event.get("eventId") + if not isinstance(sequence, int) or sequence < 1: + raise CodexWorkerError("Invalid worker event cursor") + if sequence <= cursor: + continue + cursor = sequence + yield event + if event["type"] == "turn.completed": + return + except httpx.TransportError: + pass + if attempt < 3: + await asyncio.sleep(0.1 * (attempt + 1)) + raise CodexWorkerError( + "Worker stream ended without a terminal event; do not restart the task" + ) From e206ad7a7abb16dda9c29caa2c7e6e1cb8903a39 Mon Sep 17 00:00:00 2001 From: "hanzhi.421" Date: Sun, 6 Sep 2026 21:59:12 +0800 Subject: [PATCH 2/3] fix(agents): support remote sandbox timeouts on Python 3.10 --- pyproject.toml | 1 + tests/agents/test_remote_sandbox_agent.py | 19 +++++++++++---- uv.lock | 2 ++ veadk/agents/_sandbox_code.py | 5 ++-- veadk/agents/_sandbox_skill.py | 5 ++-- veadk/agents/_sandbox_timeout.py | 23 +++++++++++++++++++ veadk/agents/agentkit_remote_sandbox_agent.py | 5 ++-- 7 files changed, 49 insertions(+), 11 deletions(-) create mode 100644 veadk/agents/_sandbox_timeout.py diff --git a/pyproject.toml b/pyproject.toml index 451616f33..f954cde94 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ authors = [ {name = "Meng Wang", email = "mengwangwm@gmail.com"} ] dependencies = [ + "async-timeout>=5,<6; python_version < '3.11'", # asyncio.timeout backport "pydantic-settings==2.10.1", # Config management "a2a-sdk>=0.3.7,<1.0.0", # For Google Agent2Agent protocol "deprecated==1.2.18", diff --git a/tests/agents/test_remote_sandbox_agent.py b/tests/agents/test_remote_sandbox_agent.py index 1353e0f4b..ce346edfd 100644 --- a/tests/agents/test_remote_sandbox_agent.py +++ b/tests/agents/test_remote_sandbox_agent.py @@ -31,6 +31,7 @@ from pydantic import PrivateAttr from veadk import Agent +from veadk.agents._sandbox_timeout import timeout from veadk.agents.agentkit_remote_sandbox_agent import ( AgentkitRemoteSandboxAgent, SandboxAgentError, @@ -476,7 +477,8 @@ async def observe(event): @pytest.mark.asyncio -async def test_cancelling_code_invocation_interrupts_remote_turn(): +@pytest.mark.parametrize("deadline", [False, True]) +async def test_cancelling_code_invocation_interrupts_remote_turn(deadline): fixture = CodeFixture() app = web.Application() app.router.add_route("*", "/{path:.*}", fixture.handle) @@ -485,6 +487,7 @@ async def test_cancelling_code_invocation_interrupts_remote_turn(): AgentkitRemoteSandboxAgent( name="sandbox", tool_type="CodeEnv", + request_timeout=1 if deadline else 900, endpoint=str(server.make_url("/?route=fixture")), ) ) @@ -496,9 +499,15 @@ async def observe(event): running = asyncio.create_task(collect(runner, observe)) await asyncio.wait_for(seen.wait(), 5) - running.cancel() - with pytest.raises(asyncio.CancelledError): - await running + if deadline: + events = await asyncio.wait_for(running, 5) + assert any( + "TimeoutError" in (event.error_message or "") for event in events + ) + else: + running.cancel() + with pytest.raises(asyncio.CancelledError): + await running assert fixture.cancelled fixture.release.set() @@ -558,7 +567,7 @@ async def generate_content_async(self, llm_request, stream=False): ) serving = asyncio.create_task(server.serve()) try: - async with asyncio.timeout(10): + async with timeout(10): while not server.started: if serving.done(): await serving diff --git a/uv.lock b/uv.lock index 3cb43effb..c368fe785 100644 --- a/uv.lock +++ b/uv.lock @@ -5786,6 +5786,7 @@ dependencies = [ { name = "agent-pilot-sdk" }, { name = "agentkit-sdk-python" }, { name = "aiomysql" }, + { name = "async-timeout", marker = "python_full_version < '3.11'" }, { name = "asyncpg" }, { name = "cookiecutter" }, { name = "deprecated" }, @@ -5895,6 +5896,7 @@ requires-dist = [ { name = "agentkit-sdk-python", marker = "extra == 'harness-sidecar'", specifier = ">=0.8.1,<0.9.0" }, { name = "aiomysql", specifier = "==0.3.2" }, { name = "anthropic", marker = "extra == 'sandbox'", specifier = ">=0.40.0" }, + { name = "async-timeout", marker = "python_full_version < '3.11'", specifier = ">=5,<6" }, { name = "asyncpg", specifier = ">=0.29.0" }, { name = "cookiecutter", specifier = "==2.6.0" }, { name = "cozeloop", marker = "extra == 'extensions'", specifier = ">=0.1.21" }, diff --git a/veadk/agents/_sandbox_code.py b/veadk/agents/_sandbox_code.py index 0d091bd3a..63d89da7d 100644 --- a/veadk/agents/_sandbox_code.py +++ b/veadk/agents/_sandbox_code.py @@ -24,6 +24,7 @@ from google.adk.events import Event from google.genai import types +from veadk.agents._sandbox_timeout import timeout from veadk.agents.agentkit_remote_sandbox_agent import SandboxAgentError, binding_key from veadk.tools.sandbox.codex_worker_client import CodexWorkerClient, CodexWorkerError @@ -84,7 +85,7 @@ async def code_events( tools = {} async with CodexWorkerClient(endpoint, headers=headers) as client: try: - async with asyncio.timeout(agent.ready_timeout): + async with timeout(agent.ready_timeout): while True: try: ready = await client.request("GET", "/readyz") @@ -104,7 +105,7 @@ async def code_events( "CodeEnv Sidecar protocol is incompatible; version 1 with tool_events is required" ) check_lease(agent, lease) - async with asyncio.timeout(agent.request_timeout): + async with timeout(agent.request_timeout): sid = await client.create_session(binding) started = await client.start_turn(sid, task, turn_key) tid = started["turnId"] diff --git a/veadk/agents/_sandbox_skill.py b/veadk/agents/_sandbox_skill.py index 839d2c416..55e904d23 100644 --- a/veadk/agents/_sandbox_skill.py +++ b/veadk/agents/_sandbox_skill.py @@ -37,6 +37,7 @@ from google.adk.a2a.converters.part_converter import convert_a2a_part_to_genai_part from google.genai import types +from veadk.agents._sandbox_timeout import timeout from veadk.agents._sandbox_code import check_lease, event_for, observation from veadk.agents.agentkit_remote_sandbox_agent import SandboxAgentError @@ -56,7 +57,7 @@ async def skill_events( async with httpx.AsyncClient( headers=headers, timeout=httpx.Timeout(30, connect=10), trust_env=False ) as http: - async with asyncio.timeout(agent.ready_timeout): + async with timeout(agent.ready_timeout): while True: try: response = await http.get( @@ -215,7 +216,7 @@ def consume(response): return events try: - async with asyncio.timeout(agent.request_timeout): + async with timeout(agent.request_timeout): try: async for response in client.send_message( message, diff --git a/veadk/agents/_sandbox_timeout.py b/veadk/agents/_sandbox_timeout.py new file mode 100644 index 000000000..84fd1ea98 --- /dev/null +++ b/veadk/agents/_sandbox_timeout.py @@ -0,0 +1,23 @@ +# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Timeout context manager for all supported Python versions.""" + +try: + from asyncio import timeout +except ImportError: # Python 3.10 + from async_timeout import timeout + +__all__ = ["timeout"] diff --git a/veadk/agents/agentkit_remote_sandbox_agent.py b/veadk/agents/agentkit_remote_sandbox_agent.py index 1ac62b1bc..636c33dd3 100644 --- a/veadk/agents/agentkit_remote_sandbox_agent.py +++ b/veadk/agents/agentkit_remote_sandbox_agent.py @@ -28,6 +28,7 @@ from google.genai import types from pydantic import Field, PrivateAttr, SecretStr +from veadk.agents._sandbox_timeout import timeout from veadk.agents._sandbox_session import ensure_agentkit_session_lease from veadk.tools.builtin_tools._agentkit import ( get_agentkit_credentials, @@ -186,7 +187,7 @@ async def _run_async_impl(self, ctx): ) async with lock: try: - async with asyncio.timeout(self.ready_timeout * 2 + 60): + async with timeout(self.ready_timeout * 2 + 60): kind, endpoint, physical, lease = await self._resolve(ctx, logical) headers = await self._headers(ctx) user_event = next( @@ -213,7 +214,7 @@ async def _run_async_impl(self, ctx): previous = {"physical": physical} self._contexts[logical] = previous backend = code_events if kind == "CodeEnv" else skill_events - async with asyncio.timeout(self.request_timeout + self.ready_timeout): + async with timeout(self.request_timeout + self.ready_timeout): async with aclosing( backend( agent=self, From c99260d8879e8c418019fbc3bb77c2ad3431edb1 Mon Sep 17 00:00:00 2001 From: "hanzhi.421" Date: Sun, 6 Sep 2026 22:06:41 +0800 Subject: [PATCH 3/3] test(create-agent): isolate skill hydration cloud provider --- .../builtin_tools/create_agent/test_create_agent_toolset.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/tools/builtin_tools/create_agent/test_create_agent_toolset.py b/tests/tools/builtin_tools/create_agent/test_create_agent_toolset.py index 216554774..978bada90 100644 --- a/tests/tools/builtin_tools/create_agent/test_create_agent_toolset.py +++ b/tests/tools/builtin_tools/create_agent/test_create_agent_toolset.py @@ -983,6 +983,10 @@ def leaf_factory(node, tools, workflow_member, parent_agent): async def test_agentkit_skill_is_hydrated_only_when_selected( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: + # This fixture supplies Volcengine state credentials, regardless of any + # provider selected by other CLI tests in the same pytest worker. + monkeypatch.setenv("AGENTKIT_CLOUD_PROVIDER", "volcengine") + monkeypatch.setenv("CLOUD_PROVIDER", "volcengine") skill = Skill( name="private-writer", description="Write private reports",