diff --git a/tests/tools/builtin_tools/test_agentkit.py b/tests/tools/builtin_tools/test_agentkit.py index c5a5aa69e..5523adc5a 100644 --- a/tests/tools/builtin_tools/test_agentkit.py +++ b/tests/tools/builtin_tools/test_agentkit.py @@ -178,8 +178,8 @@ def test_builds_exec_bash_invoke_tool_request(self): command="echo hello", exec_dir="/tmp", env={"DEMO_ENV": "from-invoke-tool"}, - timeout=30, - hard_timeout=60, + timeout=120, + hard_timeout=300, max_output_length=30000, ttl=1800, ) @@ -196,11 +196,35 @@ def test_builds_exec_bash_invoke_tool_request(self): "command": "echo hello", "exec_dir": "/tmp", "env": {"DEMO_ENV": "from-invoke-tool"}, - "timeout": 30, - "hard_timeout": 60, + "timeout": 120, + "hard_timeout": 300, "max_output_length": 30000, }, ) + self.assertEqual(ve_request.call_args.kwargs["timeout"], (10.0, 330.0)) + + +class TestInvokeAgentkitRunCode(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.agentkit_module = _load_agentkit_module() + + def test_request_timeout_covers_code_execution_timeout(self): + with patch.object( + self.agentkit_module, + "ve_request", + return_value={"Result": {"Result": "code output"}}, + ) as ve_request: + result = self.agentkit_module.invoke_agentkit_run_code( + tool_id="code-tool", + tool_user_session_id="kk", + code="print('hello')", + timeout=120, + kernel_name="python3", + ) + + self.assertEqual(result, {"Result": {"Result": "code output"}}) + self.assertEqual(ve_request.call_args.kwargs["timeout"], (10.0, 150.0)) class TestEnsureAgentkitSessionEndpoint(unittest.TestCase): diff --git a/tests/utils/test_volcengine_sign.py b/tests/utils/test_volcengine_sign.py new file mode 100644 index 000000000..2763fea3b --- /dev/null +++ b/tests/utils/test_volcengine_sign.py @@ -0,0 +1,61 @@ +# 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 unittest.mock import patch + +from veadk.utils import volcengine_sign + + +def test_ve_request_uses_bounded_default_timeout(): + with patch.object( + volcengine_sign, + "request", + return_value={"Result": "ok"}, + ) as request: + volcengine_sign.ve_request( + request_body={}, + action="ListTools", + ak="ak", + sk="sk", + service="agentkit", + version="2025-10-30", + region="cn-beijing", + host="agentkit.cn-beijing.volces.com", + ) + + assert request.call_args.kwargs["timeout"] == ( + volcengine_sign.DEFAULT_REQUEST_TIMEOUT + ) + + +def test_ve_request_forwards_custom_timeout(): + with patch.object( + volcengine_sign, + "request", + return_value={"Result": "ok"}, + ) as request: + result = volcengine_sign.ve_request( + request_body={}, + action="InvokeTool", + ak="ak", + sk="sk", + service="agentkit", + version="2025-10-30", + region="cn-beijing", + host="agentkit.cn-beijing.volces.com", + timeout=(10.0, 150.0), + ) + + assert result == {"Result": "ok"} + assert request.call_args.kwargs["timeout"] == (10.0, 150.0) diff --git a/veadk/tools/builtin_tools/_agentkit.py b/veadk/tools/builtin_tools/_agentkit.py index 10abdeec1..80b9e5830 100644 --- a/veadk/tools/builtin_tools/_agentkit.py +++ b/veadk/tools/builtin_tools/_agentkit.py @@ -28,6 +28,20 @@ _SESSION_READY_TIMEOUT = 120.0 _SESSION_POLL_INTERVAL = 1.0 _SESSION_TERMINAL_STATUSES = frozenset({"failed", "terminating", "terminated"}) +_AGENTKIT_REQUEST_CONNECT_TIMEOUT = 10.0 +_AGENTKIT_REQUEST_MIN_READ_TIMEOUT = 60.0 +_AGENTKIT_REQUEST_TIMEOUT_BUFFER = 30.0 + + +def _agentkit_request_timeout(operation_timeout: int) -> tuple[float, float]: + """Keep the synchronous request alive longer than the tool operation.""" + return ( + _AGENTKIT_REQUEST_CONNECT_TIMEOUT, + max( + _AGENTKIT_REQUEST_MIN_READ_TIMEOUT, + float(operation_timeout) + _AGENTKIT_REQUEST_TIMEOUT_BUFFER, + ), + ) def resolve_agentkit_tool_id(*preferred_env_names: str) -> str: @@ -158,6 +172,7 @@ def invoke_agentkit_run_code( host=host, header=header, scheme=scheme, + timeout=_agentkit_request_timeout(timeout), ) @@ -200,6 +215,8 @@ def invoke_agentkit_exec_bash( if ttl is not None: request_body["Ttl"] = ttl + operation_timeout = max(timeout, hard_timeout or timeout) + return ve_request( request_body=request_body, action="InvokeTool", @@ -211,6 +228,7 @@ def invoke_agentkit_exec_bash( host=host, header=header, scheme=scheme, + timeout=_agentkit_request_timeout(operation_timeout), ) diff --git a/veadk/utils/volcengine_sign.py b/veadk/utils/volcengine_sign.py index 0fcc94635..6412247e3 100644 --- a/veadk/utils/volcengine_sign.py +++ b/veadk/utils/volcengine_sign.py @@ -353,6 +353,7 @@ def ve_request( method: Literal["GET", "POST", "PUT", "DELETE"] = "POST", scheme: Literal["http", "https"] = "https", session_token: str = "", + timeout: float | tuple[float, float] | None = DEFAULT_REQUEST_TIMEOUT, ): global Service Service = service @@ -387,6 +388,7 @@ def ve_request( action, json.dumps(request_body), Scheme, + timeout=timeout, ) return response_body except Exception as e: