From 2be34ee790518d4dec4f41e462da842b26d5d2b8 Mon Sep 17 00:00:00 2001 From: "liyi.ly" Date: Fri, 4 Sep 2026 19:50:52 +0800 Subject: [PATCH] fix: bound outbound I/O and async waits --- .../docs/framework/tools/builtin.en.mdx | 1 + docs/content/docs/framework/tools/builtin.mdx | 1 + .../environment-variables.en.mdx | 24 + .../configuration/environment-variables.mdx | 24 + tests/a2a/test_hub_client_timeouts.py | 134 ++++ tests/a2a/test_ve_task_store.py | 153 +++++ tests/community/test_viking_memory_store.py | 161 +++++ .../integrations/test_integration_timeouts.py | 163 +++++ tests/knowledgebase/test_backend_timeouts.py | 121 ++++ tests/runner/test_runner_contract.py | 5 +- tests/test_no_unbounded_http_calls.py | 577 ++++++++++++++++++ tests/test_runner.py | 69 +++ .../test_resource_source_deadlines.py | 405 ++++++++++++ .../builtin_tools/test_image_edit_bounds.py | 443 ++++++++++++++ .../builtin_tools/test_llm_shield_contract.py | 296 +++++++++ .../builtin_tools/test_mobile_run_bounds.py | 489 +++++++++++++++ tests/tools/builtin_tools/test_tts_bounds.py | 362 +++++++++++ .../skills_tools/test_skills_timeouts.py | 166 +++++ tests/utils/test_http_defaults.py | 262 ++++++++ tests/utils/test_misc_download_bounds.py | 291 +++++++++ veadk/a2a/hub/a2a_hub_client.py | 14 +- veadk/a2a/remote_ve_agent.py | 2 + veadk/a2a/ve_task_store.py | 41 +- veadk/agent.py | 9 + veadk/cli/cli_uploadevalset.py | 2 + .../store/memory/viking_memory.py | 10 +- veadk/integrations/ve_apig/ve_apig_utils.py | 3 + .../ve_code_pipeline/ve_code_pipeline.py | 2 + veadk/integrations/ve_cozeloop/ve_cozeloop.py | 9 +- veadk/integrations/ve_faas/ve_faas_utils.py | 3 + .../backends/context_search_backend.py | 14 +- .../backends/vikingdb_knowledge_backend.py | 2 + veadk/runner.py | 15 +- .../sources/agentkit_knowledge.py | 66 +- .../create_agent/sources/skills.py | 108 +++- veadk/tools/builtin_tools/image_edit.py | 25 +- veadk/tools/builtin_tools/llm_shield.py | 184 +++++- veadk/tools/builtin_tools/mobile_run.py | 90 ++- veadk/tools/builtin_tools/tts.py | 85 ++- .../skills_tools/download_skills_tool.py | 9 +- .../skills_tools/register_skills_tool.py | 5 +- veadk/tools/skills_tools/skills_tool.py | 15 +- veadk/utils/http_defaults.py | 69 +++ veadk/utils/misc.py | 102 +++- veadk/utils/volcengine_sign.py | 4 +- 45 files changed, 4918 insertions(+), 117 deletions(-) create mode 100644 tests/a2a/test_hub_client_timeouts.py create mode 100644 tests/a2a/test_ve_task_store.py create mode 100644 tests/community/test_viking_memory_store.py create mode 100644 tests/integrations/test_integration_timeouts.py create mode 100644 tests/knowledgebase/test_backend_timeouts.py create mode 100644 tests/test_no_unbounded_http_calls.py create mode 100644 tests/tools/builtin_tools/create_agent/test_resource_source_deadlines.py create mode 100644 tests/tools/builtin_tools/test_image_edit_bounds.py create mode 100644 tests/tools/builtin_tools/test_llm_shield_contract.py create mode 100644 tests/tools/builtin_tools/test_mobile_run_bounds.py create mode 100644 tests/tools/builtin_tools/test_tts_bounds.py create mode 100644 tests/tools/skills_tools/test_skills_timeouts.py create mode 100644 tests/utils/test_http_defaults.py create mode 100644 tests/utils/test_misc_download_bounds.py create mode 100644 veadk/utils/http_defaults.py diff --git a/docs/content/docs/framework/tools/builtin.en.mdx b/docs/content/docs/framework/tools/builtin.en.mdx index 881971503..f29cd201a 100644 --- a/docs/content/docs/framework/tools/builtin.en.mdx +++ b/docs/content/docs/framework/tools/builtin.en.mdx @@ -468,6 +468,7 @@ Environment variables: - `TOOL_VESPEECH_API_KEY`: VeSpeech service API key - `TOOL_VESPEECH_SPEAKER`: voice, defaults to `zh_female_vv_uranus_bigtts` - `TOOL_VESPEECH_AUDIO_OUTPUT_PATH`: audio output directory, defaults to the system temp directory +- `VEADK_HTTP_STREAM_BUDGET`: total wall-clock ceiling for consuming the streamed synthesis response, defaults to 300 seconds; synthesis aborts with an error once it is exceeded. See [Environment variables · HTTP timeouts](/en/docs/references/configuration/environment-variables#http-timeouts) ## Code sandboxes diff --git a/docs/content/docs/framework/tools/builtin.mdx b/docs/content/docs/framework/tools/builtin.mdx index a9086b09d..2cf59793a 100644 --- a/docs/content/docs/framework/tools/builtin.mdx +++ b/docs/content/docs/framework/tools/builtin.mdx @@ -465,6 +465,7 @@ if __name__ == "__main__": - `TOOL_VESPEECH_API_KEY`：VeSpeech 服务的 API Key - `TOOL_VESPEECH_SPEAKER`：音色，默认为 `zh_female_vv_uranus_bigtts` - `TOOL_VESPEECH_AUDIO_OUTPUT_PATH`：音频输出目录，默认为系统临时目录 +- `VEADK_HTTP_STREAM_BUDGET`：消费流式合成响应的墙钟总时长上限，默认 300 秒；超时后合成中止并报错。详见[环境变量 · HTTP 超时](/cn/docs/references/configuration/environment-variables#http-超时) ## 代码沙箱 diff --git a/docs/content/docs/references/configuration/environment-variables.en.mdx b/docs/content/docs/references/configuration/environment-variables.en.mdx index f9bdc9977..0318cabe4 100644 --- a/docs/content/docs/references/configuration/environment-variables.en.mdx +++ b/docs/content/docs/references/configuration/environment-variables.en.mdx @@ -187,6 +187,30 @@ Prefix `OBSERVABILITY_`, configuring where traces and metrics are reported. | | `OBSERVABILITY_PROMETHEUS_USERNAME` | Username | | | `OBSERVABILITY_PROMETHEUS_PASSWORD` | Password | +## HTTP timeouts + +The `VEADK_HTTP_*` variables bound outbound `requests` calls VeADK makes internally (built-in tools, knowledge base backends, deployment and integration SDKs). `VEADK_MAX_DOWNLOAD_BYTES` separately caps remote bodies downloaded by media and Skill tooling. `requests` has no default timeout of its own: a peer that accepts the connection and then goes silent blocks the call forever. + +| Variable | Default | Meaning | +| :- | :- | :- | +| `VEADK_HTTP_CONNECT_TIMEOUT` | `10` | Seconds allowed to establish a TCP/TLS connection. An unreachable peer fails fast instead of occupying a worker. | +| `VEADK_HTTP_READ_TIMEOUT` | `60` | Maximum seconds allowed **between two consecutive socket reads**. It caps silence, not the total duration of the call, so a large body is not cut off merely for taking a long time to transfer. | +| `VEADK_HTTP_STREAM_BUDGET` | `300` | **Total wall-clock ceiling** in seconds for consuming a streamed response end to end; used by TTS plus remote media and Skill downloads. This is a different quantity from the read timeout above: a server emitting an endless trickle of valid frames resets the read gap forever and is only caught by a total budget. | +| `VEADK_MAX_DOWNLOAD_BYTES` | `268435456` | Maximum size of one remote media or Skill download in bytes (default 256 MiB). Oversized bodies are stopped before replacing a destination file. | + +Timeout values are clamped to a minimum of `1.0` second and the byte limit to at least one byte. Invalid values fall back to their defaults. + + +These four variables are evaluated at module import time, so they **must come from the real process environment** (`export`, or environment variables injected by a container or FaaS platform). Setting them in `.env` or `config.yaml` has no effect — VeADK loads those into `os.environ` later than these constants are evaluated, so the "every variable can also be expressed in `config.yaml`" rule stated at the top of this page does not apply here. + + +```bash +export VEADK_HTTP_CONNECT_TIMEOUT=10 +export VEADK_HTTP_READ_TIMEOUT=60 +export VEADK_HTTP_STREAM_BUDGET=300 +export VEADK_MAX_DOWNLOAD_BYTES=268435456 +``` + ## Others | Variable | Meaning | diff --git a/docs/content/docs/references/configuration/environment-variables.mdx b/docs/content/docs/references/configuration/environment-variables.mdx index 4dd64cfd3..12330e1f6 100644 --- a/docs/content/docs/references/configuration/environment-variables.mdx +++ b/docs/content/docs/references/configuration/environment-variables.mdx @@ -185,6 +185,30 @@ database: | | `OBSERVABILITY_PROMETHEUS_USERNAME` | 用户名 | | | `OBSERVABILITY_PROMETHEUS_PASSWORD` | 密码 | +## HTTP 超时 + +`VEADK_HTTP_*` 用于约束 VeADK 内部基于 `requests` 的出站调用（内置工具、知识库后端、部署与集成 SDK）；`VEADK_MAX_DOWNLOAD_BYTES` 另行限制媒体和 Skill 工具下载的远程响应体大小。`requests` 自身没有默认超时：对端接受连接后不再返回数据时，调用会永久挂起。 + +| 环境变量 | 默认值 | 释义 | +| :- | :- | :- | +| `VEADK_HTTP_CONNECT_TIMEOUT` | `10` | 建立 TCP/TLS 连接的秒数上限。对端不可达时快速失败，而不是长时间占用一个 worker。 | +| `VEADK_HTTP_READ_TIMEOUT` | `60` | **两次相邻 socket 读取之间**允许的最大间隔秒数。它约束的是「静默」而非响应总时长，因此大响应体不会仅因为传输耗时长而被中断。 | +| `VEADK_HTTP_STREAM_BUDGET` | `300` | 完整消费一个流式响应的**墙钟总时长**上限（秒），用于 TTS、远程媒体及 Skill 下载。这与上面的读超时是两个不同的量：持续吐出合法分片的服务端会无限刷新读间隔，只有总预算能兜住这种情况。 | +| `VEADK_MAX_DOWNLOAD_BYTES` | `268435456` | 单次远程媒体或 Skill 下载的最大字节数，默认 256 MiB。超限响应会在替换目标文件前终止。 | + +超时值最小为 `1.0` 秒，字节上限最小为 1 字节；无法解析时回退到各自默认值。 + + +这四个变量在模块导入时求值，因此**必须来自真实的进程环境变量**（`export`、容器或函数计算注入的环境变量等）。写在 `.env` 或 `config.yaml` 中不会生效——VeADK 将它们写入 `os.environ` 的时机晚于这些常量的求值，此处不适用本页开头「环境变量与 `config.yaml` 等价」的约定。 + + +```bash +export VEADK_HTTP_CONNECT_TIMEOUT=10 +export VEADK_HTTP_READ_TIMEOUT=60 +export VEADK_HTTP_STREAM_BUDGET=300 +export VEADK_MAX_DOWNLOAD_BYTES=268435456 +``` + ## 其他 | 环境变量 | 释义 | diff --git a/tests/a2a/test_hub_client_timeouts.py b/tests/a2a/test_hub_client_timeouts.py new file mode 100644 index 000000000..49226d3d5 --- /dev/null +++ b/tests/a2a/test_hub_client_timeouts.py @@ -0,0 +1,134 @@ +# 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. + +"""Regression tests: A2A hub/remote-agent HTTP calls must be bounded. + +The hub liveness probe deliberately runs on a tighter read budget than the +ordinary calls -- a hub that has gone silent should be reported as down +quickly, not after a full minute. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest +from a2a.types import AgentCapabilities, AgentCard + +from veadk.utils.http_defaults import DEFAULT_CONNECT_TIMEOUT, DEFAULT_HTTP_TIMEOUT + + +def _agent_card() -> AgentCard: + return AgentCard( + name="weather-agent", + description="Weather agent", + url="http://127.0.0.1:8000", + version="1.0.0", + capabilities=AgentCapabilities(), + default_input_modes=["text/plain"], + default_output_modes=["text/plain"], + skills=[], + ) + + +def _response(payload: dict | None = None, status_code: int = 200) -> MagicMock: + response = MagicMock() + response.status_code = status_code + response.json.return_value = payload if payload is not None else {} + return response + + +def test_health_check_read_budget_is_tighter_than_default() -> None: + from veadk.a2a.hub.a2a_hub_client import HEALTH_CHECK_TIMEOUT + + assert HEALTH_CHECK_TIMEOUT[0] == DEFAULT_CONNECT_TIMEOUT + assert HEALTH_CHECK_TIMEOUT[1] < DEFAULT_HTTP_TIMEOUT[1] + + +def test_health_check_passes_health_check_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.a2a.hub import a2a_hub_client as module + + get = MagicMock(return_value=_response()) + monkeypatch.setattr(module.requests, "get", get) + + module.A2AHubClient(server_host="127.0.0.1", server_port=8888) + + assert get.call_count == 1 + assert get.call_args.kwargs["timeout"] == module.HEALTH_CHECK_TIMEOUT + + +def test_get_agent_cards_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.a2a.hub import a2a_hub_client as module + + get = MagicMock(return_value=_response({"agent_infos": [{"agent_id": "a1"}]})) + monkeypatch.setattr(module.requests, "get", get) + + client = module.A2AHubClient(server_host="127.0.0.1", server_port=8888) + assert client.get_agent_cards(group_id="g1") == [{"agent_id": "a1"}] + + # First call is the constructor's health check. + assert get.call_count == 2 + assert get.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_register_agent_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.a2a.hub import a2a_hub_client as module + + monkeypatch.setattr(module.requests, "get", MagicMock(return_value=_response())) + post = MagicMock(return_value=_response()) + monkeypatch.setattr(module.requests, "post", post) + + client = module.A2AHubClient(server_host="127.0.0.1", server_port=8888) + client.register_agent(group_id="g1", agent_id="a1", agent_card=_agent_card()) + + assert post.call_count == 1 + assert post.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_create_group_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.a2a.hub import a2a_hub_client as module + + monkeypatch.setattr(module.requests, "get", MagicMock(return_value=_response())) + post = MagicMock(return_value=_response()) + monkeypatch.setattr(module.requests, "post", post) + + client = module.A2AHubClient(server_host="127.0.0.1", server_port=8888) + client.create_group(group_id="g1") + + assert post.call_count == 1 + assert post.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_remote_ve_agent_card_fetch_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.a2a import remote_ve_agent as module + + card = _agent_card().model_dump(mode="json", by_alias=True) + get = MagicMock(return_value=_response(card)) + monkeypatch.setattr(module.requests, "get", get) + + agent = module.RemoteVeAgent(name="weather_agent", url="http://127.0.0.1:8000") + + assert agent.name == "weather_agent" + assert get.call_count == 1 + assert get.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT diff --git a/tests/a2a/test_ve_task_store.py b/tests/a2a/test_ve_task_store.py new file mode 100644 index 000000000..fdfa63e8d --- /dev/null +++ b/tests/a2a/test_ve_task_store.py @@ -0,0 +1,153 @@ +# 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. + +"""Regression tests: `VeTaskStore` must honour the A2A `TaskStore` contract. + +Two things used to be wrong. The overrides dropped the `context` parameter, +which the A2A request handlers pass *positionally* +(`await self.task_store.get(params.id, context)`), so every call raised +`TypeError`. And the bodies were `return None` stubs, so a store that did get +called would have silently discarded every task. + +The signature expectations below are derived from `TaskStore` itself rather +than hardcoded, so an upstream parameter change fails these tests instead of +letting the override quietly drift out of sync again. +""" + +from __future__ import annotations + +import inspect + +import pytest +from a2a.server.context import ServerCallContext +from a2a.server.tasks import TaskStore +from a2a.types import Task, TaskState, TaskStatus + +from veadk.a2a.ve_task_store import VeTaskStore + +TASK_STORE_METHODS = ["save", "get", "delete"] + + +def _call_contract(func) -> list[tuple]: + """The part of a signature a caller must satisfy: name, kind, default.""" + return [ + (param.name, param.kind, param.default) + for param in inspect.signature(func).parameters.values() + ] + + +def _task(task_id: str, state: TaskState = TaskState.submitted) -> Task: + return Task( + id=task_id, + context_id="test-context", + status=TaskStatus(state=state), + ) + + +@pytest.mark.parametrize("method_name", TASK_STORE_METHODS) +def test_override_matches_base_call_contract(method_name: str) -> None: + """Each override accepts exactly what the base class promises callers.""" + expected = _call_contract(getattr(TaskStore, method_name)) + actual = _call_contract(getattr(VeTaskStore, method_name)) + + assert actual == expected + + +@pytest.mark.asyncio +async def test_context_accepted_positionally() -> None: + """A2A's request handlers and task manager pass `context` positionally.""" + store = VeTaskStore() + context = ServerCallContext() + task = _task("positional-task") + + await store.save(task, context) + assert await store.get(task.id, context) == task + + await store.delete(task.id, context) + assert await store.get(task.id, context) is None + + +@pytest.mark.asyncio +async def test_context_accepted_by_keyword() -> None: + store = VeTaskStore() + context = ServerCallContext() + task = _task("keyword-task") + + await store.save(task, context=context) + assert await store.get(task.id, context=context) == task + + await store.delete(task.id, context=context) + assert await store.get(task.id, context=context) is None + + +@pytest.mark.asyncio +async def test_context_is_optional() -> None: + """`context` defaults to None, so callers may omit it entirely.""" + store = VeTaskStore() + task = _task("no-context-task") + + await store.save(task) + assert await store.get(task.id) == task + + await store.delete(task.id) + assert await store.get(task.id) is None + + +@pytest.mark.asyncio +async def test_save_updates_existing_task() -> None: + """Re-saving an id replaces the stored task rather than duplicating it.""" + store = VeTaskStore() + await store.save(_task("task-1", TaskState.submitted)) + await store.save(_task("task-1", TaskState.completed)) + + stored = await store.get("task-1") + + assert stored is not None + assert stored.status.state == TaskState.completed + + +@pytest.mark.asyncio +async def test_get_returns_none_only_for_unknown_ids() -> None: + """A miss must mean "absent", not "this store never returns anything".""" + store = VeTaskStore() + task = _task("stored-task") + await store.save(task) + + assert await store.get("never-saved") is None + assert await store.get(task.id) == task + + +@pytest.mark.asyncio +async def test_delete_of_unknown_id_leaves_other_tasks_intact() -> None: + """Deleting a nonexistent task is tolerated and touches nothing else.""" + store = VeTaskStore() + task = _task("survivor-task") + await store.save(task) + + await store.delete("never-saved") + + assert await store.get(task.id) == task + + +@pytest.mark.asyncio +async def test_tasks_are_not_shared_between_instances() -> None: + """Each store owns its tasks; nothing leaks through class-level state.""" + first = VeTaskStore() + second = VeTaskStore() + task = _task("instance-task") + + await first.save(task) + + assert await first.get(task.id) == task + assert await second.get(task.id) is None diff --git a/tests/community/test_viking_memory_store.py b/tests/community/test_viking_memory_store.py new file mode 100644 index 000000000..ad5aef2b2 --- /dev/null +++ b/tests/community/test_viking_memory_store.py @@ -0,0 +1,161 @@ +# Copyright (c) 2026 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. + +"""Regression tests for `VikingMemoryStore.abatch`. + +`abatch` used to be a *synchronous* stub with a `...` body overriding +`BaseStore.abatch`, which is an `async def` abstract method. Every async +accessor on the base (`aget`/`asearch`/`aput`/`adelete`/`alist_namespaces`) +does a bare `await self.abatch(...)`, so the stub made all of them raise +`TypeError: object NoneType can't be used in 'await' expression`. +""" + +import inspect +import json +import threading + +import pytest + +pytest.importorskip("langgraph") + +from langgraph.store.base import ( # noqa: E402 + BaseStore, + GetOp, + SearchOp, +) + +from veadk.community.langchain_ai.store.memory import viking_memory # noqa: E402 + +INDEX = "test_index" +USER_ID = "test_user" + + +class _FakeVikingBackend: + """Network-free stand-in for the (synchronous) `VikingDBLTMBackend`.""" + + def __init__(self, index: str = INDEX): + self.index = index + self.saved: list[dict] = [] + self.searched: list[dict] = [] + self.call_threads: list[int] = [] + + def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool: + self.call_threads.append(threading.get_ident()) + self.saved.append( + { + "user_id": user_id, + "session_id": kwargs.get("session_id"), + "event_strings": list(event_strings), + } + ) + return True + + def search_memory( + self, user_id: str, query: str, top_k: int, **kwargs + ) -> list[str]: + self.call_threads.append(threading.get_ident()) + self.searched.append({"user_id": user_id, "query": query, "top_k": top_k}) + return [f"memory about {query}"] + + +@pytest.fixture +def backend() -> _FakeVikingBackend: + return _FakeVikingBackend() + + +@pytest.fixture +def store(monkeypatch, backend: _FakeVikingBackend): + # `VikingDBLTMBackend.__init__` talks to VikingDB over the network, so the + # store's only collaborator is replaced before construction. + monkeypatch.setattr( + viking_memory, "VikingDBLTMBackend", lambda index: backend, raising=True + ) + return viking_memory.VikingMemoryStore(index=INDEX) + + +def test_abatch_is_a_coroutine_function(): + assert inspect.iscoroutinefunction(viking_memory.VikingMemoryStore.abatch) + + +def test_store_is_instantiable_with_both_batch_methods_defined(store): + # `batch`/`abatch` are the only abstract methods on `BaseStore`, so defining + # both is what makes the store constructible. A refactor that drops either + # one turns `VikingMemoryStore` back into an abstract class. + assert BaseStore.__abstractmethods__ == frozenset({"batch", "abatch"}) + assert "abatch" in viking_memory.VikingMemoryStore.__dict__ + assert "batch" in viking_memory.VikingMemoryStore.__dict__ + assert viking_memory.VikingMemoryStore.__abstractmethods__ == frozenset() + assert isinstance(store, BaseStore) + + +@pytest.mark.asyncio +async def test_asearch_returns_results(store, backend): + results = await store.asearch((INDEX, USER_ID), query="pizza", limit=3) + + assert results == ["memory about pizza"] + assert backend.searched == [ + {"user_id": USER_ID, "query": "pizza", "top_k": 3}, + ] + + +@pytest.mark.asyncio +async def test_aget_returns_a_result(store): + got = await store.aget((INDEX, USER_ID), "session-1") + + # `_apply_get_op` is still a placeholder; what matters here is that the + # async accessor returns the sync result instead of raising. + assert got is not None + assert got == store.get((INDEX, USER_ID), "session-1") + + +@pytest.mark.asyncio +async def test_aput_reaches_the_backend(store, backend): + event = {"role": "user", "parts": [{"text": "hello"}]} + + await store.aput((INDEX, USER_ID), "session-1", {"event-0": event}) + + assert backend.saved == [ + { + "user_id": USER_ID, + "session_id": "session-1", + "event_strings": [json.dumps(event)], + } + ] + + +@pytest.mark.asyncio +async def test_abatch_matches_batch(store): + # Read-only ops, so running them twice is side-effect free. + ops = [ + SearchOp(namespace_prefix=(INDEX, USER_ID), query="pizza", limit=2), + GetOp(namespace=(INDEX, USER_ID), key="session-1"), + ] + + assert await store.abatch(ops) == store.batch(ops) + + +@pytest.mark.asyncio +async def test_abatch_runs_off_the_event_loop(store, backend): + loop_thread = threading.get_ident() + + await store.asearch((INDEX, USER_ID), query="pizza") + + assert backend.call_threads + assert all(ident != loop_thread for ident in backend.call_threads) + + +def test_batch_stays_on_the_calling_thread(store, backend): + store.search((INDEX, USER_ID), query="pizza") + + assert backend.call_threads == [threading.get_ident()] diff --git a/tests/integrations/test_integration_timeouts.py b/tests/integrations/test_integration_timeouts.py new file mode 100644 index 000000000..4a837b52d --- /dev/null +++ b/tests/integrations/test_integration_timeouts.py @@ -0,0 +1,163 @@ +# 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. + +"""Regression tests: Volcengine integration helpers must be bounded. + +These are the signed control-plane calls (APIG, FaaS, CozeLoop) plus the GitHub +webhook registration. All of them are made synchronously from deployment paths, +so an unbounded `requests` call hangs the caller indefinitely. +""" + +from __future__ import annotations + +import datetime +from unittest.mock import MagicMock + +import pytest + +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT + + +def _response(payload: dict | None = None, status_code: int = 200) -> MagicMock: + response = MagicMock() + response.status_code = status_code + response.text = "" + response.json.return_value = payload if payload is not None else {} + return response + + +def test_ve_apig_request_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.integrations.ve_apig import ve_apig_utils as module + + request = MagicMock(return_value=_response({"Result": {}})) + monkeypatch.setattr(module.requests, "request", request) + + module.request( + method="GET", + date=datetime.datetime(2025, 1, 1, 0, 0, 0), + query={}, + header={}, + region="cn-beijing", + ak="ak", + sk="sk", + token="token", + action="ListRoutes", + body="", + ) + + assert request.call_count == 1 + assert request.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_ve_faas_request_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.integrations.ve_faas import ve_faas_utils as module + + request = MagicMock(return_value=_response({"Result": {}})) + monkeypatch.setattr(module.requests, "request", request) + + module.request( + method="GET", + date=datetime.datetime(2025, 1, 1, 0, 0, 0), + query={}, + header={}, + ak="ak", + sk="sk", + token="token", + action="ListGateways", + body="", + ) + + assert request.call_count == 1 + assert request.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_cozeloop_search_workspace_id_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.integrations.ve_cozeloop import ve_cozeloop as module + + get = MagicMock( + return_value=_response( + {"code": 0, "data": {"workspaces": [{"name": "veadk", "id": "ws-1"}]}} + ) + ) + monkeypatch.setattr(module.requests, "get", get) + + workspace_id = module.VeCozeloop(api_key="key").search_workspace_id( + workspace_name="veadk" + ) + + assert workspace_id == "ws-1" + assert get.call_count == 1 + assert get.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_cozeloop_create_workspace_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.integrations.ve_cozeloop import ve_cozeloop as module + + # A failing lookup is what pushes `create_workspace` past its early return + # and onto the POST we care about. + monkeypatch.setattr( + module.requests, + "get", + MagicMock(return_value=_response({"code": 1})), + ) + post = MagicMock(return_value=_response({"code": 0, "data": {"id": "ws-2"}})) + monkeypatch.setattr(module.requests, "post", post) + + workspace_id = module.VeCozeloop(api_key="key").create_workspace( + workspace_name="veadk" + ) + + assert workspace_id == "ws-2" + assert post.call_count == 1 + assert post.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_code_pipeline_github_webhook_passes_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.integrations.ve_code_pipeline import ve_code_pipeline as module + + post = MagicMock( + return_value=_response( + {"id": 1, "url": "https://api.github.com/hooks/1", "events": ["push"]}, + status_code=201, + ) + ) + monkeypatch.setattr(module.requests, "post", post) + + pipeline = module.VeCodePipeline( + volcengine_access_key="ak", + volcengine_secret_key="sk", + region="cn-beijing", + ) + result = pipeline._set_github_webhook( + webhook_url="https://cp.test/webhook", + github_url="https://github.com/owner/repo", + github_token="gh-token", + ) + + assert result is not None + assert post.call_count == 1 + assert post.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + assert ( + post.call_args.kwargs["url"] == "https://api.github.com/repos/owner/repo/hooks" + ) diff --git a/tests/knowledgebase/test_backend_timeouts.py b/tests/knowledgebase/test_backend_timeouts.py new file mode 100644 index 000000000..4bad1c146 --- /dev/null +++ b/tests/knowledgebase/test_backend_timeouts.py @@ -0,0 +1,121 @@ +# 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. + +"""Regression tests: knowledgebase backends must bound their outbound calls. + +Search, control-plane, and the pre-signed TOS upload all carry the same shared +`DEFAULT_HTTP_TIMEOUT`. There is no separate, longer allowance for bulk +transfers: one socket-level default covers every call, on the reasoning that a +gap anywhere near a minute means the peer is unhealthy whatever the payload is. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT + + +def _response(payload: dict | None = None, status_code: int = 200) -> MagicMock: + response = MagicMock() + response.status_code = status_code + response.ok = 200 <= status_code < 300 + response.raise_for_status.return_value = None + response.json.return_value = payload if payload is not None else {} + return response + + +def test_context_search_search_uses_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.knowledgebase.backends import context_search_backend as module + + backend = module.ContextSearchBackend( + index="123456789", + volcengine_access_key="ak", + volcengine_secret_key="sk", + context_search_engine_endpoint="https://ctxsearch.test/engine", + context_search_engine_apikey="apikey", + ) + + post = MagicMock( + return_value=_response({"documents": [{"content": {"sys.content": "doc"}}]}) + ) + monkeypatch.setattr(module.requests, "post", post) + + assert backend.search("hello", top_k=3) == ["doc"] + assert post.call_count == 1 + assert post.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_context_search_upload_file_uses_default_http_timeout( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from veadk.knowledgebase.backends import context_search_backend as module + + backend = module.ContextSearchBackend( + index="123456789", + volcengine_access_key="ak", + volcengine_secret_key="sk", + ) + + upload = tmp_path / "doc.txt" + upload.write_text("payload", encoding="utf-8") + + put = MagicMock(return_value=_response()) + monkeypatch.setattr(module.requests, "put", put) + + backend._upload_file( + file_path=str(upload), + upload_url="https://tos.test/signed", + headers={"Content-Type": "text/plain"}, + ) + + assert put.call_count == 1 + assert put.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_vikingdb_do_request_uses_default_http_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from veadk.knowledgebase.backends import vikingdb_knowledge_backend as module + + # `model_post_init` probes the remote collection; skip it so the backend can + # be built offline. + monkeypatch.setattr( + module.VikingDBKnowledgeBackend, + "model_post_init", + lambda self, __context: None, + ) + + backend = module.VikingDBKnowledgeBackend( + index="vikingkl_timeout", + volcengine_access_key="ak", + volcengine_secret_key="sk", + region="cn-beijing", + base_url="https://api-knowledgebase.mlp.cn-beijing.volces.com", + host="api-knowledgebase.mlp.cn-beijing.volces.com", + ) + + request: Any = MagicMock(return_value=_response({"data": {}})) + monkeypatch.setattr(module.requests, "request", request) + + backend._do_request(body={"collection_name": "vikingkl_timeout"}, path="/api/info") + + assert request.call_count == 1 + assert request.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT diff --git a/tests/runner/test_runner_contract.py b/tests/runner/test_runner_contract.py index 04563ef41..782ed9662 100644 --- a/tests/runner/test_runner_contract.py +++ b/tests/runner/test_runner_contract.py @@ -84,7 +84,10 @@ def test_parameters(self): def test_defaults(self): params = inspect.signature(Runner.run).parameters assert params["user_id"].default == "" - assert isinstance(params["session_id"].default, str) + # `None`, not a string: a string default is evaluated once at import, + # so every run omitting `session_id` would share one frozen id. The + # per-call fallback is generated in the body instead. + assert params["session_id"].default is None assert params["run_config"].default is None assert params["save_tracing_data"].default is False assert params["upload_inline_data_to_tos"].default is False diff --git a/tests/test_no_unbounded_http_calls.py b/tests/test_no_unbounded_http_calls.py new file mode 100644 index 000000000..00e58be12 --- /dev/null +++ b/tests/test_no_unbounded_http_calls.py @@ -0,0 +1,577 @@ +# 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. + +"""Repo-wide guard: no outbound HTTP call may omit an explicit timeout. + +`requests` has no default timeout. A call without one blocks forever if the +peer accepts the connection and then goes quiet, which stalls the whole +process for the synchronous tools ADK invokes inline on the event loop. The +timeout sweep already fixed every call site; this test keeps them fixed. + +The checker is deliberately narrow, because a noisy guard gets deleted: + +* the receiver must be the `requests` module itself, a local name that was + bound to `requests.Session()` in an enclosing scope, or an attribute that + was bound to one in the same class (`self._session = requests.Session()` + in `__init__`, then `self._session.request(...)` from any method); +* `session.get("state")` on a dict, `db.session.delete(post)` on SQLAlchemy, + and every other `.get`/`.delete` on an unrelated object are ignored; +* a `**kwargs` splat counts as "has a timeout" -- it cannot be disproved. + +One shape stays out of reach on purpose: a session inherited from a base class +in another module (`VikingDBMemoryClient(Service)` gets its `self.session` from +the volcengine SDK) has no binding to find in the file being parsed, and +trusting every `self.session.get(...)` without one would flag dicts. + +Unparseable files are skipped loudly: the set of files that fail to parse must +match `_EXPECTED_UNPARSEABLE` exactly, so a newly broken file fails this test +instead of silently shrinking its coverage. +""" + +import ast +from pathlib import Path + +import pytest + +_REPO_ROOT = Path(__file__).resolve().parents[1] +_PACKAGE_ROOT = _REPO_ROOT / "veadk" + +# `requests` verbs plus `request()` itself. `timeout` is keyword-only in every +# one of these signatures, so a keyword scan is sufficient. +_HTTP_METHODS = frozenset( + {"get", "post", "put", "delete", "patch", "head", "options", "request"} +) + +# Known-broken sources that predate this test. Anything else that fails to +# parse is a regression, not a skip. +_EXPECTED_UNPARSEABLE = frozenset( + { + "veadk/integrations/ve_faas/template/" + "{{cookiecutter.local_dir_name}}/src/agent.py", + } +) + +# Real files in this repo that call `.get`/`.delete` on something that is not a +# `requests` session. They must stay green. +_KNOWN_FALSE_POSITIVE_FILES = ( + "veadk/cli/cli_frontend.py", + "veadk/integrations/agentkit/evaluation/feedback.py", + "veadk/integrations/ve_faas/web_template/" + "{{cookiecutter.local_dir_name}}/src/app.py", +) + +# Real file that keeps its session on an attribute (`self._session`, bound in +# `GitHubClient.__init__` and used from `_request`). It passes a timeout today, +# so the scan above is green either way -- `test_attribute_session_file_is_...` +# below is what proves it is green for the right reason. +_REAL_ATTRIBUTE_SESSION_FILE = "veadk/cli/github_cicd.py" + +_SCOPE_NODES = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Lambda) + + +def _is_requests_session_call(node: ast.AST | None) -> bool: + """True for the expression `requests.Session(...)`.""" + return ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "Session" + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "requests" + ) + + +def _binds_a_requests_session(node: ast.AST | None) -> bool: + """True for any expression that can evaluate to a `requests.Session()`. + + `session or requests.Session()` is the standard injectable-client idiom + (`GitHubClient.__init__`), and it hides the constructor inside a `BoolOp`. + """ + if _is_requests_session_call(node): + return True + if isinstance(node, ast.BoolOp): + return any(_binds_a_requests_session(value) for value in node.values) + if isinstance(node, ast.IfExp): + return _binds_a_requests_session(node.body) or _binds_a_requests_session( + node.orelse + ) + return False + + +def _attribute_path(node: ast.AST) -> str | None: + """Dotted path for a plain attribute chain: `self._session` -> "self._session". + + Returns None for anything with a computed base (`clients[0].session`), which + cannot be matched against a binding by name. + """ + parts: list[str] = [] + while isinstance(node, ast.Attribute): + parts.append(node.attr) + node = node.value + if not parts or not isinstance(node, ast.Name): + return None + parts.append(node.id) + return ".".join(reversed(parts)) + + +def _iter_own_scope(node: ast.AST): + """Yield descendants of `node` without descending into nested scopes.""" + for child in ast.iter_child_nodes(node): + if isinstance(child, _SCOPE_NODES): + continue + yield child + yield from _iter_own_scope(child) + + +def _session_binding_targets(node: ast.AST) -> list[ast.expr]: + """Assignment targets that `node` binds to a `requests.Session()`.""" + if isinstance(node, ast.Assign) and _binds_a_requests_session(node.value): + return node.targets + if isinstance(node, ast.AnnAssign) and _binds_a_requests_session(node.value): + return [node.target] + if isinstance(node, ast.withitem) and _binds_a_requests_session(node.context_expr): + return [node.optional_vars] if node.optional_vars else [] + return [] + + +def _session_names_in_scope(scope: ast.AST) -> set[str]: + """Local names bound to `requests.Session()` directly inside `scope`.""" + names: set[str] = set() + for node in _iter_own_scope(scope): + names.update( + target.id + for target in _session_binding_targets(node) + if isinstance(target, ast.Name) + ) + return names + + +def _session_attributes_in_scope(scope: ast.AST) -> set[str]: + """Dotted attribute paths bound to `requests.Session()` inside `scope`. + + A class body is walked in full, nested scopes included: the binding lives in + `__init__` and every use lives in a sibling method, so anything narrower + would miss the only shape this pattern takes. Bare names are deliberately + *not* collected that way -- `self._session` is qualified by its owner, a + local `session` is not, and leaking locals between sibling methods is + exactly what turns a guard noisy enough to get deleted. + """ + walker = ( + ast.walk(scope) if isinstance(scope, ast.ClassDef) else _iter_own_scope(scope) + ) + paths: set[str] = set() + for node in walker: + for target in _session_binding_targets(node): + path = _attribute_path(target) + if path is not None: + paths.add(path) + return paths + + +def _call_has_timeout(node: ast.Call) -> bool: + for keyword in node.keywords: + # `keyword.arg is None` is a `**kwargs` splat: the timeout may well be + # in there, so give the call the benefit of the doubt. + if keyword.arg in ("timeout", None): + return True + return False + + +class _UnboundedHttpCallFinder(ast.NodeVisitor): + """Collect line numbers of `requests` calls that carry no `timeout=`.""" + + def __init__(self) -> None: + self.offender_lines: list[int] = [] + # Stack of session receivers visible in the current scope: bare local + # names plus dotted attribute paths such as `self._session`. + self._session_targets: list[set[str]] = [] + + def _visit_scope(self, node: ast.AST) -> None: + inherited = set(self._session_targets[-1]) if self._session_targets else set() + self._session_targets.append( + inherited + | _session_names_in_scope(node) + | _session_attributes_in_scope(node) + ) + try: + self.generic_visit(node) + finally: + self._session_targets.pop() + + visit_Module = _visit_scope + visit_FunctionDef = _visit_scope + visit_AsyncFunctionDef = _visit_scope + visit_ClassDef = _visit_scope + visit_Lambda = _visit_scope + + def _is_tracked(self, target: str) -> bool: + return bool(self._session_targets) and target in self._session_targets[-1] + + def _is_http_call(self, node: ast.Call) -> bool: + func = node.func + if not isinstance(func, ast.Attribute) or func.attr not in _HTTP_METHODS: + return False + receiver = func.value + if isinstance(receiver, ast.Name): + return receiver.id == "requests" or self._is_tracked(receiver.id) + if isinstance(receiver, ast.Attribute): + path = _attribute_path(receiver) + return path is not None and self._is_tracked(path) + # `requests.Session().get(...)` without an intermediate name. + return _binds_a_requests_session(receiver) + + def visit_Call(self, node: ast.Call) -> None: + if self._is_http_call(node) and not _call_has_timeout(node): + self.offender_lines.append(node.lineno) + self.generic_visit(node) + + +def find_unbounded_http_calls(source: str) -> list[int]: + """Line numbers of timeout-less `requests` calls in `source`.""" + finder = _UnboundedHttpCallFinder() + finder.visit(ast.parse(source)) + return sorted(finder.offender_lines) + + +def _iter_package_files() -> list[Path]: + return sorted(_PACKAGE_ROOT.rglob("*.py")) + + +def _scan_package() -> tuple[list[str], set[str]]: + """Return (offenders as `path:line`, relative paths that failed to parse).""" + offenders: list[str] = [] + unparseable: set[str] = set() + for path in _iter_package_files(): + relative = path.relative_to(_REPO_ROOT).as_posix() + try: + tree = ast.parse(path.read_text(encoding="utf-8")) + except (SyntaxError, UnicodeDecodeError, ValueError): + unparseable.add(relative) + continue + finder = _UnboundedHttpCallFinder() + finder.visit(tree) + offenders.extend(f"{relative}:{line}" for line in sorted(finder.offender_lines)) + return offenders, unparseable + + +def test_package_has_python_files_to_scan(): + # Cheap tripwire: a broken path would make every other assertion vacuous. + assert len(_iter_package_files()) > 100 + + +def test_no_unbounded_http_calls_under_veadk(): + offenders, _ = _scan_package() + + assert not offenders, ( + "Outbound HTTP call(s) without an explicit `timeout=`; `requests` has " + "no default timeout, so these can block forever. Pass " + "`DEFAULT_HTTP_TIMEOUT` (or an explicit per-call value for bulk " + "transfers) from `veadk.utils.http_defaults`:\n " + "\n ".join(offenders) + ) + + +def test_unparseable_files_match_the_allowlist(): + _, unparseable = _scan_package() + + assert unparseable == set(_EXPECTED_UNPARSEABLE), ( + "The set of files this guard cannot parse changed, so its coverage " + "changed too. Newly unparseable: " + f"{sorted(unparseable - set(_EXPECTED_UNPARSEABLE))}; no longer " + f"unparseable (drop from the allowlist): " + f"{sorted(set(_EXPECTED_UNPARSEABLE) - unparseable)}" + ) + + +@pytest.mark.parametrize("relative_path", _KNOWN_FALSE_POSITIVE_FILES) +def test_known_false_positive_files_stay_green(relative_path): + path = _REPO_ROOT / relative_path + assert path.is_file(), f"missing fixture file: {relative_path}" + + assert find_unbounded_http_calls(path.read_text(encoding="utf-8")) == [] + + +def _strip_timeout_evidence(tree: ast.AST) -> ast.AST: + """Remove everything that makes a call look bounded, keeping line numbers. + + That means `timeout=` and also `**kwargs`, which `_call_has_timeout` gives + the benefit of the doubt. `GitHubClient._request` forwards a splat, so + dropping only `timeout=` would leave it green through the splat rule and + the assertion below would prove nothing about the receiver. + """ + for node in ast.walk(tree): + if isinstance(node, ast.Call): + node.keywords = [ + kw for kw in node.keywords if kw.arg not in ("timeout", None) + ] + return tree + + +def test_attribute_session_file_is_green_for_the_right_reason(): + """The real `self._session` call site must be *seen*, not merely bounded. + + Asserting the file is clean proves nothing on its own: it was clean before + the checker could resolve attribute-held sessions at all. So take the same + file, remove what marks its calls as bounded, and require the checker to + flag exactly the calls made on `self._session`. + """ + path = _REPO_ROOT / _REAL_ATTRIBUTE_SESSION_FILE + tree = ast.parse(path.read_text(encoding="utf-8")) + + expected = sorted( + node.lineno + for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr in _HTTP_METHODS + and _attribute_path(node.func.value) == "self._session" + ) + assert expected, ( + f"{_REAL_ATTRIBUTE_SESSION_FILE} no longer holds a session on " + "`self._session`; point this test at a file that does" + ) + + finder = _UnboundedHttpCallFinder() + finder.visit(tree) + assert finder.offender_lines == [], "fixture file must be clean as committed" + + finder = _UnboundedHttpCallFinder() + finder.visit(_strip_timeout_evidence(tree)) + assert sorted(finder.offender_lines) == expected + + +@pytest.mark.parametrize( + "source", + [ + "import requests\nrequests.get(url)\n", + "import requests\nrequests.post(url, json=payload)\n", + "import requests\nrequests.put(url, data=body)\n", + "import requests\nrequests.delete(url)\n", + "import requests\nrequests.patch(url, json=payload)\n", + "import requests\nrequests.head(url)\n", + "import requests\nrequests.request('GET', url)\n", + ], +) +def test_checker_flags_module_level_calls_without_timeout(source): + assert find_unbounded_http_calls(source) == [2] + + +def test_checker_accepts_module_level_calls_with_timeout(): + source = ( + "import requests\n" + "requests.get(url, timeout=(10.0, 60.0))\n" + "requests.post(url, json=payload, timeout=30)\n" + ) + + assert find_unbounded_http_calls(source) == [] + + +def test_checker_flags_session_calls_without_timeout(): + source = ( + "import requests\n" + "def fetch(url):\n" + " session = requests.Session()\n" + " return session.post(url, json={})\n" + ) + + assert find_unbounded_http_calls(source) == [4] + + +def test_checker_accepts_session_calls_with_timeout(): + source = ( + "import requests\n" + "def fetch(url):\n" + " session = requests.Session()\n" + " return session.post(url, json={}, timeout=(10.0, 300.0))\n" + ) + + assert find_unbounded_http_calls(source) == [] + + +def test_checker_flags_attribute_session_calls_without_timeout(): + source = ( + "import requests\n" + "class Client:\n" + " def __init__(self):\n" + " self._session = requests.Session()\n" + " def fetch(self, url):\n" + " return self._session.request('GET', url)\n" + ) + + assert find_unbounded_http_calls(source) == [6] + + +def test_checker_accepts_attribute_session_calls_with_timeout(): + source = ( + "import requests\n" + "class Client:\n" + " def __init__(self):\n" + " self._session = requests.Session()\n" + " def fetch(self, url):\n" + " return self._session.request('GET', url, timeout=30)\n" + ) + + assert find_unbounded_http_calls(source) == [] + + +def test_checker_sees_through_the_injectable_session_idiom(): + """`session or requests.Session()` still binds a session.""" + source = ( + "import requests\n" + "class Client:\n" + " def __init__(self, session=None):\n" + " self._session = session or requests.Session()\n" + " def fetch(self, url):\n" + " return self._session.get(url)\n" + ) + + assert find_unbounded_http_calls(source) == [6] + + +def test_attribute_sessions_do_not_leak_across_classes(): + """A `self._session` in one class says nothing about another's.""" + source = ( + "import requests\n" + "class Http:\n" + " def __init__(self):\n" + " self._session = requests.Session()\n" + "class Store:\n" + " def __init__(self, cache):\n" + " self._session = cache\n" + " def read(self, key):\n" + " return self._session.get(key)\n" + ) + + assert find_unbounded_http_calls(source) == [] + + +def test_checker_tracks_attribute_sessions_bound_outside_a_class(): + source = ( + "import requests\n" + "def configure(client):\n" + " client.session = requests.Session()\n" + " return client.session.post(url)\n" + ) + + assert find_unbounded_http_calls(source) == [4] + + +def test_checker_tracks_sessions_opened_in_a_with_block(): + source = ( + "import requests\n" + "def fetch(url):\n" + " with requests.Session() as s:\n" + " return s.get(url)\n" + ) + + assert find_unbounded_http_calls(source) == [4] + + +def test_checker_tracks_sessions_through_nested_functions(): + source = ( + "import requests\n" + "def outer(url):\n" + " session = requests.Session()\n" + " def inner():\n" + " return session.get(url)\n" + " return inner()\n" + ) + + assert find_unbounded_http_calls(source) == [5] + + +def test_checker_flags_inline_session_calls(): + source = "import requests\nrequests.Session().get(url)\n" + + assert find_unbounded_http_calls(source) == [2] + + +def test_session_names_do_not_leak_across_functions(): + source = ( + "import requests\n" + "def a(url):\n" + " session = requests.Session()\n" + " return session.get(url, timeout=5)\n" + "def b(session):\n" + " return session.get('events')\n" + ) + + assert find_unbounded_http_calls(source) == [] + + +@pytest.mark.parametrize( + "source", + [ + # Flask/Werkzeug session dict. + "def view():\n if not session.get('admin_logged_in'):\n return 401\n", + # Plain dict payloads. + "def read(session):\n return session.get('state')\n", + "def read(session):\n return session.get('events')\n", + # SQLAlchemy. + "def drop(post):\n db.session.delete(post)\n", + "def load(model, pk):\n return db.session.get(model, pk)\n", + # Similarly named locals that are not requests sessions. + "import requests\ndef f(url):\n s = build()\n return s.get(url)\n", + # A module named like `requests` but not it. + "def f(call_id):\n return auth_requests.get(call_id)\n", + "def f(self, rid):\n return self._pending_requests.get(rid)\n", + # An attribute session with no binding in this file: inherited from a + # base class elsewhere, or just a dict on `self`. + "class C:\n def f(self, url):\n return self.session.get(url)\n", + # A computed base cannot be matched against a binding by name. + "import requests\n" + "def f(clients, url):\n" + " clients[0].session = requests.Session()\n" + " return clients[0].session.get(url)\n", + # Non-HTTP attribute on the real module. + "import requests\nrequests.Session()\n", + ], +) +def test_checker_ignores_non_requests_receivers(source): + assert find_unbounded_http_calls(source) == [] + + +def test_checker_treats_kwargs_splat_as_bounded(): + source = ( + "import requests\n" + "def fetch(url, **kwargs):\n" + " return requests.get(url, **kwargs)\n" + ) + + assert find_unbounded_http_calls(source) == [] + + +def test_checker_reports_every_offender_in_a_file(): + source = ( + "import requests\n" + "def a(url):\n" + " return requests.get(url)\n" + "def b(url):\n" + " session = requests.Session()\n" + " return session.post(url, timeout=1)\n" + "def c(url):\n" + " return requests.request('POST', url)\n" + ) + + assert find_unbounded_http_calls(source) == [3, 8] + + +def test_failure_message_names_file_and_line(tmp_path): + # End-to-end shape check on the message the guard would print: it must + # point at `file:line` so the fix is obvious. + offender = tmp_path / "broken.py" + offender.write_text("import requests\nrequests.get(url)\n", encoding="utf-8") + + lines = find_unbounded_http_calls(offender.read_text(encoding="utf-8")) + rendered = [f"{offender.name}:{line}" for line in lines] + + assert rendered == ["broken.py:2"] diff --git a/tests/test_runner.py b/tests/test_runner.py index c45903cba..dee78f2e0 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import inspect import os import pytest @@ -99,3 +100,71 @@ def test_runner(): # Run message conversion tests _test_convert_messages(runner) + + +def _make_offline_runner() -> tuple[Runner, list[str]]: + """Build a Runner whose ``run_async`` is stubbed out. + + ``Runner.run`` still executes end to end; only the LLM-backed event stream + is replaced, so the returned list records the exact ``session_id`` that each + run forwarded downstream. + """ + agent = Agent( + model_name="test_model_name", + model_provider="test_model_provider", + model_api_key="test_model_api_key", + model_api_base="test_model_api_base", + ) + runner = Runner(agent=agent, short_term_memory=ShortTermMemory()) + + seen_session_ids: list[str] = [] + + async def fake_run_async(*, user_id, session_id, new_message, **kwargs): + seen_session_ids.append(session_id) + return + yield # pragma: no cover - makes this a generator, never reached + + runner.run_async = fake_run_async + return runner, seen_session_ids + + +def test_run_session_id_default_is_a_sentinel(): + """``session_id`` must default to ``None``, not to a call expression. + + A default such as ``f"tmp-session-{formatted_timestamp()}"`` is evaluated + once, when the function is defined, so every run omitting ``session_id`` + would share a single id frozen at import time. + """ + default = inspect.signature(Runner.run).parameters["session_id"].default + assert default is None + + +@pytest.mark.asyncio +async def test_run_generates_a_fresh_session_id_per_call(monkeypatch): + """Two runs that omit ``session_id`` must get different ids. + + UUID generation is patched with deterministic values so the test pins both + per-call evaluation and the public ``tmp-session-`` prefix. + """ + runner, seen_session_ids = _make_offline_runner() + + generated = iter(["uuid-one", "uuid-two"]) + monkeypatch.setattr( + "veadk.runner.uuid.uuid4", + lambda: type("UUID", (), {"hex": next(generated)})(), + ) + + await runner.run(messages="first") + await runner.run(messages="second") + + assert seen_session_ids == ["tmp-session-uuid-one", "tmp-session-uuid-two"] + + +@pytest.mark.asyncio +async def test_run_uses_an_explicit_session_id_unchanged(): + """An explicitly passed ``session_id`` is forwarded verbatim.""" + runner, seen_session_ids = _make_offline_runner() + + await runner.run(messages="hi", session_id="explicit-session-id") + + assert seen_session_ids == ["explicit-session-id"] diff --git a/tests/tools/builtin_tools/create_agent/test_resource_source_deadlines.py b/tests/tools/builtin_tools/create_agent/test_resource_source_deadlines.py new file mode 100644 index 000000000..156c0e24c --- /dev/null +++ b/tests/tools/builtin_tools/create_agent/test_resource_source_deadlines.py @@ -0,0 +1,405 @@ +# 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. + +"""Wall-clock deadlines on the paginated create-agent resource sweeps. + +Two things are pinned here: + +1. a breached sweep returns what it collected with ``status="error"`` -- never + an exception, and never ``status="ok"`` over a truncated list; +2. that verdict belongs to one ``collect()`` call. The toolset builds each + source once and outlives every session, and ``collect_resources`` fans the + sources out with ``asyncio.gather``, so overlapping sweeps on a single + source object are the normal case. A deadline kept on the instance lets one + call answer for the other -- a partial list passed off as complete is + exactly what makes the model conclude a resource does not exist. +""" + +from __future__ import annotations + +import asyncio +import contextvars +import threading +from types import SimpleNamespace + +import pytest + +from veadk.tools.builtin_tools.create_agent.sources import ( + AgentKitKnowledgeSource, + AgentKitSkillCenterSource, + CloudCredentials, +) +from veadk.tools.builtin_tools.create_agent.sources import ( + agentkit_knowledge as knowledge_module, +) +from veadk.tools.builtin_tools.create_agent.sources import skills as skills_module + +_KNOWLEDGE_BUDGET = knowledge_module._SWEEP_DEADLINE_SECONDS +_SKILLS_BUDGET = skills_module._SWEEP_DEADLINE_SECONDS + +# Every wait below is bounded, so a regression fails the test instead of +# hanging the suite. +_TIMEOUT = 5.0 + +_SOLO = "solo" +_BREACHING = "breaching" +_CLEAN = "clean" + +_CURRENT_CALL: contextvars.ContextVar[str] = contextvars.ContextVar( + "collect_call", default=_SOLO +) + + +class _CallClock: + """A ``time.monotonic`` stand-in whose reading depends on who asks. + + A ContextVar tells overlapping ``collect()`` calls apart: both + ``asyncio.to_thread`` and ``asyncio.gather`` copy the calling context, so a + call's coroutine and its paginating worker threads all read the same clock. + One call can then outrun its budget while the other stays inside its own. + """ + + def __init__(self) -> None: + self._readings: dict[str, float] = {} + + def set(self, call: str, seconds: float) -> None: + self._readings[call] = seconds + + def monotonic(self) -> float: + return self._readings.get(_CURRENT_CALL.get(), 0.0) + + +def _use_clock(monkeypatch, module, clock: _CallClock) -> None: + """The sweeps only ever read ``monotonic``, so that is all the fake needs.""" + monkeypatch.setattr(module, "time", SimpleNamespace(monotonic=clock.monotonic)) + + +class _Gate: + """Parks one sweep inside its worker thread until the test releases it.""" + + def __init__(self) -> None: + self._loop = asyncio.get_running_loop() + self._released = threading.Event() + self.entered = asyncio.Event() + + def enter(self) -> None: + """Called on a worker thread: hand control back, then wait.""" + self._loop.call_soon_threadsafe(self.entered.set) + if not self._released.wait(_TIMEOUT): + raise AssertionError("the parked sweep was never released") + + def release(self) -> None: + self._released.set() + + async def wait_entered(self) -> None: + await asyncio.wait_for(self.entered.wait(), _TIMEOUT) + + +def _credentials(tool_context=None) -> CloudCredentials: + return CloudCredentials("ak", "sk", "sts") + + +async def _collect_as(call: str, source): + """Run one ``collect()`` under its own clock label.""" + _CURRENT_CALL.set(call) + return await source.collect() + + +def _knowledge_page(index: int, next_token: str) -> SimpleNamespace: + return SimpleNamespace( + knowledge_bases=[ + SimpleNamespace( + knowledge_id=f"kb-{index}", + provider_knowledge_id=f"provider_{index}", + provider_type="VIKINGDB_KNOWLEDGE", + name=f"handbook-{index}", + description="", + project_name="default", + region="cn-beijing", + ) + ], + next_token=next_token, + ) + + +def _skill_space(space_id: str) -> SimpleNamespace: + return SimpleNamespace(id=space_id, name=f"Team {space_id}", project_name="p") + + +def _skill_page(skill_id: str, total_count: int) -> SimpleNamespace: + return SimpleNamespace( + items=[ + SimpleNamespace( + skill_id=skill_id, + skill_name=skill_id, + skill_description="", + version="v1", + skill_status="Published", + ) + ], + total_count=total_count, + ) + + +@pytest.mark.asyncio +async def test_knowledge_sweep_reports_partial_results_at_the_deadline( + monkeypatch, +) -> None: + clock = _CallClock() + _use_clock(monkeypatch, knowledge_module, clock) + requests = [] + + class Client: + def list_knowledge_bases(self, request): + requests.append(request) + # One page alone outlives the whole sweep budget. + clock.set(_SOLO, _KNOWLEDGE_BUDGET + 1) + return _knowledge_page(1, "next") + + source = AgentKitKnowledgeSource( + client_factory=lambda credentials, region: Client(), + credential_resolver=_credentials, + ) + + result = await asyncio.wait_for(source.collect(), _TIMEOUT) + + # The next page was promised by the token and abandoned by the deadline. + assert len(requests) == 1 + assert [resource.descriptor.ref for resource in result.resources] == [ + "agentkit_kb:kb-1" + ] + assert result.status.status == "error" + assert result.status.count == 1 + assert "collected so far" in result.status.message + + +@pytest.mark.asyncio +async def test_knowledge_slow_final_page_still_reports_deadline(monkeypatch) -> None: + """A final page crossing the budget must not be labelled successful.""" + clock = _CallClock() + _use_clock(monkeypatch, knowledge_module, clock) + + class Client: + def list_knowledge_bases(self, request): + clock.set(_SOLO, _KNOWLEDGE_BUDGET + 1) + return _knowledge_page(1, "") + + source = AgentKitKnowledgeSource( + client_factory=lambda credentials, region: Client(), + credential_resolver=_credentials, + ) + + result = await asyncio.wait_for(source.collect(), _TIMEOUT) + + assert result.status.status == "error" + assert result.status.count == 1 + + +@pytest.mark.asyncio +async def test_concurrent_knowledge_sweeps_do_not_share_a_deadline(monkeypatch) -> None: + """One instance, two sweeps: neither may answer for the other. + + The healthy sweep is parked mid-pagination while the breaching sweep runs + start to finish, so a deadline latched on the instance is still set when + the healthy sweep reads its own verdict. + """ + clock = _CallClock() + _use_clock(monkeypatch, knowledge_module, clock) + gate = _Gate() + + class Client: + def list_knowledge_bases(self, request): + call = _CURRENT_CALL.get() + first_page = not request.next_token + if call == _CLEAN and first_page: + # Hold the healthy sweep open across the breaching one. + gate.enter() + if call == _BREACHING: + clock.set(_BREACHING, _KNOWLEDGE_BUDGET + 1) + return _knowledge_page(1, "next") if first_page else _knowledge_page(2, "") + + source = AgentKitKnowledgeSource( + client_factory=lambda credentials, region: Client(), + credential_resolver=_credentials, + ) + + clean_task = asyncio.create_task(_collect_as(_CLEAN, source)) + try: + await gate.wait_entered() + breaching = await asyncio.wait_for(_collect_as(_BREACHING, source), _TIMEOUT) + finally: + gate.release() + clean = await asyncio.wait_for(clean_task, _TIMEOUT) + + assert breaching.status.status == "error" + assert breaching.status.count == 1 + assert [resource.descriptor.ref for resource in breaching.resources] == [ + "agentkit_kb:kb-1" + ] + + assert clean.status.status == "ok" + assert clean.status.count == 2 + assert clean.status.message is None + assert [resource.descriptor.ref for resource in clean.resources] == [ + "agentkit_kb:kb-1", + "agentkit_kb:kb-2", + ] + + +@pytest.mark.asyncio +async def test_skill_sweep_reports_partial_results_at_the_deadline(monkeypatch) -> None: + clock = _CallClock() + _use_clock(monkeypatch, skills_module, clock) + skill_requests = [] + + class Client: + def list_skill_spaces(self, request): + return SimpleNamespace(items=[_skill_space("ss-one")], total_count=1) + + def list_skills_by_skill_space(self, request): + skill_requests.append(request) + clock.set(_SOLO, _SKILLS_BUDGET + 1) + return _skill_page("skill-a", total_count=2) + + source = AgentKitSkillCenterSource( + client_factory=lambda credentials, region: Client(), + credential_resolver=_credentials, + ) + + result = await asyncio.wait_for(source.collect(), _TIMEOUT) + + assert [request.page_number for request in skill_requests] == [1] + assert [resource.descriptor.ref for resource in result.resources] == [ + "ss-one:skill-a" + ] + assert result.status.status == "error" + assert result.status.count == 1 + # The deadline `break` skips the `for...else` page-limit `raise`. + assert "gave up after" in result.status.message + assert "exceeded 100 pages" not in result.status.message + + +@pytest.mark.asyncio +async def test_skill_slow_final_page_still_reports_deadline(monkeypatch) -> None: + """A complete-looking final Skill page still honours the elapsed budget.""" + clock = _CallClock() + _use_clock(monkeypatch, skills_module, clock) + + class Client: + def list_skill_spaces(self, request): + return SimpleNamespace(items=[_skill_space("ss-one")], total_count=1) + + def list_skills_by_skill_space(self, request): + clock.set(_SOLO, _SKILLS_BUDGET + 1) + return _skill_page("skill-a", total_count=1) + + source = AgentKitSkillCenterSource( + client_factory=lambda credentials, region: Client(), + credential_resolver=_credentials, + ) + + result = await asyncio.wait_for(source.collect(), _TIMEOUT) + + assert result.status.status == "error" + assert result.status.count == 1 + + +@pytest.mark.asyncio +async def test_skill_space_sweep_deadline_skips_the_page_limit_error( + monkeypatch, +) -> None: + clock = _CallClock() + _use_clock(monkeypatch, skills_module, clock) + space_requests = [] + + class Client: + def list_skill_spaces(self, request): + space_requests.append(request) + clock.set(_SOLO, _SKILLS_BUDGET + 1) + # `total_count` promises a second page the deadline never fetches. + return SimpleNamespace(items=[_skill_space("ss-one")], total_count=2) + + def list_skills_by_skill_space(self, request): + raise AssertionError("a breached sweep must not keep paginating") + + clients_created = 0 + + def client_factory(credentials, region): + nonlocal clients_created + clients_created += 1 + return Client() + + source = AgentKitSkillCenterSource( + client_factory=client_factory, + credential_resolver=_credentials, + ) + + result = await asyncio.wait_for(source.collect(), _TIMEOUT) + + assert [request.page_number for request in space_requests] == [1] + assert result.resources == [] + assert result.status.status == "error" + assert result.status.count == 0 + assert clients_created == 1 + assert "gave up after" in result.status.message + assert "exceeded 100 pages" not in result.status.message + + +@pytest.mark.asyncio +async def test_concurrent_skill_sweeps_do_not_share_a_deadline(monkeypatch) -> None: + """Same instance, overlapping sweeps: one breach must not taint the other.""" + clock = _CallClock() + _use_clock(monkeypatch, skills_module, clock) + gate = _Gate() + + class Client: + def list_skill_spaces(self, request): + if _CURRENT_CALL.get() == _CLEAN: + # Hold the healthy sweep open across the breaching one. + gate.enter() + return SimpleNamespace(items=[_skill_space("ss-one")], total_count=1) + + def list_skills_by_skill_space(self, request): + if request.page_number == 1: + if _CURRENT_CALL.get() == _BREACHING: + clock.set(_BREACHING, _SKILLS_BUDGET + 1) + return _skill_page("skill-a", total_count=2) + return _skill_page("skill-b", total_count=2) + + source = AgentKitSkillCenterSource( + client_factory=lambda credentials, region: Client(), + credential_resolver=_credentials, + ) + + clean_task = asyncio.create_task(_collect_as(_CLEAN, source)) + try: + await gate.wait_entered() + breaching = await asyncio.wait_for(_collect_as(_BREACHING, source), _TIMEOUT) + finally: + gate.release() + clean = await asyncio.wait_for(clean_task, _TIMEOUT) + + assert breaching.status.status == "error" + assert breaching.status.count == 1 + assert [resource.descriptor.ref for resource in breaching.resources] == [ + "ss-one:skill-a" + ] + + assert clean.status.status == "ok" + assert clean.status.count == 2 + assert clean.status.message is None + assert [resource.descriptor.ref for resource in clean.resources] == [ + "ss-one:skill-a", + "ss-one:skill-b", + ] diff --git a/tests/tools/builtin_tools/test_image_edit_bounds.py b/tests/tools/builtin_tools/test_image_edit_bounds.py new file mode 100644 index 000000000..75f30082d --- /dev/null +++ b/tests/tools/builtin_tools/test_image_edit_bounds.py @@ -0,0 +1,443 @@ +# 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. + +"""Regression tests for the two hazards fixed in ``image_edit``. + +``image_edit`` drives the *synchronous* Ark SDK from an ``async def`` tool, and +both problems follow from that: + +1. the client was built with the SDK defaults -- a 600s read timeout, a 60s + connect timeout and two retries -- so one unresponsive ``images.generate`` + could hold the tool for half an hour, repeated for *every* item in + ``params``. ``_get_client`` now passes an explicit ``httpx.Timeout`` built + from ``DEFAULT_IMAGE_EDIT_READ_TIMEOUT`` and the shared + ``DEFAULT_CONNECT_TIMEOUT``, plus ``DEFAULT_IMAGE_EDIT_MAX_RETRIES``; +2. that blocking call, and the equally blocking ``_upload_image_to_tos`` on the + ``b64_json`` branch, ran inline on the event loop. Awaiting a coroutine that + parks the loop for minutes stalls every other agent in the process, so both + are now handed to ``asyncio.to_thread``. + +The timeout tests read their expectations from the module constants rather than +from literals, so retuning the budget stays a one-line change; what they pin is +the *shape* -- connect far tighter than read, and the whole worst-case wait far +below what the SDK would have allowed. + +The offload tests are the ones that fail loudly against the pre-fix code: the +SDK stub records ``threading.get_ident()``, and a blocked event loop is caught +directly by watching whether a concurrently scheduled task still gets to run +while the "SDK call" is in flight. + +Nothing here touches the network. ``_get_client`` is replaced wherever the tool +is driven end to end, ``MODEL_EDIT_API_KEY`` short-circuits the credential +lookup that would otherwise reach ``settings.model.api_key`` (a cached property +that can fetch a live Ark token), and an autouse tripwire makes every httpx +transport raise. Every awaited call is bounded, so a regression fails the suite +instead of freezing it. +""" + +import asyncio +import base64 +import contextlib +import threading +import time +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import httpx +import pytest +from volcenginesdkarkruntime import Ark + +from veadk.consts import DEFAULT_IMAGE_EDIT_MODEL_API_BASE +from veadk.tools.builtin_tools import image_edit as image_edit_module +from veadk.tools.builtin_tools.image_edit import image_edit +from veadk.utils.http_defaults import DEFAULT_CONNECT_TIMEOUT + +# Wall-clock ceiling on anything awaited here. Generous enough never to flake, +# short enough that a regression fails CI instead of wedging it. +_HARD_BOUND_SECONDS = 15.0 + +# How long a stubbed SDK call waits for the event loop to show a sign of life +# before giving up. Only reached when the loop is blocked, i.e. on a regression. +_LOOP_PROBE_BOUND_SECONDS = 3.0 + +# Event-loop iterations the probe below must observe while the stubbed SDK call +# is still running. One would do; three rules out a coincidental wake-up. +_REQUIRED_TICKS = 3 + +_ORIGIN_IMAGE = "https://example.invalid/origin.png" +_EDITED_URL = "https://example.invalid/edited.png" +_TOS_URL = "https://tos.example.invalid/edited.png" +_IMAGE_BYTES = b"not-really-a-png" + + +@pytest.fixture(autouse=True) +def image_edit_env(monkeypatch): + """Deterministic credentials and endpoint, with no live token lookup. + + ``_get_api_key`` falls back to ``settings.model.api_key``, a cached property + that can fetch a real Ark token, so the env var short-circuits it before it + gets there. The other two are cleared because a developer's ``.env`` is + loaded into ``os.environ`` at import time and would otherwise leak in. + """ + monkeypatch.setenv("MODEL_EDIT_API_KEY", "test_api_key") + monkeypatch.delenv("MODEL_EDIT_API_BASE", raising=False) + monkeypatch.delenv("MODEL_EDIT_NAME", raising=False) + + +@pytest.fixture(autouse=True) +def no_network(): + """Tripwire: a real request fails the test rather than dialing out. + + The timeout tests build a genuine ``Ark`` client, so this guards against a + stub going missing and the suite quietly talking to the live endpoint. + """ + + def _explode(*args, **kwargs): + raise AssertionError("image_edit tests must not perform real HTTP") + + with ( + patch.object(httpx.HTTPTransport, "handle_request", _explode), + patch.object(httpx.AsyncHTTPTransport, "handle_async_request", _explode), + ): + yield + + +def _tool_context() -> SimpleNamespace: + """A ToolContext-shaped stub. + + ``image_edit`` writes the resulting URL into ``tool_context.state``; the + rest is what ``add_span_attributes`` reads off the invocation context. That + helper swallows its own errors, so the fields exist here only to keep a + passing test from printing a traceback about them. + """ + return SimpleNamespace( + state={}, + agent_name="test_agent", + _invocation_context=SimpleNamespace( + app_name="test_app", + user_id="test_user", + session=SimpleNamespace(id="test_session"), + ), + ) + + +def _item(**overrides) -> dict: + """One entry of the ``params`` list, defaulting to the ``url`` branch.""" + item = { + "image_name": "edited", + "prompt": "make it blue", + "origin_image": _ORIGIN_IMAGE, + } + item.update(overrides) + return item + + +def _response(*, url: str = None, b64_json: str = None) -> SimpleNamespace: + """An ``images.generate`` reply, shaped as the tool consumes it.""" + return SimpleNamespace( + data=[SimpleNamespace(url=url, b64_json=b64_json)], + usage=SimpleNamespace(output_tokens=1, total_tokens=2), + ) + + +def _stub_client(generate) -> MagicMock: + """An Ark-shaped client whose only live part is ``images.generate``.""" + client = MagicMock() + client.images.generate.side_effect = generate + return client + + +@contextlib.contextmanager +def _driving(generate, upload=None): + """Patch out everything ``image_edit`` would otherwise do for real. + + ``traceback.print_exc`` is silenced too: the tool prints it on every + per-item failure, which is expected in the tests that exercise that path + and would otherwise bury the real output in noise. + """ + with contextlib.ExitStack() as stack: + stack.enter_context( + patch.object( + image_edit_module, "_get_client", return_value=_stub_client(generate) + ) + ) + stack.enter_context(patch.object(image_edit_module.traceback, "print_exc")) + if upload is not None: + stack.enter_context( + patch.object(image_edit_module, "_upload_image_to_tos", upload) + ) + yield + + +async def _run(params: list, tool_context) -> dict: + """Drive the tool under a hard bound, so a stall fails instead of hangs.""" + return await asyncio.wait_for(image_edit(params, tool_context), _HARD_BOUND_SECONDS) + + +# -------------------------------------------------------------------------- +# 1. the SDK client is built with an explicit, bounded budget +# -------------------------------------------------------------------------- + + +def test_client_carries_the_module_timeout_and_retry_budget(): + """The constructed client must actually hold the configured bounds. + + Sourced from the module constants rather than from literals: retuning the + budget must not need a test edit, only a code one. + """ + with image_edit_module._get_client() as client: + assert client.timeout.read == image_edit_module.DEFAULT_IMAGE_EDIT_READ_TIMEOUT + assert client.timeout.connect == DEFAULT_CONNECT_TIMEOUT + assert client.max_retries == image_edit_module.DEFAULT_IMAGE_EDIT_MAX_RETRIES + # httpx fans the single `timeout=` value out to the remaining phases. + assert client.timeout.write == image_edit_module.DEFAULT_IMAGE_EDIT_READ_TIMEOUT + assert client.timeout.pool == image_edit_module.DEFAULT_IMAGE_EDIT_READ_TIMEOUT + + +def test_connect_budget_is_far_tighter_than_the_read_budget(): + """Reaching an unreachable peer must fail fast, generation may take a while. + + The two halves are different quantities: a TCP/TLS handshake that has not + completed in seconds is not going to, whereas an image edit legitimately + runs for minutes. The client must use the shared connect default rather + than the SDK's minute-long one. + """ + with image_edit_module._get_client() as client: + assert client.timeout.connect == DEFAULT_CONNECT_TIMEOUT + assert client.timeout.connect * 5 <= client.timeout.read, ( + "the connect budget is not meaningfully tighter than the read budget: " + f"connect={client.timeout.connect}s read={client.timeout.read}s" + ) + + +def test_bounds_improve_substantially_on_the_ark_sdk_defaults(): + """The whole point of the change: a hung call can no longer hold ~30 minutes. + + The comparison is drawn against a client built the way ``_get_client`` used + to build one, so it tracks whatever the installed SDK's defaults happen to + be instead of hardcoding 600s and two retries. + """ + with ( + image_edit_module._get_client() as bounded, + Ark(api_key="test_api_key", base_url=DEFAULT_IMAGE_EDIT_MODEL_API_BASE) as sdk, + ): + assert bounded.timeout.read * 3 <= sdk.timeout.read, ( + "the read budget is not well below the SDK default: " + f"{bounded.timeout.read}s vs {sdk.timeout.read}s" + ) + assert bounded.max_retries < sdk.max_retries + + # Worst case for a single item, which `image_edit` pays once per entry + # in `params`: the initial attempt plus every retry, each able to run + # the full read budget. + bounded_worst = (1 + bounded.max_retries) * bounded.timeout.read + sdk_worst = (1 + sdk.max_retries) * sdk.timeout.read + assert bounded_worst * 4 <= sdk_worst, ( + "a single hung item can still hold the tool for " + f"{bounded_worst}s (SDK default: {sdk_worst}s)" + ) + + +# -------------------------------------------------------------------------- +# 2. the blocking SDK calls run off the event loop +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_image_generate_runs_off_the_event_loop(): + """``images.generate`` is synchronous, so it must not run on the loop thread.""" + loop_thread = threading.get_ident() + generate_threads: list[int] = [] + + def generate(**kwargs): + generate_threads.append(threading.get_ident()) + return _response(url=_EDITED_URL) + + tool_context = _tool_context() + with _driving(generate): + result = await _run([_item()], tool_context) + + assert generate_threads, "images.generate was never called" + assert loop_thread not in generate_threads, ( + "the blocking Ark call ran on the event-loop thread " + f"({loop_thread}); it must go through asyncio.to_thread" + ) + assert result == { + "status": "success", + "success_list": [{"edited": _EDITED_URL}], + "error_list": [], + } + assert tool_context.state["edited_url"] == _EDITED_URL + + +@pytest.mark.asyncio +async def test_event_loop_keeps_turning_while_generate_blocks(): + """A different thread is only useful if the loop is genuinely free. + + A ticker task counts loop iterations; the stubbed SDK call blocks until it + has seen the count advance. Off the loop that takes milliseconds -- on it, + the ticker cannot run at all and the call gives up at the probe bound. + """ + ticks = {"n": 0} + observed: list[int] = [] + + async def ticker(): + while True: + ticks["n"] += 1 + await asyncio.sleep(0.01) + + def generate(**kwargs): + start = ticks["n"] + deadline = time.monotonic() + _LOOP_PROBE_BOUND_SECONDS + while ticks["n"] - start < _REQUIRED_TICKS and time.monotonic() < deadline: + time.sleep(0.01) + observed.append(ticks["n"] - start) + return _response(url=_EDITED_URL) + + ticker_task = asyncio.create_task(ticker()) + try: + with _driving(generate): + result = await _run([_item()], _tool_context()) + finally: + ticker_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await ticker_task + + assert observed and observed[0] >= _REQUIRED_TICKS, ( + "the event loop made no progress while images.generate was in flight " + f"(observed {observed} iterations in {_LOOP_PROBE_BOUND_SECONDS}s): " + "the blocking SDK call is stalling the loop" + ) + assert result["status"] == "success" + + +@pytest.mark.asyncio +async def test_tos_upload_runs_off_the_event_loop(): + """The ``b64_json`` branch uploads to TOS, which blocks just as hard.""" + loop_thread = threading.get_ident() + generate_threads: list[int] = [] + upload_threads: list[int] = [] + upload_calls: list[tuple] = [] + + def generate(**kwargs): + generate_threads.append(threading.get_ident()) + return _response(b64_json=base64.b64encode(_IMAGE_BYTES).decode()) + + def upload(image_bytes, object_key): + upload_threads.append(threading.get_ident()) + # Recorded rather than asserted here: this runs on a worker thread, and + # the tool would turn an AssertionError into a silent per-item failure. + upload_calls.append((image_bytes, object_key)) + return _TOS_URL + + tool_context = _tool_context() + with _driving(generate, upload=upload): + result = await _run([_item(response_format="b64_json")], tool_context) + + assert upload_threads, "_upload_image_to_tos was never called" + assert loop_thread not in generate_threads, ( + "the blocking Ark call ran on the event-loop thread " + f"({loop_thread}); it must go through asyncio.to_thread" + ) + assert loop_thread not in upload_threads, ( + "the blocking TOS upload ran on the event-loop thread " + f"({loop_thread}); it must go through asyncio.to_thread" + ) + assert upload_calls == [(_IMAGE_BYTES, "edited.png")] + assert result == { + "status": "success", + "success_list": [{"edited": _TOS_URL}], + "error_list": [], + } + assert tool_context.state["edited_url"] == _TOS_URL + + +# -------------------------------------------------------------------------- +# 3. moving the calls off the loop did not change how failures are reported +# -------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_generate_failure_on_the_worker_thread_still_lands_in_error_list(): + """An exception raised off-thread must surface exactly as it did inline. + + ``asyncio.to_thread`` re-raises in the awaiting coroutine, so the tool's + existing ``except Exception`` still catches it and records the failure per + item: the bad entry goes to ``error_list``, the good one is unaffected, and + a partial batch is still reported as ``success``. + """ + loop_thread = threading.get_ident() + generate_threads: list[int] = [] + + def generate(**kwargs): + generate_threads.append(threading.get_ident()) + if kwargs["prompt"] == "boom": + raise RuntimeError("ark refused the edit") + return _response(url=_EDITED_URL) + + tool_context = _tool_context() + params = [ + _item(image_name="broken", prompt="boom"), + _item(image_name="fine"), + ] + with _driving(generate): + result = await _run(params, tool_context) + + assert len(generate_threads) == 2 + assert loop_thread not in generate_threads, ( + "the failing Ark call ran on the event-loop thread; a raising SDK call " + "must still be offloaded" + ) + assert result == { + "status": "success", + "success_list": [{"fine": _EDITED_URL}], + "error_list": ["broken"], + } + assert "broken_url" not in tool_context.state + assert tool_context.state["fine_url"] == _EDITED_URL + + +@pytest.mark.asyncio +async def test_failed_tos_upload_still_lands_in_error_list(): + """``_upload_image_to_tos`` swallows its own errors and returns ``None``. + + Running it on a worker thread must not change that: a falsy return is still + a per-item failure, and with nothing else in the batch the tool reports + ``error``. + """ + loop_thread = threading.get_ident() + upload_threads: list[int] = [] + + def generate(**kwargs): + return _response(b64_json=base64.b64encode(_IMAGE_BYTES).decode()) + + def upload(image_bytes, object_key): + upload_threads.append(threading.get_ident()) + return None + + tool_context = _tool_context() + with _driving(generate, upload=upload): + result = await _run([_item(response_format="b64_json")], tool_context) + + assert upload_threads, "_upload_image_to_tos was never called" + assert loop_thread not in upload_threads, ( + "the failing TOS upload ran on the event-loop thread; it must still be " + "offloaded" + ) + assert result == { + "status": "error", + "success_list": [], + "error_list": ["edited"], + } + assert tool_context.state == {} diff --git a/tests/tools/builtin_tools/test_llm_shield_contract.py b/tests/tools/builtin_tools/test_llm_shield_contract.py new file mode 100644 index 000000000..48f221671 --- /dev/null +++ b/tests/tools/builtin_tools/test_llm_shield_contract.py @@ -0,0 +1,296 @@ +# 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. + +"""Contract tests for ``LLMShieldPlugin``. + +The plugin is wired two ways, and the two wirings disagree: ADK's plugin +manager always ``await``s the hook and passes ``tool_args`` / ``result`` / +``agent``, while the agent callback path awaits only if the hook returns an +awaitable and passes ``args`` / ``tool_response``. These tests pin both calling +conventions, keep the blocking moderation request off the event loop, and pin +the fail-open behavior. No test touches the network. +""" + +import inspect +import os +import threading +from typing import Any, Dict, List, Optional + +import pytest +from google.adk.models import LlmRequest, LlmResponse +from google.adk.plugins import BasePlugin, PluginManager +from google.genai import types + +# The module builds `content_safety = LLMShieldPlugin()` at import time and the +# constructor requires the app id to be configured, so set it before importing. +os.environ.setdefault("TOOL_LLM_SHIELD_APP_ID", "test-app-id") + +from veadk.tools.builtin_tools import llm_shield # noqa: E402 +from veadk.tools.builtin_tools.llm_shield import ( # noqa: E402 + LLMShieldPlugin, + content_safety, +) + +# The hooks that actually call the moderation service. +MODERATION_HOOKS = [ + "before_model_callback", + "after_model_callback", + "before_tool_callback", + "after_tool_callback", +] + +# Every hook the class overrides, including the two no-ops. +ALL_HOOKS = ["before_agent_callback", "after_agent_callback", *MODERATION_HOOKS] + +BLOCK_MESSAGE = "Your request has been blocked due to: Prompt Injection." + + +class _DummyTool: + name = "dummy_tool" + + +class _DummyToolContext: + invocation_id = "invocation-1" + session = None + + +def _llm_request() -> LlmRequest: + return LlmRequest( + contents=[types.Content(role="user", parts=[types.Part(text="hello")])] + ) + + +def _llm_response() -> LlmResponse: + return LlmResponse( + content=types.Content(role="model", parts=[types.Part(text="hello")]) + ) + + +def _base_plugin_kwargs(hook: str) -> Dict[str, Any]: + """Build the exact keyword set ``BasePlugin`` declares for ``hook``. + + Read from the installed ADK rather than hardcoded, so this tracks upstream: + a renamed or added plugin argument fails here instead of silently drifting. + """ + values = { + "agent": object(), + "callback_context": object(), + "llm_request": _llm_request(), + "llm_response": _llm_response(), + "tool": _DummyTool(), + "tool_args": {"query": "hello"}, + "tool_context": _DummyToolContext(), + "result": {"output": "hello"}, + } + params = inspect.signature(getattr(BasePlugin, hook)).parameters + names = [name for name in params if name != "self"] + assert names, f"BasePlugin.{hook} declares no arguments" + return {name: values[name] for name in names} + + +@pytest.fixture(autouse=True) +def _no_network(monkeypatch): + """Fail loudly instead of reaching the LLM Shield endpoint.""" + + def _blocked(*args, **kwargs): + raise AssertionError("tests must not perform HTTP requests") + + monkeypatch.setattr(llm_shield.requests, "post", _blocked) + + +@pytest.fixture +def plugin() -> LLMShieldPlugin: + return LLMShieldPlugin() + + +def _stub_shield( + monkeypatch, + plugin: LLMShieldPlugin, + verdict: Optional[str] = None, + calls: Optional[List[Dict[str, Any]]] = None, +) -> None: + """Replace the blocking moderation request with a recorded stub.""" + + def _fake_request(**kwargs): + if calls is not None: + calls.append(kwargs) + return verdict + + monkeypatch.setattr(plugin, "_request_llm_shield", _fake_request) + + +def test_content_safety_is_an_adk_plugin(): + assert isinstance(content_safety, LLMShieldPlugin) + assert isinstance(content_safety, BasePlugin) + + +@pytest.mark.parametrize("hook", ALL_HOOKS) +def test_hook_is_a_coroutine_function(hook): + # The plugin manager does a bare `await callback_method(...)`, so a plain + # `def` hook returning None raises `TypeError: object NoneType can't be + # used in 'await' expression`. + assert inspect.iscoroutinefunction(getattr(LLMShieldPlugin, hook)) + + +@pytest.mark.parametrize("hook", ALL_HOOKS) +@pytest.mark.asyncio +async def test_hook_accepts_base_plugin_keywords(plugin, monkeypatch, hook): + _stub_shield(monkeypatch, plugin) + + assert await getattr(plugin, hook)(**_base_plugin_kwargs(hook)) is None + + +@pytest.mark.parametrize("hook", ALL_HOOKS) +@pytest.mark.asyncio +async def test_plugin_manager_invokes_hook(plugin, monkeypatch, hook): + # The `Runner(plugins=[content_safety])` wiring, driven through ADK itself. + _stub_shield(monkeypatch, plugin) + manager = PluginManager(plugins=[plugin]) + + assert await getattr(manager, f"run_{hook}")(**_base_plugin_kwargs(hook)) is None + + +@pytest.mark.parametrize("hook", MODERATION_HOOKS) +@pytest.mark.asyncio +async def test_plugin_manager_short_circuits_on_block(plugin, monkeypatch, hook): + _stub_shield(monkeypatch, plugin, verdict=BLOCK_MESSAGE) + manager = PluginManager(plugins=[plugin]) + + blocked = await getattr(manager, f"run_{hook}")(**_base_plugin_kwargs(hook)) + + if hook.endswith("model_callback"): + assert isinstance(blocked, LlmResponse) + assert blocked.content.parts[0].text == BLOCK_MESSAGE + else: + assert blocked == {"result": BLOCK_MESSAGE} + + +@pytest.mark.asyncio +async def test_agent_callback_convention_model_hooks(plugin, monkeypatch): + _stub_shield(monkeypatch, plugin, verdict=BLOCK_MESSAGE) + callback_context = object() + + # Keyword form used by `base_llm_flow`. + blocked = await plugin.before_model_callback( + callback_context=callback_context, llm_request=_llm_request() + ) + assert isinstance(blocked, LlmResponse) + assert blocked.content.parts[0].text == BLOCK_MESSAGE + assert blocked.partial is True + + blocked = await plugin.after_model_callback( + callback_context=callback_context, llm_response=_llm_response() + ) + assert isinstance(blocked, LlmResponse) + assert blocked.content.parts[0].text == BLOCK_MESSAGE + assert blocked.partial is True + + # Positional form allowed by the `_SingleBeforeModelCallback` type alias. + blocked = await plugin.before_model_callback(callback_context, _llm_request()) + assert blocked.content.parts[0].text == BLOCK_MESSAGE + + blocked = await plugin.after_model_callback(callback_context, _llm_response()) + assert blocked.content.parts[0].text == BLOCK_MESSAGE + + +@pytest.mark.asyncio +async def test_agent_callback_convention_tool_hooks(plugin, monkeypatch): + calls: List[Dict[str, Any]] = [] + _stub_shield(monkeypatch, plugin, verdict=BLOCK_MESSAGE, calls=calls) + tool = _DummyTool() + tool_context = _DummyToolContext() + + # `args` / `tool_response`, the names `functions.py` uses for agent callbacks. + blocked = await plugin.before_tool_callback( + tool=tool, args={"query": "hello"}, tool_context=tool_context + ) + assert blocked == {"result": BLOCK_MESSAGE} + assert calls[-1]["message"] == "query: hello" + + blocked = await plugin.after_tool_callback( + tool=tool, + args={"query": "hello"}, + tool_context=tool_context, + tool_response="tool output", + ) + assert blocked == {"result": BLOCK_MESSAGE} + assert calls[-1]["message"] == "tool output" + + # Positional form allowed by the `_SingleBeforeToolCallback` type alias. + blocked = await plugin.before_tool_callback(tool, {"query": "hello"}, tool_context) + assert blocked == {"result": BLOCK_MESSAGE} + + blocked = await plugin.after_tool_callback( + tool, {"query": "hello"}, tool_context, "tool output" + ) + assert blocked == {"result": BLOCK_MESSAGE} + + +@pytest.mark.asyncio +async def test_tool_hooks_read_the_plugin_argument_spelling(plugin, monkeypatch): + calls: List[Dict[str, Any]] = [] + _stub_shield(monkeypatch, plugin, calls=calls) + tool = _DummyTool() + tool_context = _DummyToolContext() + + await plugin.before_tool_callback( + tool=tool, tool_args={"query": "hello"}, tool_context=tool_context + ) + assert calls[-1]["message"] == "query: hello" + + await plugin.after_tool_callback( + tool=tool, + tool_args={"query": "hello"}, + tool_context=tool_context, + result={"output": "tool output"}, + ) + assert calls[-1]["message"] == "tool output\n" + + +@pytest.mark.asyncio +async def test_agent_callback_convention_agent_hooks(plugin): + # The agent path passes only `callback_context`; both hooks stay no-ops. + assert await plugin.before_agent_callback(callback_context=object()) is None + assert await plugin.after_agent_callback(callback_context=object()) is None + + +@pytest.mark.parametrize("hook", MODERATION_HOOKS) +@pytest.mark.asyncio +async def test_moderation_request_runs_off_the_event_loop(plugin, monkeypatch, hook): + # `_request_llm_shield` is a blocking `requests.post` with a 50s default + # timeout, fired up to four times per turn: it must not stall the loop. + loop_thread = threading.get_ident() + request_threads: List[int] = [] + + def _fake_request(**kwargs): + request_threads.append(threading.get_ident()) + return None + + monkeypatch.setattr(plugin, "_request_llm_shield", _fake_request) + + await getattr(plugin, hook)(**_base_plugin_kwargs(hook)) + + assert request_threads, "the moderation request was never made" + assert request_threads[0] != loop_thread + + +@pytest.mark.parametrize("hook", MODERATION_HOOKS) +@pytest.mark.asyncio +async def test_fail_open_when_moderation_returns_none(plugin, monkeypatch, hook): + # `None` from the shield means "proceed", including on error or timeout. + calls: List[Dict[str, Any]] = [] + _stub_shield(monkeypatch, plugin, verdict=None, calls=calls) + + assert await getattr(plugin, hook)(**_base_plugin_kwargs(hook)) is None + assert calls, "the moderation request was never made" diff --git a/tests/tools/builtin_tools/test_mobile_run_bounds.py b/tests/tools/builtin_tools/test_mobile_run_bounds.py new file mode 100644 index 000000000..eacf12f7e --- /dev/null +++ b/tests/tools/builtin_tools/test_mobile_run_bounds.py @@ -0,0 +1,489 @@ +# 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. + +"""Regression tests for the bounded waits in the mobile-use tool. + +Every wait exercised here used to be unbounded, so each test is written so that +the old behaviour *fails* (loudly, and quickly) instead of hanging forever: + +* ``PodPool.acquire_pod`` is driven from a worker thread with a hard + ``result(timeout=...)`` bound, and the queue is unblocked before the test + gives up so a stuck worker can never wedge the suite. +* the async tests drive the tool with a fake clock plus a fake ``asyncio.sleep`` + that still yields to the real event loop, so an outer ``asyncio.wait_for`` + can cancel a runaway polling loop. + +No network is involved: every ``ve_request`` caller in the module is patched. +""" + +import asyncio +import concurrent.futures +import time +import types + +import pytest + +from veadk.tools.builtin_tools import mobile_run +from veadk.tools.builtin_tools.mobile_run import ( + GetAgentResultResponse, + GetAgentResultResult, + ListAgentRunCurrentResponse, + ListAgentRunCurrentResponseResult, + PodPool, + ResponseMetadata, + RunAgentTaskResponse, + RunAgentTaskResult, +) + +# Real wall-clock ceiling for anything that could regress into an endless wait. +# Generous enough to never flake, short enough that CI fails fast. +_HARD_BOUND_SECONDS = 15 + +_UNBLOCK_SENTINEL = "unblock-sentinel" + + +def _acquire_pod_bounded(pool: PodPool, hint: str): + """Call ``pool.acquire_pod()`` off-thread under a hard real-time bound. + + If ``acquire_pod`` blocks (the pre-fix ``Queue.get(block=True)`` behaviour) + the test fails instead of hanging, and a sentinel is pushed into the queue + first so the stranded worker thread can finish and be joined. + """ + executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + try: + future = executor.submit(pool.acquire_pod) + try: + return future.result(timeout=_HARD_BOUND_SECONDS) + except concurrent.futures.TimeoutError: + pool.available_pods.put(_UNBLOCK_SENTINEL) + pytest.fail( + f"acquire_pod() blocked for more than {_HARD_BOUND_SECONDS}s " + f"({hint}); the bounded wait on the pod queue has regressed" + ) + finally: + executor.shutdown(wait=True) + + +def test_acquire_pod_returns_none_when_pool_is_empty(): + """An empty pool must give up instead of parking the caller forever.""" + pool = PodPool([]) + + started = time.monotonic() + pod = _acquire_pod_bounded(pool, hint="empty pool") + elapsed = time.monotonic() - started + + assert pod is None + # The wait must be governed by the module constant, not by luck. + assert elapsed < mobile_run.POD_ACQUIRE_TIMEOUT_SECONDS + 5 + + +def test_acquire_pod_round_trips_a_single_pod(): + """Exhausting the pool yields ``None``; releasing makes the pod available again.""" + pool = PodPool(["pod-1"]) + + assert _acquire_pod_bounded(pool, hint="first acquire") == "pod-1" + assert pool.get_pod_status("pod-1") == "pending" + assert pool.get_available_count() == 0 + + assert _acquire_pod_bounded(pool, hint="exhausted pool") is None + + pool.release_pod("pod-1") + assert pool.get_pod_status("pod-1") == "available" + + assert _acquire_pod_bounded(pool, hint="after release") == "pod-1" + + +def test_acquire_pod_propagates_non_empty_queue_errors(): + """Only ``queue.Empty`` maps to ``None``. + + The handler used to be a blanket ``except Exception`` that swallowed real + failures and reported them as "no pod available". + """ + + class _ExplodingQueue: + def get(self, block=True, timeout=None): + raise RuntimeError("queue backend exploded") + + def qsize(self): + return 0 + + pool = PodPool(["pod-1"]) + pool.available_pods = _ExplodingQueue() # type: ignore[assignment] + + with pytest.raises(RuntimeError, match="queue backend exploded"): + pool.acquire_pod() + + +class _FakeClock: + """Deterministic stand-in for the module-level ``time`` module. + + ``monotonic()`` and ``time()`` are tracked separately on purpose. Every + deadline in the tool is armed on the monotonic clock, so a test can step the + wall clock the way an NTP correction would and assert that no deadline + moves. Keeping ``time()`` here (rather than dropping it once the module + stopped calling it) is what makes a regression back to wall-clock arithmetic + fail loudly instead of going unnoticed. + """ + + def __init__(self, start: float = 1_000.0) -> None: + self._now = start + self._wall_skew = 0.0 + + def monotonic(self) -> float: + return self._now + + def time(self) -> float: + return self._now + self._wall_skew + + def advance(self, seconds: float) -> None: + self._now += seconds + + def step_wall_clock(self, seconds: float) -> None: + """Jump ``time()`` only, leaving ``monotonic()`` untouched.""" + self._wall_skew += seconds + + +def _metadata(action: str) -> ResponseMetadata: + return ResponseMetadata( + RequestId="req-1", + Action=action, + Version="2023-08-01", + Service="ipaas", + Region="cn-north-1", + ) + + +def _install_fake_backend( + monkeypatch, + *, + is_success, + timeout_seconds: int, + content: str = "still working", + pods_available: bool = True, + wall_clock_step_per_sleep: float = 0.0, + max_sleeps: int = 100, +): + """Build a fully offline ``mobile_use_tool`` whose result poll never finishes. + + ``_get_task_result`` always reports ``is_success``; time only moves when the + (faked) ``asyncio.sleep`` inside a wait loop is awaited, so the tests run + instantly while still exercising the real deadline arithmetic. + + ``pods_available=False`` starves ``acquire_pod`` so the pod-wait loop is the + one under test, ``wall_clock_step_per_sleep`` skews ``time()`` on every + retry without touching ``monotonic()``, and ``max_sleeps`` turns a loop that + stopped honouring its deadline into a fast failure instead of a spin. + """ + # The module reads these at call time via ``_require_env_vars`` / + # ``_get_product_and_pod``; monkeypatch keeps a clean machine clean. + monkeypatch.setenv("VOLCENGINE_ACCESS_KEY", "test-ak") + monkeypatch.setenv("VOLCENGINE_SECRET_KEY", "test-sk") + monkeypatch.setenv("TOOL_MOBILE_USE_TOOL_ID", "['product1-pod1']") + monkeypatch.setattr(mobile_run, "tool_ids", ["product1-pod1"]) + monkeypatch.setattr(mobile_run, "product_id", None) + monkeypatch.setattr(mobile_run, "pod_ids", None) + + calls: dict = { + "run": 0, + "result": 0, + "step": 0, + "cancel": 0, + "cancelled": [], + "sleeps": [], + } + clock = _FakeClock() + + if not pods_available: + # Nothing to hand out, so the tool has to fall through to the pod-wait + # deadline instead of starting a run. + monkeypatch.setattr(mobile_run.PodPool, "acquire_pod", lambda self: None) + + def fake_run_agent_task(*_args, **_kwargs) -> RunAgentTaskResponse: + calls["run"] += 1 + return RunAgentTaskResponse( + ResponseMetadata=_metadata("RunAgentTaskOneStep"), + Result=RunAgentTaskResult( + RunId="run-1", RunName="test-run", ThreadId="thread-1" + ), + ) + + def fake_get_task_result(_task_id: str) -> GetAgentResultResponse: + calls["result"] += 1 + return GetAgentResultResponse( + ResponseMetadata=_metadata("GetAgentResult"), + Result=GetAgentResultResult( + IsSuccess=is_success, + Content=content, + StructOutput="", + ScreenShots=[], + ), + ) + + def fake_get_current_step(_task_id: str) -> ListAgentRunCurrentResponse: + calls["step"] += 1 + return ListAgentRunCurrentResponse( + ResponseMetadata=_metadata("ListAgentRunCurrentStep"), + Result=ListAgentRunCurrentResponseResult( + RunId="run-1", ThreadId="thread-1", Results=[] + ), + ) + + def fake_cancel_task(_task_id: str) -> None: + calls["cancel"] += 1 + calls["cancelled"].append(_task_id) + + async def fake_sleep(delay): + calls["sleeps"].append(delay) + if len(calls["sleeps"]) > max_sleeps: + # Fail fast rather than spin: a loop that ignores its deadline must + # not burn the outer ``asyncio.wait_for`` bound to be noticed. + raise AssertionError( + f"faked asyncio.sleep called more than {max_sleeps} times; " + "a wait loop is no longer honouring its deadline" + ) + clock.advance(delay) + clock.step_wall_clock(wall_clock_step_per_sleep) + # A real yield, so an outer ``asyncio.wait_for`` can still cancel a + # polling loop that refuses to terminate. + await asyncio.sleep(0) + + monkeypatch.setattr(mobile_run, "_run_agent_task", fake_run_agent_task) + monkeypatch.setattr(mobile_run, "_get_task_result", fake_get_task_result) + monkeypatch.setattr(mobile_run, "_get_current_step", fake_get_current_step) + monkeypatch.setattr(mobile_run, "_cancel_task", fake_cancel_task) + monkeypatch.setattr(mobile_run, "time", clock) + monkeypatch.setattr( + mobile_run, + "asyncio", + types.SimpleNamespace( + sleep=fake_sleep, + gather=asyncio.gather, + # The blocking ACEP helpers are awaited via `asyncio.to_thread`, so + # the stand-in namespace must carry it or the tool fails before it + # can ever reach the bounds these tests pin. + to_thread=asyncio.to_thread, + ), + ) + + tool = mobile_run.create_mobile_use_tool( + system_prompt="you are a test agent", + timeout_seconds=timeout_seconds, + max_step=3, + step_interval_seconds=1, + ) + return tool, calls + + +async def _run_tool_bounded(tool, prompts): + """Await the tool under a hard real-time bound so a spin cannot hang CI.""" + try: + return await asyncio.wait_for(tool(prompts), timeout=_HARD_BOUND_SECONDS) + except asyncio.TimeoutError: + pytest.fail( + f"mobile_use_tool did not return within {_HARD_BOUND_SECONDS}s; " + "the result-polling loop no longer honours its deadline" + ) + + +@pytest.mark.asyncio +async def test_result_polling_gives_up_at_the_deadline(monkeypatch): + """A task that never finishes must end at ``timeout_seconds``, not spin forever.""" + tool, calls = _install_fake_backend(monkeypatch, is_success=0, timeout_seconds=12) + + results = await _run_tool_bounded(tool, ["open the app"]) + + assert len(results) == 1 + assert "timed out waiting for result after 12s" in results[0] + # Three polls, each followed by the loop's 5s sleep, cross the 12s deadline. + assert calls["result"] == 3 + assert calls["sleeps"] == [5, 5, 5] + # The pod is still handed back on the timeout path. + assert calls["cancel"] == 1 + + +@pytest.mark.asyncio +async def test_unknown_terminal_status_ends_task_without_spinning(monkeypatch): + """An unrecognised ``IsSuccess`` is terminal: stop at once, do not keep polling.""" + tool, calls = _install_fake_backend( + monkeypatch, is_success=3, timeout_seconds=12, content="device offline" + ) + + results = await _run_tool_bounded(tool, ["open the app"]) + + assert "unknown status 3" in results[0] + assert "device offline" in results[0] + # Terminal means terminal: one poll, no sleep, no fall-through to the deadline. + assert calls["result"] == 1 + assert calls["sleeps"] == [] + assert "timed out" not in results[0] + assert calls["cancel"] == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("in_progress_status", [0, None]) +async def test_in_progress_status_keeps_polling_until_deadline( + monkeypatch, in_progress_status +): + """``IsSuccess`` of ``0``/``None`` means "still running", never "unknown status". + + The service reports ``1`` for success and ``2`` for failure; ``0`` is the + in-progress marker, and a missing ``IsSuccess`` field is decoded as ``None`` + by ``_dict_to_dataclass``. Both are therefore excluded on purpose from the + unknown-terminal-status branch: treating them as terminal would abort every + task on its very first poll. This test pins that decision, so the task must + end via the deadline, not via the unknown-status error. + """ + tool, calls = _install_fake_backend( + monkeypatch, is_success=in_progress_status, timeout_seconds=12 + ) + + results = await _run_tool_bounded(tool, ["open the app"]) + + assert "unknown status" not in results[0] + assert "timed out waiting for result after 12s" in results[0] + # It really kept polling instead of bailing out on the first response. + assert calls["result"] > 1 + assert calls["step"] == calls["result"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wall_clock_step", [600.0, -600.0]) +async def test_pod_acquire_deadline_ignores_wall_clock_jumps( + monkeypatch, wall_clock_step +): + """The pod wait is timed on the monotonic clock, exactly like the result poll. + + Both phases spend the same ``timeout_seconds`` budget, so both have to be + armed on the same clock. Wall-clock arithmetic would let an NTP step forward + abort the wait after a single 1s retry, and an NTP step backwards stretch it + past its budget indefinitely. Here ``monotonic()`` advances by the real 1s + per retry while ``time()`` jumps ten minutes in either direction, so either + regression changes the retry count -- and the faked-sleep budget aborts the + run long before the outer hard bound, so a regression never hangs CI. + """ + tool, calls = _install_fake_backend( + monkeypatch, + is_success=0, + timeout_seconds=12, + pods_available=False, + wall_clock_step_per_sleep=wall_clock_step, + ) + + results = await _run_tool_bounded(tool, ["open the app"]) + + assert "timed out acquiring pod after 12s" in results[0] + # Twelve 1s retries: the budget is spent in monotonic seconds, nothing else. + assert calls["sleeps"] == [1] * 12 + # The task never held a pod, so nothing was started and nothing cancelled. + assert calls["run"] == 0 + assert calls["result"] == 0 + assert calls["cancel"] == 0 + + +@pytest.mark.asyncio +async def test_run_id_is_recorded_and_read_back_through_pod_pool(monkeypatch): + """``task_map`` is reached only through ``PodPool``'s locked accessors. + + ``acquire_pod`` runs on a worker thread (``asyncio.to_thread``), so the run + id written from the event loop and the read in the ``finally`` are the two + places that could skip ``pod_lock``. Spying on the accessors pins the + routing, and the recorded ids show the cancel still targets the run this + task actually started. + """ + seen: dict = {"set": [], "get": []} + real_set = PodPool.set_pod_task + real_get = PodPool.get_pod_task + + def spy_set(self, pid, task_id): + seen["set"].append((pid, task_id)) + return real_set(self, pid, task_id) + + def spy_get(self, pid): + value = real_get(self, pid) + seen["get"].append((pid, value)) + return value + + monkeypatch.setattr(PodPool, "set_pod_task", spy_set) + monkeypatch.setattr(PodPool, "get_pod_task", spy_get) + + tool, calls = _install_fake_backend( + monkeypatch, is_success=1, timeout_seconds=12, content="all done" + ) + + results = await _run_tool_bounded(tool, ["open the app"]) + + assert "task success: all done" in results[0] + assert ("pod1", "run-1") in seen["set"] + assert seen["get"] == [("pod1", "run-1")] + assert calls["cancelled"] == ["run-1"] + + +def test_task_map_stays_consistent_under_concurrent_acquire_release(): + """Concurrent workers never observe another worker's entry in ``task_map``. + + Each worker writes a run id only it can produce and reads it straight back + while it still owns the pod, so a lost or interleaved entry surfaces as a + mismatch. The whole fan-out is bounded once and the executor is torn down + with ``wait=False``: an accessor that deadlocks against ``pod_lock`` (say, + by being called while the lock is already held) fails the test instead of + wedging the suite. + """ + pods = [f"pod-{i}" for i in range(3)] + pool = PodPool(pods) + workers, rounds = 6, 5 + mismatches: list = [] + + def worker(worker_id: int) -> int: + completed = 0 + for round_id in range(rounds): + # Each attempt is already bounded by POD_ACQUIRE_TIMEOUT_SECONDS, so + # a couple of tries cannot hang and leave plenty of slack for a + # loaded machine; a genuinely lost round just lowers `completed`. + pid = pool.acquire_pod() or pool.acquire_pod() + if pid is None: + continue + try: + run_id = f"run-{worker_id}-{round_id}" + pool.set_pod_task(pid, run_id) + observed = pool.get_pod_task(pid) + if observed != run_id: + mismatches.append((pid, run_id, observed)) + completed += 1 + finally: + pool.release_pod(pid) + return completed + + executor = concurrent.futures.ThreadPoolExecutor(max_workers=workers) + try: + futures = [executor.submit(worker, w) for w in range(workers)] + _, not_done = concurrent.futures.wait(futures, timeout=_HARD_BOUND_SECONDS) + if not_done: + pytest.fail( + f"concurrent acquire/release did not settle within " + f"{_HARD_BOUND_SECONDS}s; a task_map accessor is blocking on " + "PodPool.pod_lock" + ) + completed = [future.result() for future in futures] + finally: + # Never wait on a possibly deadlocked worker during teardown. + executor.shutdown(wait=False) + + assert mismatches == [] + assert sum(completed) == workers * rounds + # Every pod came back and no task entry outlived its holder. + assert pool.get_available_count() == len(pods) + assert pool.task_map == {} + for pid in pods: + assert pool.get_pod_status(pid) == "available" + assert pool.get_pod_task(pid) is None diff --git a/tests/tools/builtin_tools/test_tts_bounds.py b/tests/tools/builtin_tools/test_tts_bounds.py new file mode 100644 index 000000000..fc9918b14 --- /dev/null +++ b/tests/tools/builtin_tools/test_tts_bounds.py @@ -0,0 +1,362 @@ +# 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. + +"""Regression tests for the unbounded waits removed from ``text_to_speech``. + +``text_to_speech`` is a synchronous function tool: ADK invokes it inline on the +event loop, so anything that blocks forever inside it stalls the whole process. +Four places could do exactly that, and each is pinned here: + +1. the streaming ``session.post`` carried no timeout at all; +2. the ``iter_lines`` loop had no wall-clock deadline -- a server that streams + valid ``code == 0`` frames forever resets the socket read timeout on every + frame, so only a deadline ends the loop; +3. the player thread called ``audio_queue.task_done()`` only on the success + path, so a raising ``output_stream.write`` left the queue with unfinished + tasks and wedged ``audio_queue.join()`` in the ``finally``; +4. that same ``audio_queue.join()`` -- which takes no timeout -- ran *before* + the stop event was set, so a player wedged *inside* a blocking + ``output_stream.write`` (a stalled device: the realistic hang) never + reached ``task_done()``, never saw a reason to stop, and left the bounded + thread join below it unreachable. + +Nothing here touches the network or an audio device. ``requests.Session`` is +mocked, and ``veadk.utils.audio_manager`` is replaced by a stub module injected +into ``sys.modules``: that module reads ``pyaudio.paInt16`` at import time, so +it cannot even be imported without pyaudio (absent in CI) and *would* open a +real output device on a box that has it. Stubbing -- rather than +``importorskip`` -- keeps the audio-side branches genuinely exercised in CI +instead of silently skipped. ``settings.tool.vespeech.api_key`` is stubbed for +the same reason: it is a ``cached_property`` that falls back to a live token +fetch. + +Every test that could hang on a regression is bounded, so a regression fails +the suite instead of freezing it. +""" + +import base64 +import contextlib +import json +import queue +import sys +import threading +import time +import types +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from veadk.tools.builtin_tools import tts +from veadk.tools.builtin_tools.tts import _audio_player_thread, text_to_speech +from veadk.utils.http_defaults import ( + DEFAULT_HTTP_TIMEOUT, + DEFAULT_STREAM_BUDGET_SECONDS, +) + +_AUDIO_FRAME = json.dumps({"code": 0, "data": base64.b64encode(b"pcm").decode()}) +_DONE_FRAME = json.dumps({"code": 20000000}) + + +def _tool_context() -> SimpleNamespace: + """A ToolContext-shaped stub. + + ``text_to_speech`` reads exactly one attribute off the tool context -- + ``tool_context._invocation_context.user_id`` -- so that is all this stubs. + """ + return SimpleNamespace(_invocation_context=SimpleNamespace(user_id="test_user")) + + +@pytest.fixture +def tts_env(monkeypatch, tmp_path): + """Credentials and an output directory, with no network access anywhere.""" + monkeypatch.setenv("TOOL_VESPEECH_APP_ID", "test_app_id") + monkeypatch.setenv("TOOL_VESPEECH_SPEAKER", "test_speaker") + monkeypatch.setenv("TOOL_VESPEECH_AUDIO_OUTPUT_PATH", str(tmp_path)) + monkeypatch.setattr( + tts, + "settings", + SimpleNamespace(tool=SimpleNamespace(vespeech=SimpleNamespace(api_key="k"))), + ) + + +@contextlib.contextmanager +def _stubbed_audio(output_stream=None): + """Swap in a fake ``veadk.utils.audio_manager`` for the duration of a block. + + With ``output_stream=None`` the device fails to open, which is the headless + path the tool already takes in CI; passing a stream exercises the player + thread against a mock instead of real hardware. + """ + module = types.ModuleType("veadk.utils.audio_manager") + module.AudioConfig = lambda **kwargs: SimpleNamespace(**kwargs) + module.input_audio_config = {} + module.output_audio_config = {} + if output_stream is None: + module.AudioDeviceManager = MagicMock( + side_effect=RuntimeError("no audio device under test") + ) + else: + device = MagicMock() + device.open_output_stream.return_value = output_stream + module.AudioDeviceManager = MagicMock(return_value=device) + with patch.dict(sys.modules, {"veadk.utils.audio_manager": module}): + yield module + + +def _run_bounded(func, timeout: float = 15.0): + """Run ``func`` on a daemon thread and fail -- never hang -- on a regression. + + Each of these calls loops forever against the pre-fix code, so the bound is + what turns a regression into a red test rather than a stuck CI job. + """ + box: dict = {} + + def target(): + try: + box["value"] = func() + except BaseException as exc: # noqa: BLE001 - re-raised on the main thread + box["error"] = exc + + worker = threading.Thread(target=target, daemon=True) + worker.start() + worker.join(timeout) + assert not worker.is_alive(), f"call did not return within {timeout}s" + if "error" in box: + raise box["error"] + return box["value"] + + +def _assert_queue_drains(audio_queue: queue.Queue, timeout: float = 15.0) -> None: + """Assert ``audio_queue.join()`` returns, without blocking the test on it.""" + joiner = threading.Thread(target=audio_queue.join, daemon=True) + joiner.start() + joiner.join(timeout) + assert not joiner.is_alive(), ( + f"audio_queue.join() still blocked after {timeout}s: " + f"{audio_queue.unfinished_tasks} task(s) never marked done" + ) + + +class _JumpingClock: + """A monotonic clock that leaps forward on every read. + + Deterministic and instant: the deadline trips after a known number of + frames instead of after a real wall-clock wait. + """ + + def __init__(self, step: float): + self._now = 0.0 + self._step = step + + def monotonic(self) -> float: + value = self._now + self._now += self._step + return value + + +def _endless_frames(stop: threading.Event): + """A server that never stops sending valid audio frames.""" + while True: + if stop.is_set(): + # Let the worker thread unwind once the assertions are done. + raise RuntimeError("test tore down the endless stream") + yield _AUDIO_FRAME + + +def test_streaming_post_carries_transfer_timeout(tts_env): + """The streaming POST must pass the shared transfer timeout, not None.""" + response = MagicMock() + response.iter_lines.return_value = [_AUDIO_FRAME, _DONE_FRAME] + + with _stubbed_audio(), patch.object(tts.requests, "Session") as session_cls: + session_cls.return_value.post.return_value = response + result = text_to_speech("hello", _tool_context()) + + assert "saved_audio_path" in result + post = session_cls.return_value.post + post.assert_called_once() + assert post.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + + +def test_endless_frame_stream_trips_wall_clock_deadline(tts_env): + """Valid frames forever must still end the read loop, via the deadline. + + Every frame resets the socket read timeout, so the ``requests`` timeout + alone can never fire here -- only the ``time.monotonic()`` deadline does. + """ + stop = threading.Event() + response = MagicMock() + response.iter_lines.return_value = _endless_frames(stop) + + # First read seeds the deadline (0 + budget); each later read jumps a full + # budget, so the deadline trips on the second pass through the loop. + clock = _JumpingClock(step=DEFAULT_STREAM_BUDGET_SECONDS) + fake_time = SimpleNamespace(monotonic=clock.monotonic, sleep=time.sleep) + + try: + with ( + _stubbed_audio(), + patch.object(tts.requests, "Session") as session_cls, + patch.object(tts, "time", fake_time), + ): + session_cls.return_value.post.return_value = response + result = _run_bounded(lambda: text_to_speech("hello", _tool_context())) + finally: + stop.set() + + # The new TimeoutError must flow through the existing handler, so the tool + # keeps its documented error shape instead of raising into the event loop. + assert isinstance(result, dict) + assert set(result) == {"error"} + assert "not finished within" in result["error"] + + +def test_player_thread_marks_task_done_when_playback_raises(): + """``task_done()`` runs on the failure path too, so the queue can drain.""" + audio_queue: queue.Queue = queue.Queue() + for _ in range(3): + audio_queue.put(b"pcm") + + output_stream = MagicMock() + output_stream.write.side_effect = RuntimeError("audio device went away") + stop_event = threading.Event() + + worker = threading.Thread( + target=_audio_player_thread, + args=(audio_queue, output_stream, stop_event), + daemon=True, + ) + worker.start() + try: + _assert_queue_drains(audio_queue) + finally: + stop_event.set() + worker.join(timeout=15.0) + + assert audio_queue.unfinished_tasks == 0 + assert output_stream.write.call_count == 3 + assert not worker.is_alive() + + +def test_failing_playback_does_not_wedge_the_tool(tts_env): + """End to end: playback raising every time must not block the finally.""" + output_stream = MagicMock() + output_stream.write.side_effect = RuntimeError("audio device went away") + + response = MagicMock() + response.iter_lines.return_value = [_AUDIO_FRAME] * 3 + [_DONE_FRAME] + + with ( + _stubbed_audio(output_stream), + patch.object(tts.requests, "Session") as session_cls, + ): + session_cls.return_value.post.return_value = response + result = _run_bounded(lambda: text_to_speech("hello", _tool_context())) + + assert "saved_audio_path" in result + assert output_stream.write.call_count == 3 + output_stream.close.assert_called_once() + + +def test_blocking_playback_teardown_is_bounded_and_logged(tts_env): + """A player wedged *inside* ``output_stream.write`` must not stall teardown. + + This is the realistic hang, and it is the one nothing queue-shaped can + catch: the write never returns, so ``task_done()`` never runs and + ``audio_queue.join()`` -- which takes no timeout -- never returns either. + Only a wait on the thread itself is bounded, and only if the stop flag is + already set when it is entered, since a wedged thread cannot observe a flag + raised after the wait it is blocking. The write parks on an event the test + releases only once the assertions are done, so the hang here is real rather + than simulated with a mock thread. + """ + entered_write = threading.Event() + release_write = threading.Event() + + def blocking_write(_chunk): + entered_write.set() + # Not released while the tool runs: this is a device that went away + # mid-write, holding the player thread inside PortAudio. + release_write.wait(30.0) + + output_stream = MagicMock() + output_stream.write.side_effect = blocking_write + + response = MagicMock() + response.iter_lines.return_value = [_AUDIO_FRAME] * 3 + [_DONE_FRAME] + + try: + with ( + _stubbed_audio(output_stream), + patch.object(tts.requests, "Session") as session_cls, + patch.object(tts, "_PLAYER_JOIN_TIMEOUT", 0.2), + patch.object(tts, "logger") as logger, + ): + session_cls.return_value.post.return_value = response + result = _run_bounded( + lambda: text_to_speech("hello", _tool_context()), timeout=5.0 + ) + finally: + # Let the wedged daemon player unwind instead of leaking into later tests. + release_write.set() + + assert entered_write.wait(5.0), "playback never reached the blocking write" + assert "saved_audio_path" in result + # Chunks are still unplayed and the thread is still stuck: a correct + # teardown gave up on the thread and closed the device anyway. + output_stream.close.assert_called_once() + assert any( + "did not exit in time" in str(call) for call in logger.error.call_args_list + ) + + +def test_healthy_playback_drains_every_queued_chunk(tts_env): + """A slow but working device still plays out everything already queued. + + This is what put the queue join first in the original teardown, and it has + to survive the fix: the wait on the player is renewed while the player is + still marking chunks done, so a backlog that takes longer than a single + ``_PLAYER_JOIN_TIMEOUT`` window is not cut off mid-sentence. The writes + sleep for real -- a fake clock cannot make another thread take time -- but + only briefly, and ``_run_bounded`` still caps the test. + """ + chunks = 8 + write_seconds = 0.15 + # Deliberately shorter than chunks * write_seconds: a flat bound on the + # join would truncate the backlog here. + join_timeout = 0.5 + + output_stream = MagicMock() + output_stream.write.side_effect = lambda _chunk: time.sleep(write_seconds) + played_before_close: list = [] + output_stream.close.side_effect = lambda: played_before_close.append( + output_stream.write.call_count + ) + + response = MagicMock() + response.iter_lines.return_value = [_AUDIO_FRAME] * chunks + [_DONE_FRAME] + + with ( + _stubbed_audio(output_stream), + patch.object(tts.requests, "Session") as session_cls, + patch.object(tts, "_PLAYER_JOIN_TIMEOUT", join_timeout), + ): + session_cls.return_value.post.return_value = response + result = _run_bounded(lambda: text_to_speech("hello", _tool_context())) + + assert "saved_audio_path" in result + assert output_stream.write.call_count == chunks + # The device is closed only after the last queued chunk has been played. + assert played_before_close == [chunks] diff --git a/tests/tools/skills_tools/test_skills_timeouts.py b/tests/tools/skills_tools/test_skills_timeouts.py new file mode 100644 index 000000000..4cd4f4a37 --- /dev/null +++ b/tests/tools/skills_tools/test_skills_timeouts.py @@ -0,0 +1,166 @@ +# 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. + +"""Regression tests: skills tooling must bound its object-storage transfers. + +The signed-URL up/downloads below talk straight to TOS/minio, which `requests` +would otherwise wait on forever. Every site must carry the shared +`DEFAULT_HTTP_TIMEOUT` -- object storage gets no longer allowance than any +other call. The values themselves live in `veadk.utils.http_defaults` so they +stay tunable without touching these assertions. +""" + +from __future__ import annotations + +from importlib import import_module +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT + + +def _ok_response(**attrs: Any) -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.raise_for_status.return_value = None + for name, value in attrs.items(): + setattr(response, name, value) + return response + + +def test_skills_tool_vestack_download_passes_default_http_timeout( + tmp_path: Path, +) -> None: + from veadk.tools.skills_tools.skills_tool import SkillsTool + + save_path = tmp_path / "web-search.zip" + tool = SkillsTool(skills={}) + + with ( + patch( + "veadk.utils.volcengine_sign.ve_request", + return_value={"Result": {"SignedUrl": "https://minio.test/web-search.zip"}}, + ), + patch( + "veadk.tools.skills_tools.skills_tool.download_url_to_file", + side_effect=lambda url, path: Path(path).write_bytes(b"zip-bytes"), + ) as download, + ): + success = tool._download_skill_via_vestack( + skill=SimpleNamespace(id="s-skillid"), + tos_path="skills/s-skillid/v1/web-search.zip", + cloud_provider="vestack", + access_key="ak", + secret_key="sk", + session_token="token", + skill_name="web-search", + save_path=save_path, + ) + + assert success is True + download.assert_called_once_with("https://minio.test/web-search.zip", save_path) + assert save_path.read_bytes() == b"zip-bytes" + + +def test_download_skills_tool_vestack_download_passes_default_http_timeout( + tmp_path: Path, +) -> None: + # NOTE: this module-level helper is currently unreferenced outside its own + # module, but it is a live copy of the download path above and must not + # regress if it is wired back up. + from veadk.tools.skills_tools.download_skills_tool import ( + _download_skill_via_vestack, + ) + + zip_path = tmp_path / "web-search.zip" + + with ( + patch( + "veadk.utils.volcengine_sign.ve_request", + return_value={"Result": {"SignedUrl": "https://minio.test/web-search.zip"}}, + ), + patch( + "veadk.tools.skills_tools.download_skills_tool.download_url_to_file", + side_effect=lambda url, path: Path(path).write_bytes(b"zip-bytes"), + ) as download, + ): + success = _download_skill_via_vestack( + tos_path="skills/s-skillid/v1/web-search.zip", + skill_name="web-search", + access_key="ak", + secret_key="sk", + session_token="token", + service="agentkit", + region="cn-beijing", + host="agentkit.cn-beijing.volcengineapi.com", + scheme="https", + zip_path=zip_path, + ) + + assert success is True + download.assert_called_once_with("https://minio.test/web-search.zip", zip_path) + assert zip_path.read_bytes() == b"zip-bytes" + + +def test_register_skills_tool_signed_url_upload_passes_default_http_timeout( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # `veadk.tools.skills_tools.__init__` re-exports the function under the + # same name as its module, so plain attribute access hands back the + # function. Go through the module registry to reach the module object. + module = import_module("veadk.tools.skills_tools.register_skills_tool") + + skill_dir = tmp_path / "demo-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_text( + "---\nname: demo-skill\ndescription: demo\n---\n\nbody\n", + encoding="utf-8", + ) + + session_dir = tmp_path / "session" + (session_dir / "outputs").mkdir(parents=True) + monkeypatch.setattr(module, "get_session_path", lambda session_id: session_dir) + + monkeypatch.setenv("CLOUD_PROVIDER", "vestack") + monkeypatch.setenv("VOLCENGINE_ACCESS_KEY", "ak") + monkeypatch.setenv("VOLCENGINE_SECRET_KEY", "sk") + monkeypatch.setenv("SKILL_SPACE_ID", "space-1") + + def fake_ve_request(**kwargs: Any) -> dict: + if kwargs["action"] == "GenTempTosObjectUrl": + return { + "Result": { + "SignedUrl": "https://minio.test/upload", + "TosUrl": "tos://bucket/demo-skill.zip", + } + } + return {"Result": {"SkillId": "s-1"}} + + monkeypatch.setattr(module, "ve_request", fake_ve_request) + + tool_context = SimpleNamespace(session=SimpleNamespace(id="session-1")) + + with patch("requests.put", return_value=_ok_response()) as put: + result = module.register_skills_tool( + skill_local_path=str(skill_dir), + tool_context=tool_context, + ) + + assert result.startswith("Successfully registered skill 'demo-skill'"), result + assert put.call_count == 1 + assert put.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT diff --git a/tests/utils/test_http_defaults.py b/tests/utils/test_http_defaults.py new file mode 100644 index 000000000..8b063f550 --- /dev/null +++ b/tests/utils/test_http_defaults.py @@ -0,0 +1,262 @@ +# 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. + +"""Tests for the shared HTTP timeout defaults. + +The public constants are computed at import time, so anything that exercises +the environment overrides has to reload the module. Every reload here goes +through the `reload_http_defaults` fixture, whose teardown reloads the module +one final time under the ambient environment -- otherwise a mutated module +object would leak into every later test in the same session. +""" + +import importlib +import os +from unittest.mock import patch + +import pytest + +from veadk.utils import http_defaults + + +def _env_vars_read_at_import() -> tuple[str, ...]: + """Environment variables the module reads, recorded from a live reload. + + Deriving the list beats writing one down. A hardcoded list desynchronizes + the moment a constant is renamed: the stale name keeps getting scrubbed + while the name actually read is left alone, so an ambient value for it + fails this file on any machine that sets it. + """ + seen: list[str] = [] + real_getenv = os.getenv + + def _recording_getenv(name: str, default: str | None = None) -> str | None: + seen.append(name) + return real_getenv(name, default) + + with patch.object(http_defaults.os, "getenv", _recording_getenv): + importlib.reload(http_defaults) + # Drop the module built under the patched `getenv`, exactly as the fixture + # teardown does, so discovery leaves no trace. + importlib.reload(http_defaults) + return tuple(dict.fromkeys(seen)) + + +_ENV_VARS = _env_vars_read_at_import() + + +@pytest.fixture +def reload_http_defaults(): + """Reload `http_defaults` under a scrubbed + overridden environment. + + Every variable the module reads is scrubbed unless the test overrides it, + so an ambient value on the developer's machine cannot reach the assertions. + + The teardown reload is unconditional: it runs even if the test body raises, + so the module is always restored to the state it had before the test. + """ + + def _reload(**overrides: str): + with patch.dict(os.environ, overrides, clear=False): + for name in _ENV_VARS: + if name not in overrides: + os.environ.pop(name, None) + return importlib.reload(http_defaults) + + try: + yield _reload + finally: + # `patch.dict` has already restored the ambient environment, so this + # rebuilds exactly the module state the session started with. + importlib.reload(http_defaults) + + +def test_default_tuples_have_the_documented_values(reload_http_defaults): + module = reload_http_defaults() + + assert module.DEFAULT_HTTP_TIMEOUT == (10.0, 60.0) + assert module.DEFAULT_STREAM_BUDGET_SECONDS == 300.0 + + +def test_default_tuples_are_pairs_of_floats(reload_http_defaults): + module = reload_http_defaults() + + for timeout in (module.DEFAULT_HTTP_TIMEOUT,): + assert isinstance(timeout, tuple) + assert len(timeout) == 2 + connect, read = timeout + assert isinstance(connect, float) + assert isinstance(read, float) + + +def test_scalar_defaults_are_floats(reload_http_defaults): + module = reload_http_defaults() + + assert module.DEFAULT_CONNECT_TIMEOUT == 10.0 + assert module.DEFAULT_READ_TIMEOUT == 60.0 + assert module.DEFAULT_STREAM_BUDGET_SECONDS == 300.0 + assert isinstance(module.DEFAULT_CONNECT_TIMEOUT, float) + assert isinstance(module.DEFAULT_READ_TIMEOUT, float) + assert isinstance(module.DEFAULT_STREAM_BUDGET_SECONDS, float) + + +def test_http_timeout_is_built_from_the_scalar_halves(reload_http_defaults): + module = reload_http_defaults() + + assert module.DEFAULT_HTTP_TIMEOUT == ( + module.DEFAULT_CONNECT_TIMEOUT, + module.DEFAULT_READ_TIMEOUT, + ) + + +def test_stream_budget_is_a_total_not_a_socket_gap(reload_http_defaults): + # The stream budget bounds a whole streamed response; the read timeout only + # bounds the gap between two reads. Conflating them is what let an endless + # trickle of valid frames hang `tts` forever, so keep the budget a scalar + # and keep it the larger of the two. + module = reload_http_defaults() + + assert not isinstance(module.DEFAULT_STREAM_BUDGET_SECONDS, tuple) + assert module.DEFAULT_STREAM_BUDGET_SECONDS > module.DEFAULT_READ_TIMEOUT + + +def test_env_float_reads_the_environment_variable(): + with patch.dict(os.environ, {"VEADK_TEST_TIMEOUT": "42.5"}, clear=False): + assert http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0) == 42.5 + + +def test_env_float_parses_integer_strings(): + with patch.dict(os.environ, {"VEADK_TEST_TIMEOUT": "7"}, clear=False): + value = http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0) + + assert value == 7.0 + assert isinstance(value, float) + + +def test_env_float_falls_back_when_unset(): + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("VEADK_TEST_TIMEOUT", None) + assert http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0) == 10.0 + + +def test_env_float_falls_back_when_empty(): + with patch.dict(os.environ, {"VEADK_TEST_TIMEOUT": ""}, clear=False): + assert http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0) == 10.0 + + +@pytest.mark.parametrize("raw", ["abc", "10s", "1,5", "nan-ish", " "]) +def test_env_float_falls_back_on_unparseable_input(raw): + with patch.dict(os.environ, {"VEADK_TEST_TIMEOUT": raw}, clear=False): + assert http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0) == 10.0 + + +def test_env_float_falls_back_on_non_string_value(): + # `float(object())` raises TypeError rather than ValueError; the helper + # must swallow that too instead of exploding at import time. + with patch.object(http_defaults.os, "getenv", return_value=object()): + assert http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0) == 10.0 + + +def test_env_float_clamps_to_the_minimum(): + with patch.dict(os.environ, {"VEADK_TEST_TIMEOUT": "0.001"}, clear=False): + assert http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0) == 1.0 + assert http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0, minimum=5.0) == 5.0 + + +def test_env_float_clamps_negative_values(): + with patch.dict(os.environ, {"VEADK_TEST_TIMEOUT": "-30"}, clear=False): + assert http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0) == 1.0 + + +def test_env_float_does_not_clamp_values_above_the_minimum(): + with patch.dict(os.environ, {"VEADK_TEST_TIMEOUT": "120"}, clear=False): + assert ( + http_defaults._env_float("VEADK_TEST_TIMEOUT", 10.0, minimum=5.0) == 120.0 + ) + + +def test_env_overrides_apply_at_import_time(reload_http_defaults): + module = reload_http_defaults( + VEADK_HTTP_CONNECT_TIMEOUT="3", + VEADK_HTTP_READ_TIMEOUT="17.5", + VEADK_HTTP_STREAM_BUDGET="900", + ) + + assert module.DEFAULT_CONNECT_TIMEOUT == 3.0 + assert module.DEFAULT_READ_TIMEOUT == 17.5 + assert module.DEFAULT_STREAM_BUDGET_SECONDS == 900.0 + assert module.DEFAULT_HTTP_TIMEOUT == (3.0, 17.5) + + +def test_env_overrides_are_clamped_at_import_time(reload_http_defaults): + module = reload_http_defaults( + VEADK_HTTP_CONNECT_TIMEOUT="0", + VEADK_HTTP_READ_TIMEOUT="0.25", + ) + + assert module.DEFAULT_HTTP_TIMEOUT == (1.0, 1.0) + + +def test_bad_env_overrides_keep_the_defaults(reload_http_defaults): + module = reload_http_defaults( + VEADK_HTTP_CONNECT_TIMEOUT="abc", + VEADK_HTTP_READ_TIMEOUT="", + VEADK_HTTP_STREAM_BUDGET="not-a-number", + ) + + assert module.DEFAULT_HTTP_TIMEOUT == (10.0, 60.0) + assert module.DEFAULT_STREAM_BUDGET_SECONDS == 300.0 + + +def test_module_restored_after_reload_fixture(reload_http_defaults): + # Guards the fixture itself: a mutated module must not survive a test. + module = reload_http_defaults(VEADK_HTTP_READ_TIMEOUT="999") + + assert module.DEFAULT_READ_TIMEOUT == 999.0 + assert importlib.reload(http_defaults).DEFAULT_HTTP_TIMEOUT[1] != 999.0 + + +def test_ambient_env_overrides_do_not_reach_the_assertions(reload_http_defaults): + # Guards the fixture's scrub list: a deployment that sets any of these -- + # which is exactly what the docs tell users to do -- must not turn this + # file red. Setting all of them at once covers whichever the module reads. + ambient = {name: "42" for name in _ENV_VARS} + + with patch.dict(os.environ, ambient, clear=False): + module = reload_http_defaults() + + assert module.DEFAULT_HTTP_TIMEOUT == (10.0, 60.0) + assert module.DEFAULT_STREAM_BUDGET_SECONDS == 300.0 + + +def test_every_discovered_env_var_moves_a_public_constant(reload_http_defaults): + # Tripwire on the discovery itself: an empty or stale list would scrub + # nothing useful and quietly restore the ambient-environment bug. Each name + # discovered has to actually change one of the exported constants. + assert _ENV_VARS + + defaults = { + name: getattr(reload_http_defaults(), name) for name in http_defaults.__all__ + } + + for env_var in _ENV_VARS: + module = reload_http_defaults(**{env_var: "123"}) + overridden = {name: getattr(module, name) for name in module.__all__} + + assert overridden != defaults, env_var + + +def test_all_exports_are_present(): + for name in http_defaults.__all__: + assert hasattr(http_defaults, name) diff --git a/tests/utils/test_misc_download_bounds.py b/tests/utils/test_misc_download_bounds.py new file mode 100644 index 000000000..5a005d897 --- /dev/null +++ b/tests/utils/test_misc_download_bounds.py @@ -0,0 +1,291 @@ +# 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. + +"""Regression tests for the unbounded download inside ``read_file_to_bytes``. + +``read_file_to_bytes`` fetches user-supplied media URLs (``Runner`` turns a +``MediaMessage`` into bytes with it, and ``image_generate`` reads generated +images back through it), so the peer -- not VeADK -- decides how long the +transfer runs and how many bytes it delivers. It used to read +``response.content`` in one go behind nothing but ``DEFAULT_HTTP_TIMEOUT``, +whose read half bounds only the gap between two socket reads. A server +dribbling one chunk every 59s resets that gap forever, so the call was bounded +in neither wall-clock time nor memory. + +Two limits are pinned here: ``DEFAULT_STREAM_BUDGET_SECONDS`` as a total +deadline, and ``MAX_DOWNLOAD_BYTES`` as a ceiling on what lands in the heap. + +Nothing here touches the network or sleeps: ``requests.get`` is replaced by a +fake response and ``time.monotonic`` by a clock that leaps a full budget on +every read, so the deadline trips deterministically and instantly. The +trickling stream stops itself after ``_RUNAWAY_CHUNKS`` and the call runs on a +joined thread, so a regression fails the suite twice over instead of hanging +CI. +""" + +import importlib +import os +import threading +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import requests + +from veadk.utils import misc +from veadk.utils.http_defaults import ( + DEFAULT_HTTP_TIMEOUT, + DEFAULT_STREAM_BUDGET_SECONDS, +) + +_URL = "https://example.invalid/media.png" + +# How far a "server" is allowed to trickle before the stream gives up on the +# code under test. Reached only when no deadline is enforced at all. +_RUNAWAY_CHUNKS = 50_000 + + +class _FakeResponse: + """The slice of ``requests.Response`` that ``read_file_to_bytes`` uses. + + ``content`` is implemented the way ``requests`` implements it -- buffer the + whole body -- so the pre-fix code path is reproduced faithfully rather than + quietly returning a mock. + """ + + def __init__(self, chunks, status_code: int = 200): + self._chunks = chunks + self.status_code = status_code + self.closed = False + self.chunks_served = 0 + + def __enter__(self) -> "_FakeResponse": + return self + + def __exit__(self, *exc_info) -> bool: + self.close() + return False + + def close(self) -> None: + self.closed = True + + def raise_for_status(self) -> None: + if self.status_code >= 400: + raise requests.HTTPError(f"{self.status_code} for {_URL}") + + def iter_content(self, chunk_size=None): + for chunk in self._chunks: + self.chunks_served += 1 + yield chunk + + @property + def content(self) -> bytes: + return b"".join(self.iter_content()) + + +class _JumpingClock: + """A monotonic clock that leaps forward on every read. + + Models the trickling server without waiting for one: every chunk arrives + comfortably inside the socket read timeout, yet the transfer as a whole + blows the total budget. + """ + + def __init__(self, step: float): + self._now = 0.0 + self._step = step + + def monotonic(self) -> float: + value = self._now + self._now += self._step + return value + + +def _trickling_chunks(): + """A body that never ends -- until the safety valve gives up on the caller.""" + for _ in range(_RUNAWAY_CHUNKS): + yield b"\0" * 8 + raise RuntimeError( + f"the stream was consumed to exhaustion ({_RUNAWAY_CHUNKS} chunks): " + "no wall-clock budget was enforced" + ) + + +def _run_bounded(func, timeout: float = 15.0): + """Run ``func`` on a daemon thread so a regression fails instead of hanging.""" + box: dict = {} + + def target(): + try: + box["value"] = func() + except BaseException as exc: # noqa: BLE001 - re-raised on the main thread + box["error"] = exc + + worker = threading.Thread(target=target, daemon=True) + worker.start() + worker.join(timeout) + assert not worker.is_alive(), f"call did not return within {timeout}s" + if "error" in box: + raise box["error"] + return box["value"] + + +def test_normal_download_returns_the_exact_bytes(): + """The streamed path must rebuild the body byte for byte.""" + payload = bytes(range(256)) * 40 + chunks = [payload[i : i + 1000] for i in range(0, len(payload), 1000)] + response = _FakeResponse(chunks) + + with patch.object(misc.requests, "get", return_value=response): + result = _run_bounded(lambda: misc.read_file_to_bytes(_URL)) + + assert isinstance(result, bytes) + assert result == payload + + +def test_download_streams_with_the_shared_socket_timeout(): + """Streaming is what makes the budget enforceable, so pin the call shape.""" + response = _FakeResponse([b"data"]) + + with patch.object(misc.requests, "get", return_value=response) as get: + assert misc.read_file_to_bytes(_URL) == b"data" + + get.assert_called_once() + assert get.call_args.kwargs["timeout"] == DEFAULT_HTTP_TIMEOUT + assert get.call_args.kwargs["stream"] is True + assert response.closed, "the connection must be released" + + +def test_trickling_server_trips_the_wall_clock_budget(): + """Chunks forever, each within the read timeout, must still end the call. + + Every chunk resets the socket read timeout, so ``DEFAULT_HTTP_TIMEOUT`` + alone can never fire here -- only the total deadline does. + """ + response = _FakeResponse(_trickling_chunks()) + clock = _JumpingClock(step=DEFAULT_STREAM_BUDGET_SECONDS) + + with ( + patch.object(misc.requests, "get", return_value=response), + patch.object(misc, "time", SimpleNamespace(monotonic=clock.monotonic)), + pytest.raises(requests.exceptions.Timeout) as excinfo, + ): + _run_bounded(lambda: misc.read_file_to_bytes(_URL)) + + assert "not finished within" in str(excinfo.value) + assert response.chunks_served < _RUNAWAY_CHUNKS, "the stream was drained" + assert response.closed + + +def test_oversized_download_is_rejected_before_it_is_buffered(monkeypatch): + """The budget bounds time; only the size cap bounds memory.""" + monkeypatch.setattr(misc, "MAX_DOWNLOAD_BYTES", 16) + response = _FakeResponse([b"\0" * 8] * 100) + + with ( + patch.object(misc.requests, "get", return_value=response), + pytest.raises(ValueError) as excinfo, + ): + _run_bounded(lambda: misc.read_file_to_bytes(_URL)) + + assert "16 byte limit" in str(excinfo.value) + # Three chunks: the first two fill the cap, the third exceeds it. The rest + # of the body is never pulled into memory. + assert response.chunks_served == 3 + assert response.closed + + +def test_http_errors_still_propagate(): + """Streaming must not change how a failed response is reported.""" + response = _FakeResponse([b"nope"], status_code=404) + + with ( + patch.object(misc.requests, "get", return_value=response), + pytest.raises(requests.HTTPError), + ): + misc.read_file_to_bytes(_URL) + + assert response.closed + + +def test_local_paths_are_untouched(tmp_path): + """Only the http(s) branch changed; the file branch must still round-trip.""" + payload = b"\x00\x01local bytes\xff" + path = tmp_path / "media.bin" + path.write_bytes(payload) + + assert misc.read_file_to_bytes(str(path)) == payload + + +def test_download_url_to_file_streams_and_replaces_atomically(tmp_path): + """Disk downloads use the same bounded stream without buffering the body.""" + destination = tmp_path / "skill.zip" + destination.write_bytes(b"old") + response = _FakeResponse([b"zip-", b"bytes"]) + + with patch.object(misc.requests, "get", return_value=response) as get: + downloaded = misc.download_url_to_file(_URL, destination) + + assert downloaded == len(b"zip-bytes") + assert destination.read_bytes() == b"zip-bytes" + assert response.closed + assert get.call_args.kwargs == { + "timeout": DEFAULT_HTTP_TIMEOUT, + "stream": True, + } + assert list(tmp_path.glob("*.part")) == [] + + +def test_failed_disk_download_preserves_existing_destination(monkeypatch, tmp_path): + """An oversized partial response must not replace a valid cached archive.""" + destination = tmp_path / "skill.zip" + destination.write_bytes(b"known-good") + monkeypatch.setattr(misc, "MAX_DOWNLOAD_BYTES", 4) + response = _FakeResponse([b"1234", b"5"]) + + with ( + patch.object(misc.requests, "get", return_value=response), + pytest.raises(ValueError, match="4 byte limit"), + ): + misc.download_url_to_file(_URL, destination) + + assert destination.read_bytes() == b"known-good" + assert list(tmp_path.glob("*.part")) == [] + + +def _reload_misc(override: str | None = None): + """Reload `misc` with the cap override set, or scrubbed if none is given. + + Scrubbing matters: the constant is read at import time, so an ambient + `VEADK_MAX_DOWNLOAD_BYTES` on a developer's machine would otherwise decide + what the default assertions see. + """ + env = {"VEADK_MAX_DOWNLOAD_BYTES": override} if override is not None else {} + with patch.dict(os.environ, env, clear=False): + if override is None: + os.environ.pop("VEADK_MAX_DOWNLOAD_BYTES", None) + return importlib.reload(misc) + + +def test_size_cap_is_overridable_from_the_environment(): + """Same env-var style as `http_defaults`: read once, at import time.""" + try: + assert _reload_misc("1024").MAX_DOWNLOAD_BYTES == 1024 + # A malformed override falls back to the default rather than + # disabling the cap. + assert _reload_misc("not-a-number").MAX_DOWNLOAD_BYTES == 256 * 1024 * 1024 + assert _reload_misc().MAX_DOWNLOAD_BYTES == 256 * 1024 * 1024 + finally: + # Restore the module object the rest of the session imported. + importlib.reload(misc) diff --git a/veadk/a2a/hub/a2a_hub_client.py b/veadk/a2a/hub/a2a_hub_client.py index bdc77c104..8abacf1b1 100644 --- a/veadk/a2a/hub/a2a_hub_client.py +++ b/veadk/a2a/hub/a2a_hub_client.py @@ -15,6 +15,10 @@ import requests from a2a.types import AgentCard +from veadk.utils.http_defaults import DEFAULT_CONNECT_TIMEOUT, DEFAULT_HTTP_TIMEOUT + +HEALTH_CHECK_TIMEOUT: tuple[float, float] = (DEFAULT_CONNECT_TIMEOUT, 5.0) + class A2AHubClient: def __init__(self, server_host: str, server_port: int): @@ -24,7 +28,10 @@ def __init__(self, server_host: str, server_port: int): def health_check(self) -> None: """Check the health of the server.""" - response = requests.get(f"http://{self.server_host}:{self.server_port}/ping") + response = requests.get( + f"http://{self.server_host}:{self.server_port}/ping", + timeout=HEALTH_CHECK_TIMEOUT, + ) assert response.status_code == 200, ( f"unexpected status code from A2A hub server: {response.status_code}" ) @@ -36,7 +43,8 @@ def get_agent_cards( ret = [] response = requests.get( - f"http://{self.server_host}:{self.server_port}/group/{group_id}/agents" + f"http://{self.server_host}:{self.server_port}/group/{group_id}/agents", + timeout=DEFAULT_HTTP_TIMEOUT, ).json() agent_infos = response["agent_infos"] for agent_info in agent_infos: @@ -57,6 +65,7 @@ def register_agent(self, group_id: str, agent_id: str, agent_card: AgentCard): "agent_id": agent_id, "agent_card": agent_card.model_dump(), }, + timeout=DEFAULT_HTTP_TIMEOUT, ) assert response.status_code == 200, ( @@ -69,6 +78,7 @@ def create_group(self, group_id: str): params={ "group_id": group_id, }, + timeout=DEFAULT_HTTP_TIMEOUT, ) assert response.status_code == 200, ( diff --git a/veadk/a2a/remote_ve_agent.py b/veadk/a2a/remote_ve_agent.py index 1eddedaa9..e8fbc7ffe 100644 --- a/veadk/a2a/remote_ve_agent.py +++ b/veadk/a2a/remote_ve_agent.py @@ -24,6 +24,7 @@ from veadk.integrations.ve_identity.utils import generate_headers from veadk.utils.auth import VE_TIP_TOKEN_CREDENTIAL_KEY, VE_TIP_TOKEN_HEADER +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT from veadk.utils.logger import get_logger from google.adk.utils.context_utils import Aclosing from google.adk.events.event import Event @@ -172,6 +173,7 @@ def __init__( effective_url + AGENT_CARD_WELL_KNOWN_PATH, headers=req_headers, params=req_params, + timeout=DEFAULT_HTTP_TIMEOUT, ).json() # replace agent_card_url with actual host agent_card_dict["url"] = effective_url diff --git a/veadk/a2a/ve_task_store.py b/veadk/a2a/ve_task_store.py index 0c1555dda..671897716 100644 --- a/veadk/a2a/ve_task_store.py +++ b/veadk/a2a/ve_task_store.py @@ -12,26 +12,53 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio + +from a2a.server.context import ServerCallContext from a2a.server.tasks import TaskStore from a2a.types import Task from typing_extensions import override +from veadk.utils.logger import get_logger + +logger = get_logger(__name__) + class VeTaskStore(TaskStore): - def __init__(self): + """In-process implementation of the A2A `TaskStore` interface. + + Tasks are held in a dictionary owned by this instance, so they are visible + only to the server process that stored them and are lost when it stops. + + Every method takes the `context` argument required by `TaskStore` -- the + A2A request handlers pass it positionally -- but this store does not scope + tasks by caller, so the value is unused. + """ + + def __init__(self) -> None: super().__init__() + self._tasks: dict[str, Task] = {} + self._lock = asyncio.Lock() @override - async def save(self, task: Task) -> None: + async def save(self, task: Task, context: ServerCallContext | None = None) -> None: """Saves or updates a task in the store.""" - return None + async with self._lock: + self._tasks[task.id] = task @override - async def get(self, task_id: str) -> Task | None: + async def get( + self, task_id: str, context: ServerCallContext | None = None + ) -> Task | None: """Retrieves a task from the store by ID.""" - return None + async with self._lock: + return self._tasks.get(task_id) @override - async def delete(self, task_id: str) -> None: + async def delete( + self, task_id: str, context: ServerCallContext | None = None + ) -> None: """Deletes a task from the store by ID.""" - return None + async with self._lock: + if self._tasks.pop(task_id, None) is None: + logger.warning(f"Attempted to delete nonexistent task: {task_id}") diff --git a/veadk/agent.py b/veadk/agent.py index 17740dc42..393437ca7 100644 --- a/veadk/agent.py +++ b/veadk/agent.py @@ -133,6 +133,15 @@ class Agent(LlmAgent): is ignored when `model_api_key` or the MODEL_AGENT_API_KEY env is set.""" model_extra_config: dict = Field(default_factory=dict) tool_thread_pool_config: Optional[ToolThreadPoolConfig] = None + """Thread pool for tool execution. Left unset on purpose. + + ADK's pool only offloads *sync* `FunctionTool`s to a worker thread; every + other tool -- async functions, and any `BaseTool` subclass without a `func` + attribute (MCP tools, `AgentTool`, `SkillsTool`) -- is instead run via + `asyncio.run` on a **new event loop** per call. That breaks `asyncio.Lock`s + and MCP sessions held across calls, and the pool is process-global, so + nested tool calls can exhaust it and deadlock. Enable it per agent only to + hunt blocking I/O inside async tools, which is what ADK built it for.""" tools: list[ToolUnion] = [] diff --git a/veadk/cli/cli_uploadevalset.py b/veadk/cli/cli_uploadevalset.py index 1e58526be..01f353697 100644 --- a/veadk/cli/cli_uploadevalset.py +++ b/veadk/cli/cli_uploadevalset.py @@ -52,6 +52,7 @@ def uploadevalset( import json import requests from veadk.config import getenv + from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT from pathlib import Path if not cozeloop_workspace_id: @@ -130,6 +131,7 @@ def uploadevalset( "is_skip_invalid_items": True, "items": items, }, + timeout=DEFAULT_HTTP_TIMEOUT, ) if response.status_code == 200: diff --git a/veadk/community/langchain_ai/store/memory/viking_memory.py b/veadk/community/langchain_ai/store/memory/viking_memory.py index acc75502b..0dd920178 100644 --- a/veadk/community/langchain_ai/store/memory/viking_memory.py +++ b/veadk/community/langchain_ai/store/memory/viking_memory.py @@ -14,13 +14,13 @@ from __future__ import annotations +import asyncio import json from collections.abc import Iterable from langgraph.store.base import ( BaseStore, GetOp, - ListNamespacesOp, Op, PutOp, Result, @@ -60,9 +60,11 @@ def batch(self, ops: Iterable[Op]) -> list[Result]: return results - def abatch( - self, ops: Iterable[GetOp | SearchOp | PutOp | ListNamespacesOp] - ) -> list[Result]: ... + async def abatch(self, ops: Iterable[Op]) -> list[Result]: + # The VikingDB backend client is fully synchronous and performs + # blocking HTTP calls, so run `batch` in a worker thread instead of + # blocking the running event loop. + return await asyncio.to_thread(self.batch, ops) def _apply_put_op(self, op: PutOp) -> None: index, user_id = op.namespace diff --git a/veadk/integrations/ve_apig/ve_apig_utils.py b/veadk/integrations/ve_apig/ve_apig_utils.py index f5cfef932..222176827 100644 --- a/veadk/integrations/ve_apig/ve_apig_utils.py +++ b/veadk/integrations/ve_apig/ve_apig_utils.py @@ -23,6 +23,8 @@ import requests +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT + Service = "apig" Version = "2021-03-03" Region = os.getenv("REGION") or "cn-beijing" @@ -159,6 +161,7 @@ def request(method, date, query, header, region, ak, sk, token, action, body): headers=header, params=request_param["query"], data=request_param["body"], + timeout=DEFAULT_HTTP_TIMEOUT, ) return r.json() diff --git a/veadk/integrations/ve_code_pipeline/ve_code_pipeline.py b/veadk/integrations/ve_code_pipeline/ve_code_pipeline.py index 858c3ddc6..b49e4d6c7 100644 --- a/veadk/integrations/ve_code_pipeline/ve_code_pipeline.py +++ b/veadk/integrations/ve_code_pipeline/ve_code_pipeline.py @@ -25,6 +25,7 @@ default_region, normalize_cloud_provider, ) +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT from veadk.utils.volcengine_sign import ve_request logger = get_logger(__name__) @@ -309,6 +310,7 @@ def _set_github_webhook( url=f"https://api.github.com/repos/{owner}/{repo}/hooks", headers=headers, data=json.dumps(webhook_config), + timeout=DEFAULT_HTTP_TIMEOUT, ) if response.status_code == 201: diff --git a/veadk/integrations/ve_cozeloop/ve_cozeloop.py b/veadk/integrations/ve_cozeloop/ve_cozeloop.py index 1471f8815..64bf43a64 100644 --- a/veadk/integrations/ve_cozeloop/ve_cozeloop.py +++ b/veadk/integrations/ve_cozeloop/ve_cozeloop.py @@ -15,6 +15,7 @@ import requests from veadk.consts import DEFAULT_COZELOOP_SPACE_NAME +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT from veadk.utils.logger import get_logger logger = get_logger(__name__) @@ -49,7 +50,9 @@ def create_workspace( "description": "Created by Volcengine Agent Development Kit (VeADK)", } - response = requests.post(URL, headers=headers, json=data) + response = requests.post( + URL, headers=headers, json=data, timeout=DEFAULT_HTTP_TIMEOUT + ) if response.json().get("code") == 0: workspace_id = response.json().get("data").get("id") @@ -79,7 +82,9 @@ def search_workspace_id( "page_size": 50, } - response = requests.get(URL, headers=headers, json=data) + response = requests.get( + URL, headers=headers, json=data, timeout=DEFAULT_HTTP_TIMEOUT + ) if response.json().get("code") == 0: workspaces = response.json().get("data").get("workspaces", []) diff --git a/veadk/integrations/ve_faas/ve_faas_utils.py b/veadk/integrations/ve_faas/ve_faas_utils.py index f540e0a71..d16a1f32e 100644 --- a/veadk/integrations/ve_faas/ve_faas_utils.py +++ b/veadk/integrations/ve_faas/ve_faas_utils.py @@ -31,6 +31,8 @@ import requests from volcenginesdkcore.rest import ApiException +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT + Service = "apig" Version = "2021-03-03" Region = os.getenv("REGION") or "cn-beijing" @@ -318,6 +320,7 @@ def request( headers=header, params=request_param["query"], data=request_param["body"], + timeout=DEFAULT_HTTP_TIMEOUT, ) return r.json() diff --git a/veadk/knowledgebase/backends/context_search_backend.py b/veadk/knowledgebase/backends/context_search_backend.py index 0eef12ee8..aa21ae11b 100644 --- a/veadk/knowledgebase/backends/context_search_backend.py +++ b/veadk/knowledgebase/backends/context_search_backend.py @@ -25,6 +25,9 @@ import veadk.config # noqa E401 from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend +from veadk.utils.http_defaults import ( + DEFAULT_HTTP_TIMEOUT, +) from veadk.utils.logger import get_logger from veadk.utils.volcengine_sign import ve_request import requests @@ -412,7 +415,9 @@ def search(self, query: str, top_k: int = 5) -> list[str]: } url = f"{self.context_search_engine_endpoint}/v2/search" json_data = {"text": query, "size": top_k} - response = requests.post(url, json=json_data, headers=headers) + response = requests.post( + url, json=json_data, headers=headers, timeout=DEFAULT_HTTP_TIMEOUT + ) try: result = response.json() except ValueError: @@ -596,7 +601,12 @@ def _upload_file(self, file_path: str, upload_url: str, headers: dict) -> None: The file is opened in binary mode and streamed to minimize memory usage. """ with open(file_path, "rb") as file_handle: - response = requests.put(upload_url, data=file_handle, headers=headers) + response = requests.put( + upload_url, + data=file_handle, + headers=headers, + timeout=DEFAULT_HTTP_TIMEOUT, + ) if response.status_code not in (200, 201, 204): raise ValueError( f"Upload failed for {file_path}: {response.status_code} {response.text}" diff --git a/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py b/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py index e5c6e5508..5c9a76982 100644 --- a/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py +++ b/veadk/knowledgebase/backends/vikingdb_knowledge_backend.py @@ -36,6 +36,7 @@ from veadk.integrations.ve_tos.ve_tos import VeTOS from veadk.knowledgebase.backends.base_backend import BaseKnowledgebaseBackend from veadk.knowledgebase.entry import KnowledgebaseEntry +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT from veadk.utils.logger import get_logger from veadk.utils.misc import formatted_timestamp, getenv @@ -775,6 +776,7 @@ def _do_request( url=full_path, headers=request.headers, data=request.body, + timeout=DEFAULT_HTTP_TIMEOUT, ) if not response.ok: logger.error( diff --git a/veadk/runner.py b/veadk/runner.py index f4cd6163c..f366895c8 100644 --- a/veadk/runner.py +++ b/veadk/runner.py @@ -14,6 +14,7 @@ import functools import os +import uuid from types import MethodType from typing import Union @@ -39,7 +40,7 @@ get_event_function_responses, ) from veadk.utils.logger import get_logger -from veadk.utils.misc import formatted_timestamp, read_file_to_bytes +from veadk.utils.misc import read_file_to_bytes logger = get_logger(__name__) @@ -469,7 +470,7 @@ async def run( self, messages: RunnerMessage, user_id: str = "", - session_id: str = f"tmp-session-{formatted_timestamp()}", + session_id: str | None = None, run_config: RunConfig | None = None, save_tracing_data: bool = False, upload_inline_data_to_tos: bool = False, @@ -485,7 +486,8 @@ async def run( Args: messages (RunnerMessage): Input messages (``str``, ``MediaMessage`` or a list of them). user_id (str): Override default user ID; if empty, uses the constructed ``user_id``. - session_id (str): Session ID. Defaults to a timestamp-based temporary ID. + session_id (str | None): Session ID. If ``None``, a fresh UUID-based + temporary ID (``tmp-session-``) is generated for this call. run_config (google.adk.agents.RunConfig | None): Run config; if ``None``, a default config is created using the environment var ``MODEL_AGENT_MAX_LLM_CALLS``. save_tracing_data (bool): Whether to dump tracing data to disk after the run. Defaults to ``False``. @@ -501,6 +503,13 @@ async def run( AssertionError: If a media MIME type is not among ``image/*`` or ``video/*``. Exception: Exceptions from the underlying ADK/Agent execution may propagate. """ + # Resolved per call: a default argument would be evaluated once at + # import time, so every run omitting `session_id` would share one id. + # A UUID also prevents concurrent calls within the same second from + # colliding, which the old timestamp format could not guarantee. + if session_id is None: + session_id = f"tmp-session-{uuid.uuid4().hex}" + if upload_inline_data_to_tos: _upload_inline_data_to_tos = self.upload_inline_data_to_tos self.upload_inline_data_to_tos = upload_inline_data_to_tos diff --git a/veadk/tools/builtin_tools/create_agent/sources/agentkit_knowledge.py b/veadk/tools/builtin_tools/create_agent/sources/agentkit_knowledge.py index 5b92db822..69d2cb54a 100644 --- a/veadk/tools/builtin_tools/create_agent/sources/agentkit_knowledge.py +++ b/veadk/tools/builtin_tools/create_agent/sources/agentkit_knowledge.py @@ -19,6 +19,7 @@ import asyncio import os import re +import time from dataclasses import dataclass from typing import Any @@ -34,6 +35,14 @@ resolve_cloud_credentials, ) from veadk.utils.cloud_provider import cloud_provider_from_env +from veadk.utils.http_defaults import ( + DEFAULT_CONNECT_TIMEOUT, + DEFAULT_READ_TIMEOUT, +) + +# Wall-clock ceiling for the whole paginated sweep. A per-request timeout +# bounds one call; only this bounds a hundred of them. +_SWEEP_DEADLINE_SECONDS = 120.0 @dataclass(frozen=True) @@ -88,8 +97,29 @@ async def collect(self, tool_context: Any = None) -> SourceCollection: ) ) + # Per call, never on `self`: one source object serves every session, + # so overlapping sweeps would otherwise reset each other's deadline and + # answer for each other's breach. + deadline = time.monotonic() + _SWEEP_DEADLINE_SECONDS try: - resources = await asyncio.to_thread(self._list_all, credentials) + resources, deadline_exceeded = await asyncio.to_thread( + self._list_all, credentials, deadline + ) + if deadline_exceeded: + return SourceCollection( + resources=resources, + status=ResourceSourceStatus( + source=self.name, + status="error", + count=len(resources), + message=( + "Listing gave up after " + f"{_SWEEP_DEADLINE_SECONDS:.0f}s; returning " + f"{len(resources)} knowledge base(s) collected " + "so far." + ), + ), + ) return SourceCollection( resources=resources, status=ResourceSourceStatus( @@ -103,15 +133,26 @@ async def collect(self, tool_context: Any = None) -> SourceCollection: ) ) - def _list_all(self, credentials: CloudCredentials) -> list[StoredResource]: + def _list_all( + self, + credentials: CloudCredentials, + deadline: float, + ) -> tuple[list[StoredResource], bool]: + """Return the pages collected, and whether the deadline cut them short.""" from agentkit.sdk.knowledge import types as knowledge_types client = self._client_factory(credentials, self.region) resources: list[StoredResource] = [] next_token = "" seen_tokens: set[str] = set() + deadline_exceeded = False for _ in range(100): + if time.monotonic() >= deadline: + # Abandon the remaining pages; `collect` reports the partial + # result rather than blocking the caller any longer. + deadline_exceeded = True + break response = client.list_knowledge_bases( knowledge_types.ListKnowledgeBasesRequest( MaxResults=100, @@ -130,13 +171,19 @@ def _list_all(self, credentials: CloudCredentials) -> list[StoredResource]: if resource is not None: resources.append(resource) + # A slow final page can cross the sweep deadline even when it has + # no continuation token, so re-check before declaring success. + if time.monotonic() >= deadline: + deadline_exceeded = True + break + token = str(response.next_token or "") if not token or token in seen_tokens: break seen_tokens.add(token) next_token = token - return resources + return resources, deadline_exceeded def _to_resource(self, item: Any) -> StoredResource | None: knowledge_id = str(getattr(item, "knowledge_id", "") or "").strip() @@ -179,12 +226,23 @@ def _default_client_factory(credentials: CloudCredentials, region: str): from agentkit.sdk.knowledge.client import AgentkitKnowledgeClient with default_cloud_provider(cloud_provider_from_env()): - return AgentkitKnowledgeClient( + client = AgentkitKnowledgeClient( access_key=credentials.access_key, secret_key=credentials.secret_key, session_token=credentials.session_token, region=region, ) + # The client takes no timeout argument, but it subclasses the volcengine + # `Service`, whose setters are re-read on every request. Guarded so an SDK + # shape change degrades to the SDK default instead of raising. + for setter, seconds in ( + ("set_connection_timeout", DEFAULT_CONNECT_TIMEOUT), + ("set_socket_timeout", DEFAULT_READ_TIMEOUT), + ): + apply_timeout = getattr(client, setter, None) + if callable(apply_timeout): + apply_timeout(int(seconds)) + return client def agentkit_viking_index(payload: AgentKitKnowledgePayload) -> str: diff --git a/veadk/tools/builtin_tools/create_agent/sources/skills.py b/veadk/tools/builtin_tools/create_agent/sources/skills.py index 53cb6c54a..ae913f945 100644 --- a/veadk/tools/builtin_tools/create_agent/sources/skills.py +++ b/veadk/tools/builtin_tools/create_agent/sources/skills.py @@ -18,6 +18,7 @@ import asyncio import os +import time from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from typing import Any, Literal @@ -38,6 +39,10 @@ ) from veadk.tools.builtin_tools.create_agent.sources.base import SourceCollection from veadk.utils.cloud_provider import cloud_provider_from_env +from veadk.utils.http_defaults import ( + DEFAULT_CONNECT_TIMEOUT, + DEFAULT_READ_TIMEOUT, +) FINDSKILL_SEARCH_URL = os.getenv( "FINDSKILL_SEARCH_URL", @@ -46,6 +51,9 @@ FindSkillSearcher = Callable[[str], Awaitable[dict[str, Any]]] _PAGE_SIZE = 100 _MAX_PAGES = 100 +# Wall-clock ceiling for one sweep: spaces are paginated, then every space is +# paginated again. A per-request timeout bounds one call, not the fan-out. +_SWEEP_DEADLINE_SECONDS = 120.0 @dataclass(frozen=True) @@ -260,12 +268,36 @@ async def collect(self, tool_context: Any = None) -> SourceCollection: if credentials is None: return self._status("skipped", "AK/SK or STS credentials are unavailable.") + # Per call, never on `self`: one source object serves every session, + # so overlapping sweeps would otherwise reset each other's deadline and + # answer for each other's breach. + deadline = time.monotonic() + _SWEEP_DEADLINE_SECONDS try: - spaces = await asyncio.to_thread(self._list_spaces, credentials) + spaces, spaces_exceeded = await asyncio.to_thread( + self._list_spaces, credentials, deadline + ) if self.space_ids: allowed = set(self.space_ids) spaces = [space for space in spaces if space.id in allowed] - resources = await self._collect_space_skills(credentials, spaces) + if spaces_exceeded: + # The budget is already spent. Do not instantiate one client + # per collected Space merely to have every worker rediscover + # the same expired deadline. + return self._status( + "error", + f"Sweep gave up after {_SWEEP_DEADLINE_SECONDS:.0f}s; " + "returning 0 Skill(s) collected so far.", + ) + resources, skills_exceeded = await self._collect_space_skills( + credentials, spaces, deadline + ) + if skills_exceeded: + return self._status( + "error", + f"Sweep gave up after {_SWEEP_DEADLINE_SECONDS:.0f}s; " + f"returning {len(resources)} Skill(s) collected so far.", + resources=resources, + ) return SourceCollection( resources=resources, status=ResourceSourceStatus( @@ -281,23 +313,37 @@ def _status( self, status: Literal["skipped", "error"], message: str, + resources: Sequence[StoredResource] = (), ) -> SourceCollection: return SourceCollection( + resources=list(resources), status=ResourceSourceStatus( source=self.name, status=status, + count=len(resources), message=message, - ) + ), ) - def _list_spaces(self, credentials: CloudCredentials) -> list[_AgentKitSkillSpace]: + def _list_spaces( + self, + credentials: CloudCredentials, + deadline: float, + ) -> tuple[list[_AgentKitSkillSpace], bool]: + """Return the Spaces collected, and whether the deadline cut them short.""" from agentkit.sdk.skills.types import ListSkillSpacesRequest client = self._client_factory(credentials, self.region) spaces: list[_AgentKitSkillSpace] = [] seen_ids: set[str] = set() collected_count = 0 + deadline_exceeded = False for page in range(1, _MAX_PAGES + 1): + if time.monotonic() >= deadline: + # Abandon the remaining pages; `collect` reports the partial + # result rather than passing it off as a complete listing. + deadline_exceeded = True + break response = client.list_skill_spaces( ListSkillSpacesRequest(PageNumber=page, PageSize=_PAGE_SIZE) ) @@ -315,6 +361,11 @@ def _list_spaces(self, credentials: CloudCredentials) -> list[_AgentKitSkillSpac project_name=str(getattr(item, "project_name", "") or ""), ) ) + # The last page can finish after the deadline and still advertise + # no successor, so check the clock before reporting a complete list. + if time.monotonic() >= deadline: + deadline_exceeded = True + break if not _has_next_page( response, collected_count=collected_count, @@ -324,38 +375,57 @@ def _list_spaces(self, credentials: CloudCredentials) -> list[_AgentKitSkillSpac break else: raise RuntimeError("AgentKit Skill Space pagination exceeded 100 pages.") - return spaces + return spaces, deadline_exceeded async def _collect_space_skills( self, credentials: CloudCredentials, spaces: Sequence[_AgentKitSkillSpace], - ) -> list[StoredResource]: + deadline: float, + ) -> tuple[list[StoredResource], bool]: semaphore = asyncio.Semaphore(self._max_concurrency) - async def collect_one(space: _AgentKitSkillSpace) -> list[StoredResource]: + async def collect_one( + space: _AgentKitSkillSpace, + ) -> tuple[list[StoredResource], bool]: + if time.monotonic() >= deadline: + return [], True async with semaphore: + if time.monotonic() >= deadline: + return [], True return await asyncio.to_thread( self._list_skills, credentials, space, + deadline, ) results = await asyncio.gather(*(collect_one(space) for space in spaces)) - return [resource for group in results for resource in group] + resources = [resource for group, _ in results for resource in group] + return resources, any(exceeded for _, exceeded in results) def _list_skills( self, credentials: CloudCredentials, space: _AgentKitSkillSpace, - ) -> list[StoredResource]: + deadline: float, + ) -> tuple[list[StoredResource], bool]: + """Return the Skills collected, and whether the deadline cut them short.""" from agentkit.sdk.skills.types import ListSkillsBySkillSpaceRequest + if time.monotonic() >= deadline: + return [], True client = self._client_factory(credentials, self.region) resources: list[StoredResource] = [] seen_ids: set[str] = set() collected_count = 0 + deadline_exceeded = False for page in range(1, _MAX_PAGES + 1): + if time.monotonic() >= deadline: + # Abandon the remaining pages; `collect` reports the partial + # result rather than passing it off as a complete listing. + deadline_exceeded = True + break response = client.list_skills_by_skill_space( ListSkillsBySkillSpaceRequest( SkillSpaceId=space.id, @@ -371,6 +441,11 @@ def _list_skills( continue seen_ids.add(resource.descriptor.ref) resources.append(resource) + # Check after each blocking request as well as before it. Otherwise + # a slow final page would be labelled complete after the budget. + if time.monotonic() >= deadline: + deadline_exceeded = True + break if not _has_next_page( response, collected_count=collected_count, @@ -382,7 +457,7 @@ def _list_skills( raise RuntimeError( f"AgentKit Skill pagination exceeded 100 pages for Space '{space.id}'." ) - return resources + return resources, deadline_exceeded def _to_agentkit_resource( self, @@ -433,12 +508,23 @@ def _default_agentkit_client_factory( from agentkit.sdk.skills.client import AgentkitSkillsClient with default_cloud_provider(cloud_provider_from_env()): - return AgentkitSkillsClient( + client = AgentkitSkillsClient( access_key=credentials.access_key, secret_key=credentials.secret_key, session_token=credentials.session_token, region=region, ) + # The client takes no timeout argument, but it subclasses the volcengine + # `Service`, whose setters are re-read on every request. Guarded so an SDK + # shape change degrades to the SDK default instead of raising. + for setter, seconds in ( + ("set_connection_timeout", DEFAULT_CONNECT_TIMEOUT), + ("set_socket_timeout", DEFAULT_READ_TIMEOUT), + ): + apply_timeout = getattr(client, setter, None) + if callable(apply_timeout): + apply_timeout(int(seconds)) + return client def resolve_agentkit_skill( diff --git a/veadk/tools/builtin_tools/image_edit.py b/veadk/tools/builtin_tools/image_edit.py index 371b08384..3b9f86dcf 100644 --- a/veadk/tools/builtin_tools/image_edit.py +++ b/veadk/tools/builtin_tools/image_edit.py @@ -12,11 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import base64 import json import traceback from typing import Dict +import httpx from google.adk.tools import ToolContext from opentelemetry import trace from opentelemetry.trace import Span @@ -27,11 +29,21 @@ DEFAULT_IMAGE_EDIT_MODEL_API_BASE, DEFAULT_IMAGE_EDIT_MODEL_NAME, ) +from veadk.utils.http_defaults import DEFAULT_CONNECT_TIMEOUT from veadk.utils.logger import get_logger from veadk.version import VERSION logger = get_logger(__name__) +# Ark defaults to a 600s read timeout with two retries, so a single hung +# `images.generate` can hold the tool for ~30 minutes -- and `image_edit` +# repeats it for every item in `params`. Image generation is slow, but not +# that slow: three minutes is roughly an order of magnitude above a normal +# edit, and one retry is enough because failures are already reported +# per item through `error_list`. +DEFAULT_IMAGE_EDIT_READ_TIMEOUT: float = 180.0 +DEFAULT_IMAGE_EDIT_MAX_RETRIES: int = 1 + def _get_api_key() -> str: """Resolve credentials only when the tool is actually executed.""" @@ -49,6 +61,11 @@ def _get_client() -> Ark: return Ark( api_key=_get_api_key(), base_url=getenv("MODEL_EDIT_API_BASE", DEFAULT_IMAGE_EDIT_MODEL_API_BASE), + timeout=httpx.Timeout( + timeout=DEFAULT_IMAGE_EDIT_READ_TIMEOUT, + connect=DEFAULT_CONNECT_TIMEOUT, + ), + max_retries=DEFAULT_IMAGE_EDIT_MAX_RETRIES, ) @@ -147,7 +164,10 @@ async def image_edit( "parts.1.image_url.name": "origin_image", "parts.1.image_url.url": origin_image, } - response = client.images.generate( + # The Ark client is synchronous, so awaiting it inline would + # block the event loop for the whole generation. + response = await asyncio.to_thread( + client.images.generate, model=getenv("MODEL_EDIT_NAME", DEFAULT_IMAGE_EDIT_MODEL_NAME), **inputs, extra_headers={ @@ -176,7 +196,8 @@ async def image_edit( image = item.b64_json image_bytes = base64.b64decode(image) - tos_url = _upload_image_to_tos( + tos_url = await asyncio.to_thread( + _upload_image_to_tos, image_bytes=image_bytes, object_key=f"{image_name}.png", ) diff --git a/veadk/tools/builtin_tools/llm_shield.py b/veadk/tools/builtin_tools/llm_shield.py index 9df9922dd..971d32663 100644 --- a/veadk/tools/builtin_tools/llm_shield.py +++ b/veadk/tools/builtin_tools/llm_shield.py @@ -11,6 +11,7 @@ # 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 asyncio import json import os import time @@ -19,6 +20,7 @@ from volcenginesdkllmshield.models.llm_shield_sign import request_sign from google.adk.plugins import BasePlugin +from google.adk.agents.base_agent import BaseAgent from google.adk.agents.callback_context import CallbackContext from google.adk.tools.tool_context import ToolContext from google.adk.models import LlmRequest, LlmResponse @@ -32,6 +34,19 @@ logger = get_logger(__name__) +def _first_not_none(*values: Any) -> Any: + """Return the first value that is not None, or None if there is none. + + ADK spells some hook arguments differently on its two wirings (`args` vs + `tool_args`, `tool_response` vs `result`), so the hooks below accept both + spellings and normalize them through this helper. + """ + for value in values: + if value is not None: + return value + return None + + class LLMShieldPlugin(BasePlugin): """ LLM Shield Plugin for content moderation and safety. @@ -41,8 +56,23 @@ class LLMShieldPlugin(BasePlugin): It helps detect and block potentially harmful content including sensitive information, prompt injection attacks, and policy violations. + The hooks support both ADK wirings, which call them with different argument + names: the plugin API passes `tool_args` / `result` and an extra `agent`, + while the agent callback API passes `args` / `tool_response`. + Examples: - Basic usage with default settings: + As a runner plugin, moderating every agent in the runner: + ```python + from veadk.tools.builtin_tools.llm_shield import content_safety + runner = Runner( + agent=agent, + app_name=app_name, + session_service=session_service, + plugins=[content_safety], + ) + ``` + + As callbacks on a single agent: ```python from veadk.tools.builtin_tools.llm_shield import content_safety agent = Agent( @@ -269,18 +299,87 @@ def _request_llm_shield( return None - def before_agent_callback( - self, callback_context: CallbackContext, **kwargs + async def _request_llm_shield_async( + self, + message: str, + role: str, + hook_name: Optional[str] = None, + session_info: Optional[Dict[str, str]] = None, + ) -> Optional[str]: + """ + Run `_request_llm_shield` in a worker thread. + + The moderation call is a blocking `requests.post` with a default 50s + timeout and it fires up to four times per agent turn, so it is handed + to a thread instead of stalling the event loop. + + Args: + message (str): The content to be moderated + role (str): The role of the message sender ("user" or "assistant") + hook_name (str, optional): Hook name for the Lumen Moderate endpoint + session_info (dict, optional): Session and run ids for the Lumen endpoint + + Returns: + Optional[str]: A blocking message if content violates policies, + None if content is safe or on error + """ + return await asyncio.to_thread( + self._request_llm_shield, + message=message, + role=role, + hook_name=hook_name, + session_info=session_info, + ) + + async def before_agent_callback( + self, + callback_context: Optional[CallbackContext] = None, + *, + agent: Optional[BaseAgent] = None, + **kwargs, ) -> None: + """ + Hook placeholder run before an agent starts. + + Args: + callback_context (CallbackContext, optional): The agent invocation context + agent (BaseAgent, optional): The agent about to run, supplied by the + plugin wiring only + **kwargs: Additional keyword arguments + + Returns: + None: the agent always proceeds normally + """ # TODO: Implement agent-level input validation and context analysis return None - def after_agent_callback(self, callback_context: CallbackContext, **kwargs) -> None: + async def after_agent_callback( + self, + callback_context: Optional[CallbackContext] = None, + *, + agent: Optional[BaseAgent] = None, + **kwargs, + ) -> None: + """ + Hook placeholder run after an agent finishes. + + Args: + callback_context (CallbackContext, optional): The agent invocation context + agent (BaseAgent, optional): The agent that has just run, supplied by + the plugin wiring only + **kwargs: Additional keyword arguments + + Returns: + None: the original agent output is always used + """ # TODO: Implement post-agent analysis and context analysis return None - def before_model_callback( - self, callback_context: CallbackContext, llm_request: LlmRequest, **kwargs + async def before_model_callback( + self, + callback_context: Optional[CallbackContext] = None, + llm_request: Optional[LlmRequest] = None, + **kwargs, ) -> Optional[LlmResponse]: """ Moderate user input before sending to the language model. @@ -290,8 +389,8 @@ def before_model_callback( returns a blocking response instead of allowing the request to proceed. Args: - callback_context (CallbackContext): The callback execution context - llm_request (LlmRequest): The incoming LLM request to moderate + callback_context (CallbackContext, optional): The callback execution context + llm_request (LlmRequest, optional): The incoming LLM request to moderate **kwargs: Additional keyword arguments Returns: @@ -314,7 +413,7 @@ def before_model_callback( if not last_user_message: return None - response = self._request_llm_shield( + response = await self._request_llm_shield_async( message=last_user_message, role="user", hook_name="before_model" ) if response: @@ -328,8 +427,11 @@ def before_model_callback( ) return None - def after_model_callback( - self, callback_context: CallbackContext, llm_response: LlmResponse, **kwargs + async def after_model_callback( + self, + callback_context: Optional[CallbackContext] = None, + llm_response: Optional[LlmResponse] = None, + **kwargs, ) -> Optional[LlmResponse]: """ Moderate model output before returning to the user. @@ -339,8 +441,8 @@ def after_model_callback( instead of the original model output. Args: - callback_context (CallbackContext): The callback execution context - llm_response (LlmResponse): The model's response to moderate + callback_context (CallbackContext, optional): The callback execution context + llm_response (LlmResponse, optional): The model's response to moderate **kwargs: Additional keyword arguments Returns: @@ -362,7 +464,7 @@ def after_model_callback( if not last_model_message: return None - response = self._request_llm_shield( + response = await self._request_llm_shield_async( message=last_model_message, role="assistant", hook_name="after_model" ) if response: @@ -376,8 +478,14 @@ def after_model_callback( ) return None - def before_tool_callback( - self, tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext, **kwargs + async def before_tool_callback( + self, + tool: Optional[BaseTool] = None, + args: Optional[Dict[str, Any]] = None, + tool_context: Optional[ToolContext] = None, + *, + tool_args: Optional[Dict[str, Any]] = None, + **kwargs, ) -> Optional[Dict]: """ Moderate tool arguments before tool execution. @@ -387,22 +495,27 @@ def before_tool_callback( returns a blocking result instead of allowing tool execution. Args: - tool (BaseTool): The tool to be executed - args (Dict[str, Any]): The arguments passed to the tool - tool_context (ToolContext): The tool execution context + tool (BaseTool, optional): The tool to be executed + args (Dict[str, Any], optional): The arguments passed to the tool, + as named by the agent callback wiring + tool_context (ToolContext, optional): The tool execution context + tool_args (Dict[str, Any], optional): The same arguments, as named + by the plugin wiring **kwargs: Additional keyword arguments Returns: Optional[Dict]: A blocking result if arguments are unsafe, None if arguments are safe to proceed """ + args = _first_not_none(args, tool_args) or {} + args_list = [] for key, value in args.items(): args_list.append(f"{key}: {value}") message = "\n".join(args_list) - response = self._request_llm_shield( + response = await self._request_llm_shield_async( message=message, role="user", hook_name="before_tool_call", @@ -413,12 +526,15 @@ def before_tool_callback( return {"result": response} return None - def after_tool_callback( + async def after_tool_callback( self, - tool: BaseTool, - args: Dict[str, Any], - tool_context: CallbackContext, - tool_response: Union[str, Dict[str, Any], List[Any]], + tool: Optional[BaseTool] = None, + args: Optional[Dict[str, Any]] = None, + tool_context: Optional[ToolContext] = None, + tool_response: Optional[Union[str, Dict[str, Any], List[Any]]] = None, + *, + tool_args: Optional[Dict[str, Any]] = None, + result: Optional[Union[str, Dict[str, Any], List[Any]]] = None, **kwargs, ) -> Optional[Dict]: """ @@ -429,16 +545,24 @@ def after_tool_callback( violates safety policies, returns a blocking result. Args: - tool (BaseTool): The tool that was executed - args (Dict[str, Any]): The arguments that were passed to the tool - tool_context (CallbackContext): The tool execution context - tool_response (Union[str, Dict[str, Any], List[Any]]): The tool's response + tool (BaseTool, optional): The tool that was executed + args (Dict[str, Any], optional): The arguments that were passed to the + tool, as named by the agent callback wiring + tool_context (ToolContext, optional): The tool execution context + tool_response (Union[str, Dict[str, Any], List[Any]], optional): The + tool's response, as named by the agent callback wiring + tool_args (Dict[str, Any], optional): The same arguments, as named by + the plugin wiring + result (Union[str, Dict[str, Any], List[Any]], optional): The same tool + response, as named by the plugin wiring **kwargs: Additional keyword arguments Returns: Optional[Dict]: A blocking result if tool output is unsafe, None if output is safe to return """ + tool_response = _first_not_none(tool_response, result) + message = "" if isinstance(tool_response, str): message = tool_response @@ -449,7 +573,7 @@ def after_tool_callback( for item in tool_response: message += f"{item}\n" - response = self._request_llm_shield( + response = await self._request_llm_shield_async( message=message, role="assistant", hook_name="after_tool_call", diff --git a/veadk/tools/builtin_tools/mobile_run.py b/veadk/tools/builtin_tools/mobile_run.py index 65bf8837e..2b223e421 100644 --- a/veadk/tools/builtin_tools/mobile_run.py +++ b/veadk/tools/builtin_tools/mobile_run.py @@ -92,7 +92,7 @@ async def main(): import time from dataclasses import dataclass from typing import Type, TypeVar, List, Dict, Any, Callable -from queue import Queue +from queue import Empty, Queue from threading import Lock from veadk.utils.logger import get_logger @@ -117,6 +117,9 @@ async def main(): "TOOL_MOBILE_USE_TOOL_ID", ] +# Bounded wait for a free pod so callers can re-check their own deadline. +POD_ACQUIRE_TIMEOUT_SECONDS = 1 + class MobileUseToolError(Exception): def __init__(self, msg: str): @@ -209,6 +212,10 @@ def __init__(self, pod_ids: List[str]): self.pod_ids = pod_ids self.available_pods = Queue() self.pod_lock = Lock() + # Invariant: every read and write of `task_map` goes through the + # accessors on this class, which all hold `pod_lock`. `acquire_pod` runs + # on a worker thread (`asyncio.to_thread`), so its write really can + # interleave with an event-loop-side read. self.task_map: Dict[str, str] = {} for pid in pod_ids: @@ -217,16 +224,19 @@ def __init__(self, pod_ids: List[str]): def acquire_pod(self) -> Any | None: try: - pid = self.available_pods.get(block=True) - with self.pod_lock: - self.task_map[pid] = "pending" + pid = self.available_pods.get( + block=True, timeout=POD_ACQUIRE_TIMEOUT_SECONDS + ) + except Empty: logger.debug( - f"Acquired pod: {pid}, available pods: {self.available_pods.qsize()}" + f"Pod acquisition timeout after {POD_ACQUIRE_TIMEOUT_SECONDS}s, no available pod" ) - return pid - except Exception as e: - logger.warning(f"Pod acquisition timeout: {e}") return None + self.set_pod_task(pid, "pending") + logger.debug( + f"Acquired pod: {pid}, available pods: {self.available_pods.qsize()}" + ) + return pid def release_pod(self, pid: str) -> None: with self.pod_lock: @@ -237,6 +247,16 @@ def release_pod(self, pid: str) -> None: f"Released pod: {pid}, available pods: {self.available_pods.qsize()}" ) + # Task bookkeeping is only ever touched through these two accessors, so the + # `pod_lock` discipline lives in this class instead of at every call site. + def set_pod_task(self, pid: str, task_id: str) -> None: + with self.pod_lock: + self.task_map[pid] = task_id + + def get_pod_task(self, pid: str) -> str | None: + with self.pod_lock: + return self.task_map.get(pid) + def get_pod_status(self, pid: str) -> str: with self.pod_lock: return self.task_map.get(pid, "available") @@ -407,7 +427,7 @@ def create_mobile_use_tool( * "You are a mobile testing agent. Follow least-privilege principles and avoid unauthorized access." max_step (int): Maximum execution steps per agent. timeout_seconds (int): - Maximum wait time in seconds. Raises if not finished. Default: 600. + Maximum wait time in seconds. Raises if not finished. Default: 900. step_interval_seconds (int): Status polling interval in seconds. Default: 1. @@ -444,17 +464,22 @@ async def mobile_use_tool(user_prompts: List[str]) -> list[None]: coroutines = [] def task_worker(index: int, prompt: str) -> Callable: - wait_start = time.time() + # One budget, one clock: waiting for a pod and polling for its + # result are consecutive phases of the same `timeout_seconds`, so + # they share a single deadline armed here. Monotonic only - an NTP + # step must not stretch (step back) or cut short (step forward) a + # wait, which is what a `time.time()` deadline would allow. + deadline = time.monotonic() + timeout_seconds async def run(): nonlocal results pod_id = None try: while True: - pod_id = pod_pool.acquire_pod() + pod_id = await asyncio.to_thread(pod_pool.acquire_pod) if pod_id: break - if time.time() - wait_start >= timeout_seconds: + if time.monotonic() >= deadline: raise MobileUseToolError( f"Task {index} timed out acquiring pod after {timeout_seconds}s" ) @@ -466,7 +491,10 @@ async def run(): logger.info( f"Task {index} assigned to pod: {pod_id}, starting: {prompt}" ) - task_response = _run_agent_task( + # Every ACEP call below is blocking `requests` I/O; run it + # off the event loop so `asyncio.gather` really is concurrent. + task_response = await asyncio.to_thread( + _run_agent_task, system_prompt, prompt, pod_id, @@ -475,11 +503,19 @@ async def run(): timeout_seconds, ) task_id = task_response.Result.RunId - pod_pool.task_map[pod_id] = task_id + pod_pool.set_pod_task(pod_id, task_id) while True: - result_response = _get_task_result(task_id) - if result_response.Result.IsSuccess == 1: + if time.monotonic() >= deadline: + raise MobileUseToolError( + f"Task {index} timed out waiting for result after {timeout_seconds}s" + ) + + result_response = await asyncio.to_thread( + _get_task_result, task_id + ) + status = result_response.Result.IsSuccess + if status == 1: results[index] = ( f"task success: {result_response.Result.Content}\n" ) @@ -487,14 +523,20 @@ async def run(): f"Task {index} succeeded on pod: {pod_id}, result: {result_response.Result.Content}" ) break - elif result_response.Result.IsSuccess == 2: + elif status == 2: results[index] = ( f"task failed: {result_response.Result.Content}" ) logger.error(f"Task {index} failed on pod: {pod_id}") break + elif status not in (None, 0): + raise MobileUseToolError( + f"Task {index} returned unknown status {status} on pod: {pod_id}, content: {result_response.Result.Content}" + ) - current_step = _get_current_step(task_id) + current_step = await asyncio.to_thread( + _get_current_step, task_id + ) if current_step.Result.Results: last_step = current_step.Result.Results[-1] logger.debug( @@ -508,8 +550,16 @@ async def run(): logger.error(error_msg) finally: if pod_id: - _cancel_task(pod_pool.task_map[pod_id]) - pod_pool.release_pod(pod_id) + try: + await asyncio.to_thread( + _cancel_task, pod_pool.get_pod_task(pod_id) + ) + except Exception as e: + # A failed cancel must not skip release_pod, or the + # pod leaks for the lifetime of the process. + logger.error(f"Task {index} cancel failed: {e}") + finally: + pod_pool.release_pod(pod_id) return run diff --git a/veadk/tools/builtin_tools/tts.py b/veadk/tools/builtin_tools/tts.py index 185fe8262..d13bc57c2 100644 --- a/veadk/tools/builtin_tools/tts.py +++ b/veadk/tools/builtin_tools/tts.py @@ -23,10 +23,24 @@ from typing import Dict, Any from google.adk.tools import ToolContext from veadk.config import getenv, settings +from veadk.utils.http_defaults import ( + DEFAULT_HTTP_TIMEOUT, + DEFAULT_STREAM_BUDGET_SECONDS, +) from veadk.utils.logger import get_logger logger = get_logger(__name__) +# bound on how long the audio player thread may take to exit once it has been +# told to stop; a thread that is still marking queued chunks done is granted a +# further window of the same length, so only a stalled one is given up on +_PLAYER_JOIN_TIMEOUT = 5.0 + +# how long the player thread waits for the next chunk before re-checking the +# stop event, and so the worst case delay between the queue running dry and the +# thread exiting +_PLAYER_QUEUE_POLL_TIMEOUT = 0.1 + def text_to_speech(text: str, tool_context: ToolContext) -> Dict[str, Any]: """TTS provides users with the ability to convert text to speech, turning the text content of LLM into audio. @@ -87,7 +101,13 @@ def text_to_speech(text: str, tool_context: ToolContext) -> Dict[str, Any]: try: logger.debug(f"Request TTS server with payload: {payload}.") - response = session.post(url, headers=headers, json=payload, stream=True) + response = session.post( + url, + headers=headers, + json=payload, + stream=True, + timeout=DEFAULT_HTTP_TIMEOUT, + ) os.makedirs(temp_dir, exist_ok=True) with tempfile.NamedTemporaryFile( @@ -158,8 +178,13 @@ def handle_server_response( except Exception as e: logger.error(f"Failed to initialize audio device: {e}") + deadline = time.monotonic() + DEFAULT_STREAM_BUDGET_SECONDS try: for chunk in response.iter_lines(decode_unicode=True): + if time.monotonic() > deadline: + raise TimeoutError( + f"tts response not finished within {DEFAULT_STREAM_BUDGET_SECONDS}s" + ) if not chunk: continue data = json.loads(chunk) @@ -190,37 +215,75 @@ def handle_server_response( raise finally: if output_stream: - audio_queue.join() + # Ask the player to play out what it already holds and then exit. + # Signalling before waiting is what makes the wait bounded: a + # thread wedged inside a blocking output_stream.write can only + # observe the event once that write returns, and one that never + # returns must be abandoned rather than waited on. stop_event.set() - if player_thread and player_thread.is_alive(): - player_thread.join() + if player_thread: + _join_audio_player_thread(player_thread, audio_queue) output_stream.close() +def _join_audio_player_thread(player_thread, audio_queue) -> None: + """Wait for a stopping player thread to drain and exit, but never forever. + + The caller has already set the stop event, so the thread plays whatever is + still queued and then leaves. Each wait is bounded by _PLAYER_JOIN_TIMEOUT + and is only renewed while the thread keeps marking chunks done, so healthy + playback of any length still drains in full while a thread stalled inside + output_stream.write is given up on after a single window. audio_queue.join() + cannot do this job: it takes no timeout, and a stalled thread never reaches + task_done(). Nothing feeds the queue any more, so unfinished_tasks only + falls and the loop runs at most once per queued chunk. + + Args: + player_thread: The already signalled audio player thread. + audio_queue: The queue that thread is draining. + + Returns: + None + """ + pending = audio_queue.unfinished_tasks + while True: + player_thread.join(timeout=_PLAYER_JOIN_TIMEOUT) + if not player_thread.is_alive(): + return + remaining = audio_queue.unfinished_tasks + if remaining >= pending: + logger.error("audio player thread did not exit in time") + return + pending = remaining + + def _audio_player_thread(audio_queue, output_stream, stop_event): """ Play audio data from queue. Args: audio_queue: The queue to store audio data. output_stream: The output stream to play audio. - stop_event: The event to stop the thread. + stop_event: The event asking the thread to play out the queue and exit. Returns: """ - while not stop_event.is_set(): + # stop_event means "finish what is already queued, then exit", so the thread + # only leaves once the queue has actually run dry as well. + while not stop_event.is_set() or not audio_queue.empty(): + try: + audio_data = audio_queue.get(timeout=_PLAYER_QUEUE_POLL_TIMEOUT) + except queue.Empty: + continue try: # write audio data to output stream - audio_data = audio_queue.get(timeout=1.0) if audio_data: output_stream.write(audio_data) - audio_queue.task_done() - except queue.Empty: - # if queue is empty, sleep for a while - time.sleep(0.1) except Exception as e: logger.error(f"Failed to play audio data: {e}") time.sleep(0.1) + finally: + audio_queue.task_done() logger.debug("audio player thread exited") diff --git a/veadk/tools/skills_tools/download_skills_tool.py b/veadk/tools/skills_tools/download_skills_tool.py index 34c7effd3..c220d573f 100644 --- a/veadk/tools/skills_tools/download_skills_tool.py +++ b/veadk/tools/skills_tools/download_skills_tool.py @@ -22,8 +22,9 @@ from veadk.integrations.ve_tos.ve_tos import VeTOS -from veadk.utils.volcengine_sign import ve_request from veadk.utils.logger import get_logger +from veadk.utils.misc import download_url_to_file +from veadk.utils.volcengine_sign import ve_request logger = get_logger(__name__) @@ -41,7 +42,6 @@ def _download_skill_via_vestack( zip_path: Path, ) -> bool: import json - import requests from veadk.utils.volcengine_sign import ve_request try: @@ -96,10 +96,7 @@ def _download_skill_via_vestack( return False else: try: - response = requests.get(signed_url) - response.raise_for_status() - with open(zip_path, "wb") as f: - f.write(response.content) + download_url_to_file(signed_url, zip_path) return True except Exception as e: logger.warning( diff --git a/veadk/tools/skills_tools/register_skills_tool.py b/veadk/tools/skills_tools/register_skills_tool.py index 6c3d5b6ac..d8ed47f72 100644 --- a/veadk/tools/skills_tools/register_skills_tool.py +++ b/veadk/tools/skills_tools/register_skills_tool.py @@ -26,6 +26,7 @@ from veadk.tools.skills_tools.session_path import get_session_path from veadk.integrations.ve_tos.ve_tos import VeTOS from veadk.utils.volcengine_sign import ve_request +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT from veadk.utils.logger import get_logger logger = get_logger(__name__) @@ -185,7 +186,9 @@ def register_skills_tool( try: with open(zip_file_path, "rb") as f: - response = requests.put(signed_url, data=f) + response = requests.put( + signed_url, data=f, timeout=DEFAULT_HTTP_TIMEOUT + ) response.raise_for_status() except Exception as e: logger.error(f"Failed to upload skill '{skill_name}' to minio: {e}") diff --git a/veadk/tools/skills_tools/skills_tool.py b/veadk/tools/skills_tools/skills_tool.py index 41acffe9e..907d48981 100644 --- a/veadk/tools/skills_tools/skills_tool.py +++ b/veadk/tools/skills_tools/skills_tool.py @@ -13,6 +13,7 @@ # limitations under the License. from __future__ import annotations +import asyncio import os from pathlib import Path from typing import Any, Dict @@ -27,6 +28,7 @@ from veadk.tools.skills_tools.session_path import get_session_path from veadk.tracing.telemetry.telemetry import set_common_attributes_on_tool_span from veadk.utils.logger import get_logger +from veadk.utils.misc import download_url_to_file tracer = trace.get_tracer("veadk.skills_tool") @@ -101,7 +103,12 @@ async def run_async( return "Error: No skill name provided" with tracer.start_as_current_span(f"execute_skill {skill_name}") as span: - result = self._invoke_skill(skill_name, tool_context) + # `_invoke_skill` does blocking disk and network I/O (skill + # download). Awaiting it inline would stall the event loop for + # every other session in the process. + result = await asyncio.to_thread( + self._invoke_skill, skill_name, tool_context + ) self._add_skill_span_attributes(span, skill_name, result) self._upload_skill_metrics(span, skill_name, result) return result @@ -415,7 +422,6 @@ def _download_skill_via_vestack( ) -> bool: """Download a skill using the vestack environment GenTempTosObjectDownloadUrl API.""" import json - import requests from veadk.utils.volcengine_sign import ve_request # Extract skill_id and skill_version from TosPath @@ -486,10 +492,7 @@ def _download_skill_via_vestack( return False else: try: - response = requests.get(signed_url) - response.raise_for_status() - with open(save_path, "wb") as f: - f.write(response.content) + download_url_to_file(signed_url, save_path) return True except Exception as e: logger.error( diff --git a/veadk/utils/http_defaults.py b/veadk/utils/http_defaults.py new file mode 100644 index 000000000..71e7a1b17 --- /dev/null +++ b/veadk/utils/http_defaults.py @@ -0,0 +1,69 @@ +# 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. + +"""Shared HTTP timeout defaults for outbound requests. + +Every outbound call in VeADK should carry an explicit timeout. `requests` has +no default timeout at all, so a call without one blocks forever if the peer +accepts the connection and then goes silent -- which, for the synchronous tools +that ADK invokes inline on the event loop, stalls the whole process. + +Note on semantics: for `requests`, the read half of the tuple bounds the time +between two consecutive socket reads, not the total duration of the call. It +caps silence, not transfer size. Anything that needs a hard wall-clock ceiling +(a polling loop, a paginated sweep) still needs its own deadline on top. + +All values are overridable per deployment via environment variables. +""" + +import os + +__all__ = [ + "DEFAULT_CONNECT_TIMEOUT", + "DEFAULT_READ_TIMEOUT", + "DEFAULT_HTTP_TIMEOUT", + "DEFAULT_STREAM_BUDGET_SECONDS", +] + + +def _env_float(name: str, default: float, minimum: float = 1.0) -> float: + raw = os.getenv(name) + if not raw: + return default + try: + return max(minimum, float(raw)) + except (TypeError, ValueError): + return default + + +# Time allowed to establish a TCP/TLS connection. A peer that is unreachable +# should fail fast rather than occupy a worker. +DEFAULT_CONNECT_TIMEOUT: float = _env_float("VEADK_HTTP_CONNECT_TIMEOUT", 10.0) + +# Time allowed between two consecutive reads for ordinary control-plane calls. +DEFAULT_READ_TIMEOUT: float = _env_float("VEADK_HTTP_READ_TIMEOUT", 60.0) + +# (connect, read) tuple, ready to pass straight to `requests`. One value for +# every outbound call: a chunk gap anywhere near a minute already means the peer +# is unhealthy, whether the payload is a JSON reply or a hundred-megabyte object. +DEFAULT_HTTP_TIMEOUT: tuple[float, float] = ( + DEFAULT_CONNECT_TIMEOUT, + DEFAULT_READ_TIMEOUT, +) + +# Wall-clock ceiling for consuming a streamed response end to end. This is a +# different quantity from the read timeout above, which only bounds the gap +# between two reads: a server emitting an endless trickle of valid frames resets +# that gap forever and is only caught by a total budget. +DEFAULT_STREAM_BUDGET_SECONDS: float = _env_float("VEADK_HTTP_STREAM_BUDGET", 300.0) diff --git a/veadk/utils/misc.py b/veadk/utils/misc.py index cbde5763a..80815995e 100644 --- a/veadk/utils/misc.py +++ b/veadk/utils/misc.py @@ -16,16 +16,46 @@ import json import os import sys +import tempfile import time import types -from typing import Any, Dict, List, MutableMapping, Optional, Tuple +from pathlib import Path +from typing import Any, Callable, Dict, List, MutableMapping, Optional, Tuple import requests from yaml import safe_load +from veadk.utils.http_defaults import ( + DEFAULT_HTTP_TIMEOUT, + DEFAULT_STREAM_BUDGET_SECONDS, +) + import __main__ +def _env_int(name: str, default: int, minimum: int = 1) -> int: + """Read a positive int from the environment, mirroring `http_defaults`.""" + raw = os.getenv(name) + if not raw: + return default + try: + return max(minimum, int(raw)) + except (TypeError, ValueError): + return default + + +# Hard ceiling on a single remote download. The wall-clock budget bounds how +# long a transfer may run, not how much it may deliver: 300s on a fast link is +# tens of gigabytes. Some callers materialize the body in memory and others +# write remote skill archives to disk, so both paths share the same limit. +# 256 MiB is far above any generated image, short clip, or normal skill bundle. +MAX_DOWNLOAD_BYTES: int = _env_int("VEADK_MAX_DOWNLOAD_BYTES", 256 * 1024 * 1024) + +# Big enough that per-chunk overhead is noise, small enough that the deadline +# and the size cap are re-checked often during a transfer. +_DOWNLOAD_CHUNK_SIZE = 1024 * 1024 + + def read_file(file_path): with open(file_path, "r", encoding="utf-8") as f: data = f.readlines() @@ -38,11 +68,75 @@ def formatted_timestamp() -> str: return time.strftime("%Y%m%d%H%M%S", time.localtime()) +def _consume_remote_file(file_url: str, consume: Callable[[bytes], Any]) -> int: + """Consume a remote body in bounded chunks and return its byte count.""" + deadline = time.monotonic() + DEFAULT_STREAM_BUDGET_SECONDS + downloaded = 0 + with requests.get(file_url, timeout=DEFAULT_HTTP_TIMEOUT, stream=True) as response: + response.raise_for_status() + for chunk in response.iter_content(chunk_size=_DOWNLOAD_CHUNK_SIZE): + if time.monotonic() > deadline: + raise requests.exceptions.Timeout( + f"download of {file_url} not finished within " + f"{DEFAULT_STREAM_BUDGET_SECONDS}s" + ) + if not chunk: + continue + downloaded += len(chunk) + if downloaded > MAX_DOWNLOAD_BYTES: + raise ValueError( + f"download of {file_url} exceeds the " + f"{MAX_DOWNLOAD_BYTES} byte limit " + "(override with VEADK_MAX_DOWNLOAD_BYTES)" + ) + consume(chunk) + return downloaded + + +def download_url_to_file(file_url: str, destination: str | os.PathLike[str]) -> int: + """Download an HTTP(S) URL to a file with time and size bounds. + + The body is streamed to a temporary sibling and atomically moved into + place only after the complete response succeeds. An existing destination + therefore survives timeouts, oversized bodies, and other transfer errors. + """ + if not file_url.startswith(("http://", "https://")): + raise ValueError(f"download URL must use http(s): {file_url}") + + destination_path = Path(destination) + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + prefix=f".{destination_path.name}.", + suffix=".part", + dir=destination_path.parent, + delete=False, + ) as temporary_file: + temporary_path = Path(temporary_file.name) + downloaded = _consume_remote_file(file_url, temporary_file.write) + os.replace(temporary_path, destination_path) + temporary_path = None + return downloaded + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + + def read_file_to_bytes(file_path: str) -> bytes: + """Read a local path or an http(s) URL into memory. + + Raises `requests.exceptions.Timeout` if a remote transfer outlives + `DEFAULT_STREAM_BUDGET_SECONDS`, and `ValueError` if it exceeds + `MAX_DOWNLOAD_BYTES`. The read half of `DEFAULT_HTTP_TIMEOUT` only bounds + the gap between two socket reads, so a peer trickling a byte every 59s + never trips it and would otherwise stream into `response.content` forever: + unbounded in both wall-clock time and memory. + """ if file_path.startswith(("http://", "https://")): - response = requests.get(file_path) - response.raise_for_status() - return response.content + chunks: List[bytes] = [] + _consume_remote_file(file_path, chunks.append) + return b"".join(chunks) else: with open(file_path, "rb") as f: return f.read() diff --git a/veadk/utils/volcengine_sign.py b/veadk/utils/volcengine_sign.py index 6412247e3..cb32ac4e1 100644 --- a/veadk/utils/volcengine_sign.py +++ b/veadk/utils/volcengine_sign.py @@ -21,10 +21,12 @@ import requests +from veadk.utils.http_defaults import DEFAULT_HTTP_TIMEOUT + # Bounded default (connect, read) timeout in seconds for Volcengine API calls, so a # hung endpoint cannot block the caller forever. Callers with slow control-plane # operations (deploys, large uploads) can override via the ``timeout`` parameter. -DEFAULT_REQUEST_TIMEOUT: tuple[float, float] = (10, 60) +DEFAULT_REQUEST_TIMEOUT: tuple[float, float] = DEFAULT_HTTP_TIMEOUT Service = "" Version = ""