diff --git a/tests/tools/builtin_tools/test_invoke_poll_skill.py b/tests/tools/builtin_tools/test_invoke_poll_skill.py new file mode 100644 index 000000000..3f0f1e342 --- /dev/null +++ b/tests/tools/builtin_tools/test_invoke_poll_skill.py @@ -0,0 +1,239 @@ +# 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. + +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +def _tool_context(): + invocation_context = types.SimpleNamespace( + session=types.SimpleNamespace(id="session-1"), + agent=types.SimpleNamespace(name="agent"), + user_id="user", + ) + return types.SimpleNamespace( + state={"state-key": "state-value"}, + _invocation_context=invocation_context, + ) + + +def _load_skill_task_module( + module_name, + *, + ensure_agentkit_session_endpoint=lambda **_kwargs: "https://sandbox.test", + post_a2a_jsonrpc=lambda **_kwargs: { + "result": {"kind": "task", "id": "task-1", "status": {"state": "working"}} + }, + inbound_auth_token=lambda _tool_context: "inbound-token", +): + module_path = ( + Path(__file__).resolve().parents[3] + / "veadk" + / "tools" + / "builtin_tools" + / f"{module_name}.py" + ) + + fake_google = types.ModuleType("google") + fake_google.__path__ = [] # type: ignore[attr-defined] + fake_google_adk = types.ModuleType("google.adk") + fake_google_adk.__path__ = [] # type: ignore[attr-defined] + fake_google_adk_tools = types.ModuleType("google.adk.tools") + fake_google_adk_tools.ToolContext = object + + fake_veadk = types.ModuleType("veadk") + fake_veadk.__path__ = [] # type: ignore[attr-defined] + fake_tools = types.ModuleType("veadk.tools") + fake_tools.__path__ = [] # type: ignore[attr-defined] + fake_builtin_tools = types.ModuleType("veadk.tools.builtin_tools") + fake_builtin_tools.__path__ = [] # type: ignore[attr-defined] + fake_agentkit = types.ModuleType("veadk.tools.builtin_tools._agentkit") + fake_agentkit.resolve_agentkit_tool_id = lambda _name: "test-tool" + fake_agentkit.ensure_agentkit_session_endpoint = ensure_agentkit_session_endpoint + + fake_execute_skills = types.ModuleType("veadk.tools.builtin_tools.execute_skills") + fake_execute_skills._A2A_HISTORY_LENGTH = 20 + fake_execute_skills._a2a_request_timeout = lambda _deadline: 60 + fake_execute_skills._a2a_result_task = lambda _operation, response: response[ + "result" + ] + fake_execute_skills._inbound_auth_token = inbound_auth_token + fake_execute_skills._post_a2a_jsonrpc = post_a2a_jsonrpc + fake_execute_skills._tool_user_session_id = ( + lambda tool_context: "agent_user_session-1" + ) + + def fake_validate_timeout(timeout): + if type(timeout) is not int or not 1 <= timeout <= 1800: + raise ValueError("timeout must be an integer between 1 and 1800 seconds") + + fake_execute_skills._validate_timeout = fake_validate_timeout + + stub_modules = { + "google": fake_google, + "google.adk": fake_google_adk, + "google.adk.tools": fake_google_adk_tools, + "veadk": fake_veadk, + "veadk.tools": fake_tools, + "veadk.tools.builtin_tools": fake_builtin_tools, + "veadk.tools.builtin_tools._agentkit": fake_agentkit, + "veadk.tools.builtin_tools.execute_skills": fake_execute_skills, + } + + with patch.dict(sys.modules, stub_modules): + spec = importlib.util.spec_from_file_location( + f"test_{module_name}_module", module_path + ) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class TestInvokeSkill(unittest.TestCase): + def test_invoke_skill_posts_nonblocking_message_send_and_returns_task(self): + session_kwargs = [] + jsonrpc_calls = [] + + module = _load_skill_task_module( + "invoke_skill", + ensure_agentkit_session_endpoint=lambda **kwargs: ( + session_kwargs.append(kwargs) or "https://sandbox.test" + ), + post_a2a_jsonrpc=lambda **kwargs: ( + jsonrpc_calls.append(kwargs) + or { + "result": { + "kind": "task", + "id": "task-1", + "status": {"state": "working"}, + } + } + ), + ) + + task = module.invoke_skill("do work", tool_context=_tool_context(), timeout=30) + + self.assertEqual( + {"kind": "task", "id": "task-1", "status": {"state": "working"}}, + task, + ) + self.assertEqual(1, len(session_kwargs)) + self.assertEqual("test-tool", session_kwargs[0]["tool_id"]) + self.assertEqual( + "agent_user_session-1", session_kwargs[0]["tool_user_session_id"] + ) + self.assertEqual({"state-key": "state-value"}, session_kwargs[0]["tool_state"]) + self.assertEqual(1800, session_kwargs[0]["ttl"]) + self.assertTrue(session_kwargs[0]["wait_until_ready"]) + + self.assertEqual(1, len(jsonrpc_calls)) + call = jsonrpc_calls[0] + self.assertEqual("https://sandbox.test", call["endpoint"]) + self.assertEqual(60, call["timeout"]) + self.assertEqual("inbound-token", call["inbound_auth"]) + payload = call["payload"] + self.assertEqual("2.0", payload["jsonrpc"]) + self.assertEqual("message/send", payload["method"]) + self.assertEqual("user", payload["params"]["metadata"]["user_id"]) + self.assertEqual("session-1", payload["params"]["metadata"]["session_id"]) + self.assertFalse(payload["params"]["configuration"]["blocking"]) + self.assertEqual(20, payload["params"]["configuration"]["historyLength"]) + self.assertEqual("do work", payload["params"]["message"]["parts"][0]["text"]) + + def test_invoke_skill_requires_tool_context(self): + module = _load_skill_task_module("invoke_skill") + + with self.assertRaisesRegex(ValueError, "tool_context is required"): + module.invoke_skill("do work", tool_context=None) + + def test_invoke_skill_wraps_endpoint_resolution_error(self): + module = _load_skill_task_module( + "invoke_skill", + ensure_agentkit_session_endpoint=lambda **_kwargs: (_ for _ in ()).throw( + RuntimeError("session failed") + ), + ) + + with self.assertRaisesRegex( + RuntimeError, "AgentKit session endpoint is not available" + ): + module.invoke_skill("do work", tool_context=_tool_context()) + + +class TestPollSkill(unittest.TestCase): + def test_poll_skill_posts_one_tasks_get_request_and_returns_task(self): + session_kwargs = [] + jsonrpc_calls = [] + + module = _load_skill_task_module( + "poll_skill", + ensure_agentkit_session_endpoint=lambda **kwargs: ( + session_kwargs.append(kwargs) or "https://sandbox.test" + ), + post_a2a_jsonrpc=lambda **kwargs: ( + jsonrpc_calls.append(kwargs) + or { + "result": { + "kind": "task", + "id": "task-1", + "status": {"state": "completed"}, + } + } + ), + ) + + task = module.poll_skill("task-1", tool_context=_tool_context(), timeout=30) + + self.assertEqual( + {"kind": "task", "id": "task-1", "status": {"state": "completed"}}, + task, + ) + self.assertEqual(1, len(session_kwargs)) + self.assertEqual("test-tool", session_kwargs[0]["tool_id"]) + self.assertEqual(1, len(jsonrpc_calls)) + call = jsonrpc_calls[0] + self.assertEqual("https://sandbox.test", call["endpoint"]) + self.assertEqual(60, call["timeout"]) + self.assertEqual("inbound-token", call["inbound_auth"]) + payload = call["payload"] + self.assertEqual("tasks/get", payload["method"]) + self.assertEqual("task-1", payload["params"]["id"]) + self.assertEqual(20, payload["params"]["historyLength"]) + + def test_poll_skill_rejects_empty_task_id_before_endpoint_resolution(self): + module = _load_skill_task_module( + "poll_skill", + ensure_agentkit_session_endpoint=lambda **_kwargs: self.fail( + "task_id must be rejected before endpoint resolution" + ), + ) + + with self.assertRaisesRegex(ValueError, "task_id is required"): + module.poll_skill("", tool_context=_tool_context()) + + def test_poll_skill_requires_tool_context(self): + module = _load_skill_task_module("poll_skill") + + with self.assertRaisesRegex(ValueError, "tool_context is required"): + module.poll_skill("task-1", tool_context=None) + + +if __name__ == "__main__": + unittest.main() diff --git a/veadk/tools/builtin_tools/invoke_skill.py b/veadk/tools/builtin_tools/invoke_skill.py new file mode 100644 index 000000000..fed3bf8e9 --- /dev/null +++ b/veadk/tools/builtin_tools/invoke_skill.py @@ -0,0 +1,94 @@ +# 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. + +from __future__ import annotations + +import time +import uuid + +from google.adk.tools import ToolContext + +from veadk.tools.builtin_tools._agentkit import ( + ensure_agentkit_session_endpoint, + resolve_agentkit_tool_id, +) +from veadk.tools.builtin_tools.execute_skills import ( + _A2A_HISTORY_LENGTH, + _a2a_request_timeout, + _a2a_result_task, + _inbound_auth_token, + _post_a2a_jsonrpc, + _tool_user_session_id, + _validate_timeout, +) + + +def invoke_skill( + workflow_prompt: str, + tool_context: ToolContext = None, + timeout: int = 1800, +) -> dict: + """Create an A2A skill task and return the initial task object.""" + if tool_context is None: + raise ValueError("tool_context is required for invoke_skill") + _validate_timeout(timeout) + + tool_id = resolve_agentkit_tool_id("AGENTKIT_TOOL_ID_SKILLS") + try: + endpoint = ensure_agentkit_session_endpoint( + tool_id=tool_id, + tool_user_session_id=_tool_user_session_id(tool_context), + tool_state=tool_context.state, + ttl=max(timeout, 1800), + wait_until_ready=True, + ) + except Exception as exc: + raise RuntimeError( + f"AgentKit session endpoint is not available: {exc}" + ) from exc + + invocation_context = tool_context._invocation_context + deadline = time.monotonic() + timeout + inbound_auth = _inbound_auth_token(tool_context) + + return _a2a_result_task( + "A2ASendMessage", + _post_a2a_jsonrpc( + endpoint=endpoint, + payload={ + "jsonrpc": "2.0", + "id": uuid.uuid4().hex, + "method": "message/send", + "params": { + "message": { + "kind": "message", + "messageId": uuid.uuid4().hex, + "role": "user", + "parts": [{"kind": "text", "text": workflow_prompt}], + }, + "metadata": { + "user_id": invocation_context.user_id, + "session_id": invocation_context.session.id, + }, + "configuration": { + "blocking": False, + "historyLength": _A2A_HISTORY_LENGTH, + }, + }, + }, + timeout=_a2a_request_timeout(deadline), + retry_until=deadline, + inbound_auth=inbound_auth, + ), + ) diff --git a/veadk/tools/builtin_tools/poll_skill.py b/veadk/tools/builtin_tools/poll_skill.py new file mode 100644 index 000000000..de7614924 --- /dev/null +++ b/veadk/tools/builtin_tools/poll_skill.py @@ -0,0 +1,81 @@ +# 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. + +from __future__ import annotations + +import time +import uuid + +from google.adk.tools import ToolContext + +from veadk.tools.builtin_tools._agentkit import ( + ensure_agentkit_session_endpoint, + resolve_agentkit_tool_id, +) +from veadk.tools.builtin_tools.execute_skills import ( + _A2A_HISTORY_LENGTH, + _a2a_request_timeout, + _a2a_result_task, + _inbound_auth_token, + _post_a2a_jsonrpc, + _tool_user_session_id, + _validate_timeout, +) + + +def poll_skill( + task_id: str, + tool_context: ToolContext = None, + timeout: int = 1800, +) -> dict: + """Fetch one A2A skill task status snapshot.""" + if tool_context is None: + raise ValueError("tool_context is required for poll_skill") + if not isinstance(task_id, str) or not task_id: + raise ValueError("task_id is required for poll_skill") + _validate_timeout(timeout) + + tool_id = resolve_agentkit_tool_id("AGENTKIT_TOOL_ID_SKILLS") + try: + endpoint = ensure_agentkit_session_endpoint( + tool_id=tool_id, + tool_user_session_id=_tool_user_session_id(tool_context), + tool_state=tool_context.state, + ttl=max(timeout, 1800), + wait_until_ready=True, + ) + except Exception as exc: + raise RuntimeError( + f"AgentKit session endpoint is not available: {exc}" + ) from exc + + deadline = time.monotonic() + timeout + return _a2a_result_task( + "A2AGetTask", + _post_a2a_jsonrpc( + endpoint=endpoint, + payload={ + "jsonrpc": "2.0", + "id": uuid.uuid4().hex, + "method": "tasks/get", + "params": { + "id": task_id, + "historyLength": _A2A_HISTORY_LENGTH, + }, + }, + timeout=_a2a_request_timeout(deadline), + retry_until=deadline, + inbound_auth=_inbound_auth_token(tool_context), + ), + )