From cfd6b0a5fb9c57685dc2b02ca059fa88d8eff8ec Mon Sep 17 00:00:00 2001 From: lmx-2077 <1024010227lmx@gmail.com> Date: Mon, 17 Aug 2026 21:41:19 +0800 Subject: [PATCH] feat(agent): per-category tool timeouts with graceful degradation (#1890) (#2134) --- .env.example | 9 + docs/CHANGELOG.md | 5 + docs/full-guide.md | 4 + docs/full-guide_EN.md | 4 + main.py | 16 +- src/agent/agents/base_agent.py | 4 +- src/agent/factory.py | 268 +++++- src/agent/research.py | 5 +- src/agent/runner.py | 370 ++++++--- src/agent/tools/execution.py | 53 +- src/agent/tools/registry.py | 46 +- src/config.py | 23 + src/services/system_config_service.py | 6 + tests/agent/test_tool_timeout.py | 1083 +++++++++++++++++++++++++ 14 files changed, 1757 insertions(+), 139 deletions(-) create mode 100644 tests/agent/test_tool_timeout.py diff --git a/.env.example b/.env.example index 3d35f3534..dae0bc819 100644 --- a/.env.example +++ b/.env.example @@ -507,6 +507,15 @@ AGENT_SKILLS= # AGENT_PORTFOLIO_AGENT_TIMEOUT_S=0 # AGENT_SKILL_AGENT_TIMEOUT_S=0 +# Agent 工具按类别默认超时（秒，0=关闭；未配置时回退到全局 tool_call_timeout_seconds 预算） +# 有效超时按 first-wins 解析：显式 per-run tool_call_timeout_seconds > 单工具显式 timeout_seconds > 类别默认值 > 无限制； +# 剩余 wall-clock 预算仅作不可突破的外层 cap；inf/nan/负数视为无限制。 +# market 类工具（get_market_indices / get_sector_rankings 等网络数据调用）复用 AGENT_DATA_TOOL_TIMEOUT_S，无独立开关。 +# AGENT_DATA_TOOL_TIMEOUT_S=0 +# AGENT_SEARCH_TOOL_TIMEOUT_S=0 +# AGENT_ANALYSIS_TOOL_TIMEOUT_S=0 +# AGENT_ACTION_TOOL_TIMEOUT_S=0 + # 策略专家并发数（仅 specialist 模式生效；范围 1-4，默认 3） # AGENT_SKILL_CONCURRENCY=3 diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index ce8e1e364..6972d9962 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -15,6 +15,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/). - [修复] 单股推送模式在未配置通知渠道时仍会落盘本地个股报告；CLI 启动分析若因空股票列表、个股结果全失败或本地报告保存失败而未生成报告，会显式返回失败并记录原因。 - [修复] 合并推送模式下即使个股汇总报告落盘失败，仍会先发送已有的合并通知；仅启用大盘复盘但最终未生成任何复盘内容时，分析任务会显式返回失败。 +- [新功能] Agent 工具调用支持按类别（data/search/analysis/action/market）配置默认超时，并允许单工具声明 `timeout_seconds`；有效超时按 first-wins 优先级解析（显式 per-run `tool_call_timeout_seconds` > 单工具显式 `timeout_seconds` > 类别默认 > 无限制），剩余 wall-clock 预算仅作不可突破的外层 cap，超时后返回结构化 `{"timeout": true}` 错误（标记 `retriable: false` 并写入 `non_retriable_tool_results` 防重试重复执行）供 Agent 继续执行而非中断循环（fixes #1890）。 +- [修复] Agent 工具注册表（`src/agent/factory.get_tool_registry`）由模块级缓存改为按「类别超时映射的值」比对失效，规避 CPython 回收对象后地址复用（`id(config)` 相同）导致配置 reload 后的 `Config` 被误判为未变、沿用过期超时的真 bug；新增 `_coerce_config_timeout` 类型白名单，使调用方传入 `MagicMock` / 缺属性 stub / 脏字符串（如 `float(MagicMock())` 静默得到 1.0）时降级为「无类别限制」而非崩溃或强加 1 秒超时；`build_agent_executor(config)` / `build_agent_chat_executor(config)` 现已把调用方 `config` 透传给 `get_tool_registry(config)`（不再无参调用冻结首构 registry）；`main._reload_runtime_config` 与 `SystemConfigService._reload_runtime_singletons`（及 `update()`→`reload_now` 路径）在配置热重载时调用 `reset_tool_registry()` 强制重建；回归测试补充「传入新 config 后 registry 重建」「reload 后新超时应生效」及「builder 透传 config」三类场景（#1890 的 review follow-up，闭环 OR-COM-dd1e8fa7 / OR-COM-bff42110） +- [修复] Agent 工具超时 review 闭环（fixes #1890 的 4 个 blocker）：超时解析由 min 契约改为 first-wins（显式 per-run `tool_call_timeout_seconds` > 单工具 `ToolDefinition.timeout_seconds` > 类别默认 > 无限制，剩余 wall-clock 预算只作不可突破的外层 cap；research 路径不再传 `tool_call_timeout_seconds` 以免覆盖类别限制）；超时结果标记 `retriable: false` 并写入 `non_retriable_tool_results` 阻断 LLM 同调用重试重入，且超时触发时为仍在后台运行的 handler 武装协作取消信号（`is_tool_cancellation_requested()` 与既有 `check_tool_execution()` 检查点均响应，handler 从不轮询则行为不变），作为 review 要求的「handler 内协作取消」缓解，规避 Python 线程无法 force-stop 导致的重复执行与副作用；`_coerce_config_timeout` 对 `inf`/`nan`/负数降级为「无限制」，根绝 `future.result(timeout=inf)` 触发 `OverflowError`；`get_tool_registry` / `reset_tool_registry` 加 `threading.Lock` 双检锁，且重建后返回本次构建的局部 registry（而非全局缓存），消除并发重建竞态与跨调用超时串扰；`@tool` 装饰器将 `ToolPolicy.timeout_seconds` 折叠进 `ToolDefinition` 单一来源；统一单/并行工具超时包装（单一 executor + deadline 驱动的 wait loop，消除并行路径嵌套 executor 与线程翻倍，duration 精确到各工具自身超时值），并新增快慢工具混合并行回归；同步 `docs/full-guide_EN.md` 的超时环境变量文档；测试覆盖 first-wins、non-retriable、协作取消接线、finite 校验、缓存线程安全与快慢混合并行。 +- [修复] 按最新 review 复核收敛 3 处正确性问题（OR-COM-7f3d3f5b / 3d6b61f8 / a1e8b0c2）：`BaseAgent._filtered_registry()` 携带源 registry 的类别超时映射（工具子集仍生效类别上限，不再绕过 #1890 类别超时）；并行批次 >5 时排队调用的 per-tool 超时自 worker 实际开始起算（不再提交即烧预算导致对未启动调用的假超时）；`get_tool_registry()` 缓存命中快路径在锁内读取一致对（消除与 `reset_tool_registry()` 竞态返回 `None` 或错配 registry）。新增对应回归测试。 + ## [3.30.0] - 2026-08-09 ### 发布亮点 diff --git a/docs/full-guide.md b/docs/full-guide.md index 3f2839c0a..e0f879f52 100644 --- a/docs/full-guide.md +++ b/docs/full-guide.md @@ -246,6 +246,10 @@ daily_stock_analysis/ | `AGENT_BACKEND` | 现有问股 Chat 的运行方式：`auto`（推荐，保持默认模型）、`litellm` 或 `codex_app_server`（实验，仅 single-agent Chat） | `auto` | 否 | | `AGENT_GENERATION_BACKEND` | Agent Chat 生成后端；Web 设置页仅暴露 `auto|litellm`，手写 local CLI backend 会返回 unsupported tool-calling 诊断 | `auto` | 否 | | `AGENT_SKILL_CONCURRENCY` | `specialist` 模式策略专家 worker 并发上限，范围 `1-4`；最多选择 4 个策略，默认 3 个并发，第 4 个进入下一批次并共享整体超时预算 | `3` | 否 | +| `AGENT_DATA_TOOL_TIMEOUT_S` | Agent `data` 类工具默认超时秒数；同时作为 `market` 类工具（`get_market_indices` / `get_sector_rankings` 等网络数据调用）的类别默认；`0` 表示关闭，回退到全局预算；有效超时按 first-wins 解析：显式 per-run `tool_call_timeout_seconds` > 单工具显式 `timeout_seconds` > 类别默认 > 无限制，剩余 wall-clock 预算仅作不可突破的外层 cap，`inf`/`nan`/负数降级为「无限制」；超时为 best-effort 软中断：Python 线程无法被强制停止，handler 可能在超时后继续运行，超时结果标记 `retriable: false` 并写入 `non_retriable_tool_results` 阻断重试，且为仍在后台的 handler 武装协作取消信号（`is_tool_cancellation_requested()` 与既有 `check_tool_execution()` 检查点均响应，handler 从不轮询则行为不变）以尽量减小副作用 | `0` | 否 | +| `AGENT_SEARCH_TOOL_TIMEOUT_S` | Agent `search` 类工具默认超时秒数；`0` 表示关闭，回退到全局预算 | `0` | 否 | +| `AGENT_ANALYSIS_TOOL_TIMEOUT_S` | Agent `analysis` 类工具默认超时秒数；`0` 表示关闭，回退到全局预算 | `0` | 否 | +| `AGENT_ACTION_TOOL_TIMEOUT_S` | Agent `action` 类工具默认超时秒数；`0` 表示关闭，回退到全局预算 | `0` | 否 | | `LITELLM_MODEL` | 主模型，格式 `provider/model`（如 `gemini/gemini-3.1-pro-preview`），推荐优先使用 | - | 否 | | `AGENT_LITELLM_MODEL` | 「默认模型」问股的主模型（可选）；留空继承主模型，无 provider 前缀按 `openai/` 解析；Codex 不使用此项 | - | 否 | | `AGENT_CONTEXT_COMPRESSION_ENABLED` | 「默认模型」问股可见历史的 LLM 压缩开关；Codex 使用最近 20 条可见对话且保留该配置 | `false` | 否 | diff --git a/docs/full-guide_EN.md b/docs/full-guide_EN.md index 159957d97..6f28c6f45 100644 --- a/docs/full-guide_EN.md +++ b/docs/full-guide_EN.md @@ -217,6 +217,10 @@ Default schedule: Every weekday at **18:00 (Beijing Time)** automatic execution. | `AGENT_BACKEND` | Runtime for the existing ask-stock Chat: `auto` (recommended, preserves the default model), `litellm`, or `codex_app_server` (experimental, single-agent Chat only) | `auto` | No | | `AGENT_GENERATION_BACKEND` | Agent Chat generation backend. Web settings only expose `auto|litellm`; hand-written local CLI backends return an unsupported tool-calling diagnostic | `auto` | No | | `AGENT_SKILL_CONCURRENCY` | Specialist-mode strategy worker concurrency cap, range `1-4`. Up to four strategies are selected; the default runs three concurrently and queues the fourth under the shared pipeline budget | `3` | No | +| `AGENT_DATA_TOOL_TIMEOUT_S` | Default timeout (seconds) for Agent `data`-category tools; also the category default for `market` tools (`get_market_indices` / `get_sector_rankings` and other network-backed data calls); `0` disables and falls back to the global budget. The effective timeout is resolved first-wins: explicit per-run `tool_call_timeout_seconds` > per-tool `timeout_seconds` > category default > no limit, with the remaining wall-clock budget as an unbreakable outer cap; `inf`/`nan`/negative degrade to "no limit". Timeout is a best-effort soft interrupt: Python threads cannot be force-stopped, so a handler may keep running after the timeout; the timed-out result is marked `retriable: false` and recorded in `non_retriable_tool_results` to block immediate retries, and the runner arms a cooperative-cancel signal (`is_tool_cancellation_requested()` and the existing `check_tool_execution()` checkpoints both honor it) to reduce side effects | `0` | No | +| `AGENT_SEARCH_TOOL_TIMEOUT_S` | Default timeout (seconds) for Agent `search`-category tools; `0` disables and falls back to the global budget | `0` | No | +| `AGENT_ANALYSIS_TOOL_TIMEOUT_S` | Default timeout (seconds) for Agent `analysis`-category tools; `0` disables and falls back to the global budget | `0` | No | +| `AGENT_ACTION_TOOL_TIMEOUT_S` | Default timeout (seconds) for Agent `action`-category tools; `0` disables and falls back to the global budget | `0` | No | | `LITELLM_MODEL` | Primary model, format `provider/model` (e.g. `gemini/gemini-3.1-pro-preview`), recommended | - | No | | `AGENT_LITELLM_MODEL` | Optional primary model for **Default model** ask-stock; empty inherits the primary model and bare names become `openai/`; Codex does not use this setting | - | No | | `AGENT_CONTEXT_COMPRESSION_ENABLED` | LLM compression for visible **Default model** ask-stock history; Codex uses the 20 most recent visible messages and retains this setting | `false` | No | diff --git a/main.py b/main.py index 6945406f1..ba5e43de4 100644 --- a/main.py +++ b/main.py @@ -1326,7 +1326,21 @@ def _reload_runtime_config() -> Config: """Reload config from the latest persisted `.env` values for scheduled runs.""" _reload_env_file_values_preserving_overrides() Config.reset_instance() - return get_config() + new_config = get_config() + + # Drop the module-level ToolRegistry so the next + # ``build_agent_executor`` / ``build_agent_chat_executor`` call rebuilds + # it against the freshly-loaded Config and picks up new + # ``AGENT_*_TOOL_TIMEOUT_S`` overrides (Issue #1890). Wrap in try/except + # so a future runtime reset helper cannot crash scheduled-job bootstrap. + try: + from src.agent.factory import reset_tool_registry + + reset_tool_registry() + except Exception as exc: # pragma: no cover - defensive guard + logger.warning("Failed to reset tool registry during config reload: %s", exc) + + return new_config def _build_schedule_time_provider(default_schedule_time: str): diff --git a/src/agent/agents/base_agent.py b/src/agent/agents/base_agent.py index 8fada59e1..588826918 100644 --- a/src/agent/agents/base_agent.py +++ b/src/agent/agents/base_agent.py @@ -257,7 +257,9 @@ class BaseAgent(ABC): return self.tool_registry from src.agent.tools.registry import ToolRegistry as TR - filtered = TR() + # Carry the source registry's category-timeout map so the filtered subset + # still enforces the per-category ceilings (review OR-COM-7f3d3f5b). + filtered = TR(category_timeout_map=self.tool_registry.category_timeout_map) for name in self.tool_names: tool_def = self.tool_registry.get(name) if tool_def: diff --git a/src/agent/factory.py b/src/agent/factory.py index 641a6edd9..268df5844 100644 --- a/src/agent/factory.py +++ b/src/agent/factory.py @@ -26,6 +26,8 @@ Usage:: import copy import logging +import math +import threading from dataclasses import dataclass from typing import List, Optional @@ -37,6 +39,17 @@ logger = logging.getLogger(__name__) # Module-level caches # --------------------------------------------------------------------------- _TOOL_REGISTRY = None +# Track the per-category timeout mapping ``_TOOL_REGISTRY`` was built from, so +# callers passing a freshly-reloaded ``Config`` always observe the right +# ceilings instead of inheriting the first build's values. +# +# Deliberately keyed by *value* rather than ``id(config)``: CPython reuses the +# address of a freed object, so a ``Config.reset_instance()`` + ``get_config()`` +# round-trip can hand back a new instance with the *same* id() as the collected +# one, which would silently mark a changed config as "unchanged". This mirrors +# the ``_SKILL_MANAGER_CUSTOM_DIR`` value-comparison used by ``get_skill_manager`` +# below. Pair with :func:`reset_tool_registry` for runtime config reload paths. +_CACHED_TIMEOUT_MAP: Optional[dict] = None _SKILL_MANAGER_PROTOTYPE = None # Sentinel used as initial value so None (i.e. no custom dir) compares as "changed" # on the very first call, forcing a build rather than accidentally skipping it. @@ -45,6 +58,20 @@ _SENTINEL = object() # the cache if AGENT_SKILL_DIR changes at runtime (e.g. via config reload). _SKILL_MANAGER_CUSTOM_DIR: object = _SENTINEL +# Guards the module-level ToolRegistry cache against concurrent rebuild/reset. +# ``get_tool_registry`` builds the registry *outside* this lock (see below) and +# only takes it to atomically assign the shared ``_TOOL_REGISTRY`` / +# ``_CACHED_TIMEOUT_MAP`` pair, so two requests with different ``Config`` +# objects can never observe a registry built for the other's timeouts. +_tool_registry_lock = threading.Lock() + +# Safety ceiling for any single per-category timeout value sourced from +# ``AGENT_*_TOOL_TIMEOUT_S``. ``inf``/``nan``/negative already degrade to +# "no limit"; this clips absurd-but-finite values so they can never reach +# ``future.result(timeout=...)`` and stall a request. The global +# ``tool_call_timeout_seconds`` budget is *not* subject to this clip. +_MAX_TOOL_TIMEOUT_S = 3600.0 + @dataclass class SkillPromptState: @@ -190,12 +217,144 @@ def _should_use_legacy_default_prompt( return getattr(bull_trend_skill, "source", None) == "builtin" -def get_tool_registry(): - """Return a cached ToolRegistry (built once, shared across requests).""" - global _TOOL_REGISTRY - if _TOOL_REGISTRY is not None: - return _TOOL_REGISTRY +def _coerce_config_timeout(config, field_name: str) -> float: + """Read a timeout field off ``config`` as a float, defaulting to ``0.0``. + ``get_tool_registry`` builds its category map from the *caller-supplied* + config, so it must tolerate partial config objects and test doubles: a + missing attribute raises ``AttributeError`` and ``MagicMock() > 0`` raises + ``TypeError``. A plain ``float()`` conversion is not enough either -- + ``float(MagicMock())`` silently yields ``1.0``, which would quietly impose + a bogus 1-second ceiling on every tool. Only genuine numbers (or numeric + strings, since ``.env`` values arrive as text) are accepted; anything else + degrades to "no category limit" and defers to the global + ``tool_call_timeout_seconds`` budget. + """ + raw_value = getattr(config, field_name, None) + + if isinstance(raw_value, bool) or raw_value is None: + # ``True`` would coerce to a nonsensical 1-second ceiling. + return 0.0 + if isinstance(raw_value, (int, float)): + value = float(raw_value) + if not math.isfinite(value): + # ``inf``/``nan`` are not valid ceilings; an ``inf`` would later + # raise ``OverflowError`` at ``future.result(timeout=inf)``. + logger.warning( + "[AgentFactory] Non-finite value for %s: %r, treating as no category limit", + field_name, + raw_value, + ) + return 0.0 + if value < 0: + logger.warning( + "[AgentFactory] Negative value for %s: %r, treating as no category limit", + field_name, + raw_value, + ) + return 0.0 + if value == 0.0: + # Explicitly disabled (the default) — silent, not a warning. + return 0.0 + if value > _MAX_TOOL_TIMEOUT_S: + logger.warning( + "[AgentFactory] Clamping %s=%r to safety ceiling %ss", + field_name, value, _MAX_TOOL_TIMEOUT_S, + ) + return float(_MAX_TOOL_TIMEOUT_S) + return value + if isinstance(raw_value, str): + try: + value = float(raw_value.strip()) + except ValueError: + logger.warning( + "[AgentFactory] Invalid value for %s: %r, treating as no category limit", + field_name, + raw_value, + ) + return 0.0 + if not math.isfinite(value): + logger.warning( + "[AgentFactory] Non-finite value for %s: %r, treating as no category limit", + field_name, + raw_value, + ) + return 0.0 + if value < 0: + logger.warning( + "[AgentFactory] Negative value for %s: %r, treating as no category limit", + field_name, + raw_value, + ) + return 0.0 + if value == 0.0: + return 0.0 + if value > _MAX_TOOL_TIMEOUT_S: + return float(_MAX_TOOL_TIMEOUT_S) + return value + + # Mocks / arbitrary objects: never let them shape real timeout behaviour. + logger.debug( + "[AgentFactory] Non-numeric value for %s: %r, treating as no category limit", + field_name, + type(raw_value).__name__, + ) + return 0.0 + + +def _build_category_timeout_map(config) -> dict: + """Map tool categories to their default timeout (seconds) from config. + + Only explicit overrides (value > 0) are carried; 0.0 means "no category + limit, fall back to the global tool_call_timeout_seconds budget". + + ``market`` tools (get_market_indices / get_sector_rankings) are + network-backed data calls, so they share the ``data`` ceiling. + """ + data_timeout = _coerce_config_timeout(config, "agent_data_tool_timeout_s") + return { + k: v + for k, v in { + "data": data_timeout, + "search": _coerce_config_timeout(config, "agent_search_tool_timeout_s"), + "analysis": _coerce_config_timeout(config, "agent_analysis_tool_timeout_s"), + "action": _coerce_config_timeout(config, "agent_action_tool_timeout_s"), + "market": data_timeout, + }.items() + if v > 0 + } + + +def reset_tool_registry() -> None: + """Drop the cached :class:`ToolRegistry` so the next access rebuilds it. + + Intended for runtime config reload paths + (``main._reload_runtime_config`` / + ``SystemConfigService._reload_runtime_singletons``) so that newly-loaded + ``AGENT_*_TOOL_TIMEOUT_S`` overrides actually take effect on subsequent + :func:`build_agent_executor` / :func:`build_agent_chat_executor` calls, + instead of being silently shadowed by the registry built at first import. + + Guarded by ``_tool_registry_lock`` so a concurrent :func:`get_tool_registry` + cannot read a half-reset cache (``None`` registry with a stale map, or vice + versa). + """ + global _TOOL_REGISTRY, _CACHED_TIMEOUT_MAP + with _tool_registry_lock: + if _TOOL_REGISTRY is not None: + logger.info("[AgentFactory] ToolRegistry cache cleared; will rebuild on next access") + _TOOL_REGISTRY = None + _CACHED_TIMEOUT_MAP = None + + +def _build_tool_registry(category_timeout_map): + """Construct and populate a :class:`ToolRegistry` from the tool modules. + + Extracted from :func:`get_tool_registry` so the (potentially heavy) tool + module imports and registrations happen *outside* the cache lock — holding + the lock during imports would serialise startup and risk a deadlock if any + imported module ever imported this module back. + """ from src.agent.tools.registry import ToolRegistry from src.agent.tools.data_tools import ALL_DATA_TOOLS from src.agent.tools.analysis_tools import ALL_ANALYSIS_TOOLS @@ -203,13 +362,91 @@ def get_tool_registry(): from src.agent.tools.market_tools import ALL_MARKET_TOOLS from src.agent.tools.backtest_tools import ALL_BACKTEST_TOOLS - registry = ToolRegistry() - for tool_fn in ALL_DATA_TOOLS + ALL_ANALYSIS_TOOLS + ALL_SEARCH_TOOLS + ALL_MARKET_TOOLS + ALL_BACKTEST_TOOLS: + registry = ToolRegistry(category_timeout_map=category_timeout_map or None) + for tool_fn in ( + ALL_DATA_TOOLS + + ALL_ANALYSIS_TOOLS + + ALL_SEARCH_TOOLS + + ALL_MARKET_TOOLS + + ALL_BACKTEST_TOOLS + ): registry.register(tool_fn) + return registry - _TOOL_REGISTRY = registry - logger.info("[AgentFactory] ToolRegistry cached (%d tools)", len(registry._tools) if hasattr(registry, "_tools") else -1) - return _TOOL_REGISTRY + +def get_tool_registry(config=None): + """Return a :class:`ToolRegistry` bound to the lifetime of ``config``. + + The first call (or the first call after :func:`reset_tool_registry`) + builds and caches a registry from ``config``. The cache is keyed by the + resolved per-category timeout *values*, so a config carrying different + ``AGENT_*_TOOL_TIMEOUT_S`` settings (e.g. a freshly-reloaded ``Config`` + singleton) rebuilds the registry, while an equivalent config reuses the + cached instance and avoids re-registering the whole tool surface. + + Value comparison is intentional: ``id(config)`` is not a safe key because + CPython reuses the address of a collected object, so the instance returned + after ``Config.reset_instance()`` can share the previous instance's id and + would be misread as "unchanged". + + When ``config`` is ``None`` the function falls back to ``get_config()``, + preserving the legacy contract for callers that do not yet know about + config reload. + + Thread-safety: the fast-path cache hit reads both globals under the lock + (so the check and the return observe one consistent pair — see + ``reset_tool_registry``), and the registry build runs outside the lock with + only the final assignment of the shared ``_TOOL_REGISTRY`` / + ``_CACHED_TIMEOUT_MAP`` pair serialised (double-checked locking), so two + requests with different ``Config`` objects can never observe a registry built + for the other's timeouts. On a (re)build the function returns the registry it + just constructed for the caller's timeout map, so a concurrent rebuild cannot + leak a mismatched registry into this call's result. + """ + global _TOOL_REGISTRY, _CACHED_TIMEOUT_MAP + + if config is None: + from src.config import get_config + + config = get_config() + + category_timeout_map = _build_category_timeout_map(config) + # Fast path: take the lock so the cache-hit check and the return observe one + # consistent (registry, timeout-map) pair — an unlocked check-then-use can + # race with ``reset_tool_registry()`` and return ``None`` or a registry built + # for the wrong timeout map (review OR-COM-a1e8b0c2). + with _tool_registry_lock: + if _TOOL_REGISTRY is not None and _CACHED_TIMEOUT_MAP == category_timeout_map: + return _TOOL_REGISTRY + + # Build *outside* the lock (see ``_build_tool_registry`` rationale above). + registry = _build_tool_registry(category_timeout_map) + with _tool_registry_lock: + # Re-check under the lock: another thread may have installed an + # equivalent registry while we were importing the tool modules. + if _TOOL_REGISTRY is not None and _CACHED_TIMEOUT_MAP == category_timeout_map: + return _TOOL_REGISTRY + invalidated = _TOOL_REGISTRY is not None + _TOOL_REGISTRY = registry + _CACHED_TIMEOUT_MAP = category_timeout_map + if invalidated: + logger.info( + "[AgentFactory] ToolRegistry rebuilt for new category timeouts: %s", + category_timeout_map or "none (global budget only)", + ) + else: + logger.info( + "[AgentFactory] ToolRegistry cached (%d tools, category timeouts: %s)", + len(registry._tools) if hasattr(registry, "_tools") else -1, + category_timeout_map or "none (global budget only)", + ) + # Return the registry built for *this* call's timeout map (not the shared + # ``_TOOL_REGISTRY`` global). Under a concurrent rebuild for a different + # ``Config``/timeout map, this guarantees the caller gets the registry that + # matches the timeouts it requested, instead of a registry built for another + # request. The global is still assigned above so equivalent subsequent calls + # hit the cache fast path. + return registry def get_skill_manager(config=None): @@ -336,7 +573,11 @@ def build_agent_executor(config=None, skills: Optional[List[str]] = None): from src.agent.llm_adapter import LLMToolAdapter - registry = get_tool_registry() + # Pass ``config`` explicitly so a reloaded ``Config`` instance observes + # the new per-category timeout map (Issue #1890 compatibility concern: + # the module-level ToolRegistry would otherwise keep the first build's + # timeouts even after ``_reload_runtime_config``). + registry = get_tool_registry(config) prompt_state = resolve_skill_prompt_state(config, skills=skills) skill_manager = prompt_state.skill_manager logger.info( @@ -406,7 +647,10 @@ def build_agent_chat_executor(config=None, skills: Optional[List[str]] = None): if backend_id == "litellm" and arch == "multi": return build_agent_executor(config, skills=skills) - registry = get_tool_registry() + # Pass ``config`` explicitly so a reloaded ``Config`` instance observes + # the new per-category timeout map (see ``build_agent_executor`` for the + # same rationale). + registry = get_tool_registry(config) prompt_state = resolve_skill_prompt_state(config, skills=skills) if backend_id == "litellm": from src.agent.llm_adapter import LLMToolAdapter diff --git a/src/agent/research.py b/src/agent/research.py index a0a5e5fb7..1a859fa58 100644 --- a/src/agent/research.py +++ b/src/agent/research.py @@ -379,7 +379,10 @@ Token budget remaining: ~{remaining_budget} llm_adapter=self.llm_adapter, max_steps=4, max_wall_clock_seconds=timeout_seconds, - tool_call_timeout_seconds=timeout_seconds, + # NOTE: do not pass tool_call_timeout_seconds here — under the + # first-wins contract it would override the per-tool / category + # limits, letting one tool consume the whole sub-question budget. + # The remaining wall-clock budget alone keeps tools bounded. ) if not result.success and self._looks_like_timeout_error(result.error): return { diff --git a/src/agent/runner.py b/src/agent/runner.py index 529271cf6..c57a82ae6 100644 --- a/src/agent/runner.py +++ b/src/agent/runner.py @@ -17,10 +17,12 @@ from __future__ import annotations import json import logging +import math import re import time +import threading import contextvars -from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError, as_completed +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Optional @@ -30,6 +32,7 @@ from src.agent.protocols import StageFailureReason from src.agent.stream_events import stream_event from src.agent.tools.registry import ToolRegistry from src.agent.tools.execution import ( + TOOL_CANCEL_EVENT, _build_tool_cache_key, _guard_tool_stock_scope, _is_non_retriable_tool_result, @@ -348,7 +351,10 @@ def run_agent_loop( progress_callback: Optional callback receiving progress dicts. thinking_labels: Override map of tool_name → friendly label. max_wall_clock_seconds: Optional overall timeout budget for the loop. - tool_call_timeout_seconds: Optional timeout for one parallel tool batch. + tool_call_timeout_seconds: Optional explicit per-run tool-call timeout. + Highest-priority (first-wins) in the per-tool timeout chain — + overrides per-tool declarations and category defaults — but is + capped by ``max_wall_clock_seconds`` when both are set. emit_stage_events: Whether to emit the synthetic ``agent_loop`` stage lifecycle. Orchestrated business stages disable this so ``stage_start`` / ``stage_done`` only describe real stages. @@ -509,13 +515,10 @@ def run_agent_loop( assistant_msg["provider_blocks"] = response.provider_blocks messages.append(assistant_msg) - # Execute tools (parallel when > 1) - effective_tool_timeout = tool_call_timeout_seconds - if remaining_timeout is not None: - effective_tool_timeout = min( - remaining_timeout, - tool_call_timeout_seconds if tool_call_timeout_seconds and tool_call_timeout_seconds > 0 else remaining_timeout, - ) + # Execute tools (parallel when > 1). ``tool_call_timeout_seconds`` is + # the caller's explicit per-run override — highest priority in the + # first-wins chain (Issue #1890 contract) — while ``remaining_timeout`` + # is the unbreakable outer wall-clock cap for this batch. tool_results = _execute_tools( response.tool_calls, tool_registry, @@ -523,7 +526,8 @@ def run_agent_loop( progress_callback, tool_calls_log, non_retriable_tool_results, - tool_wait_timeout_seconds=effective_tool_timeout, + tool_call_timeout_seconds=tool_call_timeout_seconds, + tool_wait_timeout_seconds=remaining_timeout, stock_scope=stock_scope, ) @@ -601,6 +605,107 @@ def run_agent_loop( # Internal tool execution # ============================================================ +def _coerce_positive_timeout(value) -> Optional[float]: + """Coerce a timeout candidate to a positive finite float, else ``None``. + + ``None`` / non-numeric / non-positive / ``inf`` / ``nan`` all map to + ``None`` ("no limit at this level"). Rejecting non-finite values prevents + an ``OverflowError`` from ``future.result(timeout=inf)`` and avoids the + undefined ordering a ``nan`` would produce. + """ + if value is None: + return None + try: + v = float(value) + except (TypeError, ValueError): + return None + if not math.isfinite(v) or v <= 0: + return None + return v + + +def _build_timeout_result_payload( + tool_name: str, + arguments: Dict[str, Any], + timeout_s: float, + non_retriable_tool_results: Optional[Dict[str, str]], +) -> str: + """Build the timeout-shaped result string and record it as non-retriable. + + A timed-out call is marked ``"retriable": False`` and inserted into + ``non_retriable_tool_results`` so the LLM's retry of the *same* call reuses + the cached failure instead of spinning up a second (possibly side-effecting) + execution — a best-effort guard against duplicate work, since Python cannot + forcibly cancel an already-started tool thread. + """ + label = f"{timeout_s:.2f}s" + result_str = json.dumps({ + "error": f"Tool execution timed out after {label}", + "timeout": True, + "retriable": False, + }) + if non_retriable_tool_results is not None: + cache_key = _build_tool_cache_key(tool_name, arguments) + # Non-dict / missing ``arguments`` yields ``None`` here; skip the write so + # unrelated no-arg tool calls cannot collide on a shared ``None`` key. + if cache_key: + non_retriable_tool_results[cache_key] = result_str + return result_str + + +def _resolve_per_tool_timeout( + tool_call, + tool_registry: Optional[ToolRegistry], + explicit_timeout: Optional[float] = None, + wall_clock_budget: Optional[float] = None, +) -> Optional[float]: + """Resolve the effective timeout for a single tool call (Issue #1890). + + Precedence is *first-wins*, exactly as confirmed in the issue contract: + + 1. explicit per-run timeout (``tool_call_timeout_seconds``) + 2. per-tool declaration (``ToolDefinition.timeout_seconds``) + 3. category default (``AGENT_*_TOOL_TIMEOUT_S``) + 4. none (no per-tool limit) + + An earlier level is never lowered by a later one: an explicit + ``tool_call_timeout_seconds`` can *relax* a stricter per-tool or category + default (the previous ``min()`` across all levels made that impossible), and + a per-tool declaration is never overridden by a smaller category default. + + ``wall_clock_budget`` (the remaining overall loop budget) is then applied as + an *unbreakable outer cap*: the winner can never exceed it, so a long + explicit timeout cannot blow past the caller's overall wall-clock budget. + """ + explicit = _coerce_positive_timeout(explicit_timeout) + budget = _coerce_positive_timeout(wall_clock_budget) + + if tool_registry is None: + base = explicit + else: + tool_def = tool_registry.get(tool_call.name) + if tool_def is None: + base = explicit + else: + per_tool = _coerce_positive_timeout(getattr(tool_def, "timeout_seconds", None)) + category = _coerce_positive_timeout( + tool_registry.category_default_timeout(tool_def.category) + ) + # First-wins: explicit per-run > per-tool declaration > category default. + base = ( + explicit + if explicit is not None + else (per_tool if per_tool is not None else category) + ) + + if base is None: + # No explicit/per-tool/category limit; the wall-clock budget governs. + return budget + if budget is not None: + return min(base, budget) + return base + + def _execute_tools( tool_calls, tool_registry: ToolRegistry, @@ -608,12 +713,23 @@ def _execute_tools( progress_callback: Optional[Callable], tool_calls_log: List[Dict[str, Any]], non_retriable_tool_results: Optional[Dict[str, str]] = None, + tool_call_timeout_seconds: Optional[float] = None, tool_wait_timeout_seconds: Optional[float] = None, stock_scope: Optional[StockScope] = None, ) -> List[Dict[str, Any]]: """Execute one or more tool calls, returning ordered result dicts. - Single tools run inline; multiple tools run in parallel threads. + A single tool with no resolved timeout runs inline (no thread); otherwise + all tools share one executor whose per-tool timeouts are enforced by a + deadline-driven wait loop — there are no per-tool nested pools, so a batch + of N tools uses at most ``min(N, 5)`` threads (review: unify the single and + parallel timeout wrapping). ``tool_wait_timeout_seconds`` also caps the + whole batch as an outer ceiling. + + ``tool_call_timeout_seconds`` is the caller's explicit per-run override + (highest priority, first-wins) and ``tool_wait_timeout_seconds`` the + remaining wall-clock budget (outer cap) — both feed + :func:`_resolve_per_tool_timeout`. """ def _exec_single(tc_item): @@ -624,50 +740,39 @@ def _execute_tools( non_retriable_tool_results=non_retriable_tool_results, ) - results: List[Dict[str, Any]] = [] + def _exec_with_deadline(tc_item, per_tool_timeout, deadline_holder): + """Run one tool and record its per-tool deadline when the worker actually + starts. The batch pool is capped at ``min(N, 5)`` workers, so calls that + queue behind a full pool must not burn their timeout before execution + begins (review OR-COM-3d6b61f8) — otherwise a batch larger than five can + falsely time out a tool that never got a worker. + """ + if per_tool_timeout and per_tool_timeout > 0: + deadline_holder[0] = time.monotonic() + per_tool_timeout + return _exec_single(tc_item) - if len(tool_calls) == 1: - tc = tool_calls[0] - if progress_callback: - progress_callback(stream_event("tool_start", step=step, tool=tc.name)) - timeout_triggered = False - if tool_wait_timeout_seconds and tool_wait_timeout_seconds > 0: - pool = ThreadPoolExecutor(max_workers=1) - ctx = contextvars.copy_context() - try: - future = pool.submit(ctx.run, _exec_single, tc) - try: - _, result_str, success, dur, cached, guard_result = future.result(timeout=tool_wait_timeout_seconds) - except FuturesTimeoutError: - timeout_triggered = True - future.cancel() - timeout_label = f"{tool_wait_timeout_seconds:.2f}s" - logger.warning("Tool '%s' timed out after %s at step %d", tc.name, timeout_label, step) - result_str = json.dumps({ - "error": f"Tool execution timed out after {timeout_label}", - "timeout": True, - }) - success = False - dur = round(tool_wait_timeout_seconds, 2) - cached = False - guard_result = None - finally: - pool.shutdown(wait=not timeout_triggered, cancel_futures=timeout_triggered) + def _record(tc_item, *, timed_out, timeout_s=None, out=None, guard_result=None): + """Emit ``tool_done``, build the log entry, and append to ``results``.""" + if out is not None: + _, result_str, success, dur, cached, guard_result = out else: - _, result_str, success, dur, cached, guard_result = _exec_single(tc) + result_str = _build_timeout_result_payload( + tc_item.name, tc_item.arguments, timeout_s, non_retriable_tool_results, + ) + success = False + dur = round(timeout_s, 2) + cached = False if progress_callback: - progress_callback(stream_event("tool_done", step=step, tool=tc.name, success=success, duration=dur)) + progress_callback(stream_event( + "tool_done", step=step, tool=tc_item.name, success=success, duration=dur, + )) log_entry = { - "step": step, "tool": tc.name, "arguments": tc.arguments, + "step": step, "tool": tc_item.name, "arguments": tc_item.arguments, "success": success, "duration": dur, "result_length": len(result_str), "cached": cached, } - if tool_wait_timeout_seconds and tool_wait_timeout_seconds > 0 and not success: - try: - if json.loads(result_str).get("timeout") is True: - log_entry["timeout"] = True - except (TypeError, ValueError, json.JSONDecodeError): - pass + if timed_out: + log_entry["timeout"] = True if guard_result is not None: log_entry.update({ "guarded": True, @@ -676,74 +781,109 @@ def _execute_tools( "allowed_stock_codes": guard_result.get("allowed_stock_codes", []), }) tool_calls_log.append(log_entry) - results.append({"tc": tc, "result_str": result_str}) - else: - for tc in tool_calls: - if progress_callback: - progress_callback(stream_event("tool_start", step=step, tool=tc.name)) + results.append({"tc": tc_item, "result_str": result_str}) - pool = ThreadPoolExecutor(max_workers=min(len(tool_calls), 5)) - timeout_triggered = False - try: - futures = {pool.submit(contextvars.copy_context().run, _exec_single, tc): tc for tc in tool_calls} - pending = set(futures) - for future in as_completed( - futures, - timeout=tool_wait_timeout_seconds if tool_wait_timeout_seconds and tool_wait_timeout_seconds > 0 else None, - ): - pending.discard(future) - tc_item, result_str, success, dur, cached, guard_result = future.result() - if progress_callback: - progress_callback(stream_event("tool_done", step=step, tool=tc_item.name, success=success, duration=dur)) - log_entry = { - "step": step, "tool": tc_item.name, "arguments": tc_item.arguments, - "success": success, "duration": dur, "result_length": len(result_str), - "cached": cached, - } - if guard_result is not None: - log_entry.update({ - "guarded": True, - "expected_stock_code": guard_result.get("expected_stock_code"), - "requested_stock_code": guard_result.get("requested_stock_code"), - "allowed_stock_codes": guard_result.get("allowed_stock_codes", []), - }) - tool_calls_log.append(log_entry) - results.append({"tc": tc_item, "result_str": result_str}) - except FuturesTimeoutError: - timeout_triggered = True - timeout_label = ( - f"{tool_wait_timeout_seconds:.2f}s" - if tool_wait_timeout_seconds is not None - else "the configured limit" + results: List[Dict[str, Any]] = [] + if not tool_calls: + return results + + plan = [] + for tc in tool_calls: + per_tool_timeout = _resolve_per_tool_timeout( + tc, tool_registry, tool_call_timeout_seconds, tool_wait_timeout_seconds, + ) + plan.append((tc, per_tool_timeout)) + if progress_callback: + progress_callback(stream_event("tool_start", step=step, tool=tc.name)) + + # Fast path: a single tool with no resolved timeout runs inline (no thread). + if len(plan) == 1 and not (plan[0][1] and plan[0][1] > 0): + _record(plan[0][0], timed_out=False, out=_exec_single(plan[0][0])) + return results + + pool = ThreadPoolExecutor(max_workers=min(len(plan), 5)) + futures: Dict[Any, Any] = {} + cancel_of: Dict[Any, threading.Event] = {} + # ``deadline_of`` maps a Future to a single-element holder list; the worker + # writes the deadline into it when it *starts* running (see + # ``_exec_with_deadline``). ``None`` until then means "still queued / no + # per-tool timeout" and is never treated as expired. + deadline_of: Dict[Any, List[Optional[float]]] = {} + timeout_of: Dict[Any, float] = {} + timeout_triggered = False + try: + # One executor for the whole batch. Each future gets its own cancel + # event (keyed by Future, not by tool name, so two parallel calls to the + # same tool cannot shadow each other) and its own per-tool deadline. + for tc, per_tool_timeout in plan: + cancel_event = threading.Event() + ctx = contextvars.copy_context() + ctx.run(TOOL_CANCEL_EVENT.set, cancel_event) + deadline_holder: List[Optional[float]] = [None] + fut = pool.submit( + ctx.run, _exec_with_deadline, tc, per_tool_timeout, deadline_holder, ) - logger.warning("Tool batch timed out after %s at step %d", timeout_label, step) - for future, tc_item in futures.items(): - if future in pending: - future.cancel() - result_str = json.dumps({ - "error": f"Tool execution timed out after {timeout_label}", - "timeout": True, - }) - if progress_callback: - progress_callback(stream_event( - "tool_done", - step=step, - tool=tc_item.name, - success=False, - duration=round(tool_wait_timeout_seconds or 0.0, 2), - )) - tool_calls_log.append({ - "step": step, - "tool": tc_item.name, - "arguments": tc_item.arguments, - "success": False, - "duration": round(tool_wait_timeout_seconds or 0.0, 2), - "result_length": len(result_str), - "cached": False, - "timeout": True, - }) - results.append({"tc": tc_item, "result_str": result_str}) - finally: - pool.shutdown(wait=not timeout_triggered, cancel_futures=timeout_triggered) + futures[fut] = tc + cancel_of[fut] = cancel_event + timeout_of[fut] = per_tool_timeout or 0.0 + deadline_of[fut] = deadline_holder + + batch_deadline = ( + time.monotonic() + tool_wait_timeout_seconds + if tool_wait_timeout_seconds and tool_wait_timeout_seconds > 0 + else None + ) + pending = set(futures) + + while pending: + now = time.monotonic() + deadlines = [] + for f in pending: + holder = deadline_of.get(f) + if holder is not None and holder[0] is not None: + deadlines.append(holder[0]) + if batch_deadline is not None: + deadlines.append(batch_deadline) + next_deadline = min(deadlines) if deadlines else None + if next_deadline is not None: + wait_timeout = max(0.0, next_deadline - now) + elif any(timeout_of.get(f, 0.0) > 0 for f in pending): + # A pending call has a per-tool timeout but its worker has not + # started yet (it is queued behind the capped pool). Poll briefly + # so a tool that starts and then hangs still gets its own timeout + # (review OR-COM-3d6b61f8) instead of being forgotten. + wait_timeout = 0.01 + else: + # No deadline and no pending call has a timeout — nothing can time + # out, so block until a call completes (efficient for the default + # no-timeout path). + wait_timeout = None + done, _ = wait(pending, timeout=wait_timeout, return_when=FIRST_COMPLETED) + for fut in done: + pending.discard(fut) + _record(futures[fut], timed_out=False, out=fut.result()) + + now = time.monotonic() + for fut in list(pending): + holder = deadline_of.get(fut) + deadline = holder[0] if holder is not None else None + expired = (deadline is not None and now >= deadline) or ( + batch_deadline is not None and now >= batch_deadline + ) + if not expired: + continue + pending.discard(fut) + cancel_of[fut].set() + timeout_triggered = True + timeout_s = timeout_of.get(fut) or (tool_wait_timeout_seconds or 0.0) + logger.warning( + "Tool '%s' timed out after %.2fs at step %d", + futures[fut].name, timeout_s, step, + ) + _record(futures[fut], timed_out=True, timeout_s=timeout_s) + finally: + # Do not wait when a timeout fired — the timed-out handler may still be + # running in a worker and would otherwise block this call indefinitely. + pool.shutdown(wait=not timeout_triggered, cancel_futures=True) return results diff --git a/src/agent/tools/execution.py b/src/agent/tools/execution.py index 51abecc87..9d7f7ffd2 100644 --- a/src/agent/tools/execution.py +++ b/src/agent/tools/execution.py @@ -21,6 +21,39 @@ from src.agent.tools.registry import ToolRegistry logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Cooperative cancellation signal (best-effort timeouts, Issue #1890 / review) +# --------------------------------------------------------------------------- +# Python cannot forcibly stop an already-started tool thread, so when a tool +# times out its handler may keep running in the background. To give long or +# side-effecting handlers a way to honour the timeout, the runner arms a +# per-call ``threading.Event`` on timeout and publishes it through this +# contextvar for the duration of the call. Handlers MAY poll +# ``is_tool_cancellation_requested`` and abort early; the same signal is also +# honoured by the existing ``check_tool_execution()`` checkpoint that data / +# backtest / tool-surface handlers already call, so a timed-out handler stops at +# its next safe boundary even without opting into the helper. The signal is +# strictly opt-in by checkpoint: handlers that never call either are completely +# unaffected, and the runner only sets the event after a timeout has already +# fired — never during normal completion. This is the "in-handler cooperative +# cancel" mitigation requested in review, layered on top of the non-retriable +# cache (which blocks the LLM from re-launching the same call). +TOOL_CANCEL_EVENT: "contextvars.ContextVar[Optional[threading.Event]]" = contextvars.ContextVar( + "tool_cancel_event", default=None +) + + +def is_tool_cancellation_requested() -> bool: + """Return True when the current tool call has been asked to cancel. + + Tool handlers may call this cheaply inside long loops or before performing + an irreversible side effect, to honour a best-effort timeout. Returns False + unless the runner has armed the cancellation event for the running call. + """ + event = TOOL_CANCEL_EVENT.get() + return event is not None and event.is_set() + + _SUMMARY_LIMIT = 500 _TOKEN_PATTERN = re.compile( r"(?i)\b(?:sk|pk|ghp|gho|github_pat|xox[baprs]?|bearer)[-_a-z0-9]{12,}\b" @@ -100,12 +133,20 @@ def reset_tool_execution_context(token: contextvars.Token) -> None: def check_tool_execution() -> None: """Stop at a safe handler boundary when cancellation or deadline is reached.""" context = _ACTIVE_TOOL_CONTEXT.get() - if context is None: - return - if context.cancel_event is not None and context.cancel_event.is_set(): - raise ToolExecutionCancelled("Tool execution was cancelled") - if context.deadline is not None and time.monotonic() >= context.deadline: - raise ToolExecutionDeadlineExceeded("Tool execution deadline was exceeded") + if context is not None: + if context.cancel_event is not None and context.cancel_event.is_set(): + raise ToolExecutionCancelled("Tool execution was cancelled") + if context.deadline is not None and time.monotonic() >= context.deadline: + raise ToolExecutionDeadlineExceeded("Tool execution deadline was exceeded") + # Runner-armed cooperative cancel (Issue #1890): when a per-tool timeout + # fires, the runner arms ``TOOL_CANCEL_EVENT`` for the running call. Real + # tool handlers (data/backtest/tool-surface) poll this checkpoint rather than + # the opt-in ``is_tool_cancellation_requested()`` helper, so honour the signal + # here too — a still-running handler then aborts before its next side-effecting + # step instead of running to completion in the background thread. + cancel_event = TOOL_CANCEL_EVENT.get() + if cancel_event is not None and cancel_event.is_set(): + raise ToolExecutionCancelled("Tool execution timed out (cooperative cancel)") def serialize_tool_result(result: Any) -> str: diff --git a/src/agent/tools/registry.py b/src/agent/tools/registry.py index 8b877bad9..fe33222b2 100644 --- a/src/agent/tools/registry.py +++ b/src/agent/tools/registry.py @@ -44,6 +44,7 @@ class ToolPolicy: policy_status: str = "unknown" scope_dimensions: List[str] = field(default_factory=list) cancellation_safe: bool = False + timeout_seconds: Optional[float] = None @classmethod def unknown(cls) -> "ToolPolicy": @@ -58,6 +59,7 @@ class ToolPolicy: permissions: Optional[List[str]] = None, scope_dimensions: Optional[List[str]] = None, cancellation_safe: bool = False, + timeout_seconds: Optional[float] = None, ) -> "ToolPolicy": return cls( read_only=read_only, @@ -66,6 +68,7 @@ class ToolPolicy: policy_status="declared", scope_dimensions=list(scope_dimensions or []), cancellation_safe=bool(cancellation_safe), + timeout_seconds=timeout_seconds, ) def to_public_dict(self) -> Dict[str, Any]: @@ -85,8 +88,9 @@ class ToolDefinition: description: str parameters: List[ToolParameter] handler: Callable - category: str = "data" # data | analysis | search | action + category: str = "data" # data | analysis | search | action | market policy: ToolPolicy = field(default_factory=ToolPolicy.unknown) + timeout_seconds: Optional[float] = None # Optional per-tool execution timeout # ----- Multi-provider schema converters ----- @@ -172,8 +176,9 @@ class ToolRegistry: registry.execute("get_realtime_quote", stock_code="600519") """ - def __init__(self): + def __init__(self, category_timeout_map: Optional[Dict[str, float]] = None): self._tools: Dict[str, ToolDefinition] = {} + self._category_timeouts: Dict[str, float] = dict(category_timeout_map or {}) # ----- Registration ----- @@ -215,6 +220,24 @@ class ToolRegistry: def __contains__(self, name: str) -> bool: return name in self._tools + def category_default_timeout(self, category: str) -> Optional[float]: + """Return the configured default timeout (seconds) for a tool category. + + Returns ``None`` when no default is configured, letting the caller fall + back to the global ``tool_call_timeout_seconds`` budget. + """ + return self._category_timeouts.get(category) + + @property + def category_timeout_map(self) -> Dict[str, float]: + """Return a copy of the per-category default timeout map. + + Used by :meth:`ToolRegistry` consumers that build a filtered registry + (e.g. ``BaseAgent._filtered_registry``) so the category ceilings survive + the subset copy (review OR-COM-7f3d3f5b). + """ + return dict(self._category_timeouts) + # ----- Schema generation ----- def to_openai_tools(self) -> List[dict]: @@ -315,6 +338,7 @@ def tool( parameters: Optional[List[ToolParameter]] = None, registry: Optional[ToolRegistry] = None, policy: Optional[ToolPolicy] = None, + timeout_seconds: Optional[float] = None, ): """Decorator to register a function as an agent tool. @@ -333,13 +357,29 @@ def tool( if params is None: params = _infer_parameters(func) + # Single source of truth for per-tool timeout is + # ``ToolDefinition.timeout_seconds`` (the field ``runner`` actually + # reads). If the caller supplied a ``timeout_seconds`` via the policy + # but omitted the explicit ``@tool(timeout_seconds=...)`` argument, fold + # the policy value in here so the two never diverge. + effective_timeout = timeout_seconds + resolved_policy = policy or ToolPolicy.unknown() + if effective_timeout is None and getattr(resolved_policy, "timeout_seconds", None) is not None: + effective_timeout = resolved_policy.timeout_seconds + logger.debug( + "Tool '%s': using ToolPolicy.timeout_seconds=%s as the effective per-tool timeout", + name, + effective_timeout, + ) + tool_def = ToolDefinition( name=name, description=description, parameters=params, handler=func, category=category, - policy=policy or ToolPolicy.unknown(), + policy=resolved_policy, + timeout_seconds=effective_timeout, ) target_registry = registry or get_default_registry() diff --git a/src/config.py b/src/config.py index 2521f3090..657c6d9f3 100644 --- a/src/config.py +++ b/src/config.py @@ -1014,6 +1014,13 @@ class Config: agent_decision_agent_timeout_s: float = 0 agent_portfolio_agent_timeout_s: float = 0 agent_skill_agent_timeout_s: float = 0 + # Per-category default timeouts for agent tool calls (seconds). + # 0 / unset means "no category default" -> falls back to the global + # tool_call_timeout_seconds budget. + agent_data_tool_timeout_s: float = 0.0 + agent_search_tool_timeout_s: float = 0.0 + agent_analysis_tool_timeout_s: float = 0.0 + agent_action_tool_timeout_s: float = 0.0 agent_skill_concurrency: int = 3 agent_risk_override: bool = True # Allow risk agent to veto buy signals agent_deep_research_budget: int = 30000 # Max token budget for deep research @@ -1946,6 +1953,22 @@ class Config: os.getenv('AGENT_SKILL_AGENT_TIMEOUT_S'), 0, field_name='AGENT_SKILL_AGENT_TIMEOUT_S', minimum=0, ), + agent_data_tool_timeout_s=parse_env_float( + os.getenv('AGENT_DATA_TOOL_TIMEOUT_S'), 0.0, + field_name='AGENT_DATA_TOOL_TIMEOUT_S', minimum=0.0, + ), + agent_search_tool_timeout_s=parse_env_float( + os.getenv('AGENT_SEARCH_TOOL_TIMEOUT_S'), 0.0, + field_name='AGENT_SEARCH_TOOL_TIMEOUT_S', minimum=0.0, + ), + agent_analysis_tool_timeout_s=parse_env_float( + os.getenv('AGENT_ANALYSIS_TOOL_TIMEOUT_S'), 0.0, + field_name='AGENT_ANALYSIS_TOOL_TIMEOUT_S', minimum=0.0, + ), + agent_action_tool_timeout_s=parse_env_float( + os.getenv('AGENT_ACTION_TOOL_TIMEOUT_S'), 0.0, + field_name='AGENT_ACTION_TOOL_TIMEOUT_S', minimum=0.0, + ), agent_skill_concurrency=parse_env_int( os.getenv('AGENT_SKILL_CONCURRENCY'), 3, diff --git a/src/services/system_config_service.py b/src/services/system_config_service.py index 43d454458..0d89c909d 100644 --- a/src/services/system_config_service.py +++ b/src/services/system_config_service.py @@ -313,9 +313,15 @@ class SystemConfigService: @staticmethod def _reload_runtime_singletons() -> None: """Reset runtime singleton services after config reload.""" + from src.agent.factory import reset_tool_registry from src.agent.tools.data_tools import reset_fetcher_manager from src.search_service import reset_search_service + # Drop the module-level ToolRegistry so the next + # ``build_agent_executor`` / ``build_agent_chat_executor`` call + # rebuilds ``_TOOL_REGISTRY`` against the fresh Config and picks up + # new ``AGENT_*_TOOL_TIMEOUT_S`` overrides (Issue #1890). + reset_tool_registry() reset_fetcher_manager() reset_search_service() diff --git a/tests/agent/test_tool_timeout.py b/tests/agent/test_tool_timeout.py new file mode 100644 index 000000000..9efc123a4 --- /dev/null +++ b/tests/agent/test_tool_timeout.py @@ -0,0 +1,1083 @@ +# -*- coding: utf-8 -*- +"""Unit tests for agent tool timeout resolution (Issue #1890). + +Covers: +- registry / policy / definition timeout fields +- ``_resolve_per_tool_timeout`` wiring (first-wins precedence) +- end-to-end ``_execute_tools`` per-tool timeout + fail-open behaviour +- cooperative-cancel wiring into ``check_tool_execution`` + +NOTE: ``ToolRegistry`` defines ``__len__`` but not ``__bool__``, so an *empty* +registry is falsy in Python and ``@tool(registry=empty)`` falls back to the +global default registry. These tests therefore register tools directly via +``ToolDefinition`` (the same path ``factory.get_tool_registry`` uses) instead +of relying on the decorator's registry fallback. +""" +import gc +import json +import re +import sys +import threading +import time +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from src.agent.factory import _build_category_timeout_map + +# Mock heavy optional deps before importing agent modules (mirrors +# tests/agent/test_runtime_facts.py so the suite runs without litellm). +sys.modules.setdefault("litellm", MagicMock()) + +from src.agent.tools.registry import ( + ToolDefinition, + ToolPolicy, + ToolRegistry, + tool, +) +from src.agent.runner import _execute_tools, _resolve_per_tool_timeout +from src.agent.tools.execution import _build_tool_cache_key + + +# --------------------------------------------------------------------------- +# 1. Registry / policy / definition fields +# --------------------------------------------------------------------------- +class TestRegistryTimeoutFields: + def test_category_default_timeout_lookup(self): + reg = ToolRegistry(category_timeout_map={"data": 30.0, "action": 5.0}) + assert reg.category_default_timeout("data") == 30.0 + assert reg.category_default_timeout("action") == 5.0 + assert reg.category_default_timeout("analysis") is None + + def test_policy_declared_carries_timeout(self): + assert ToolPolicy.declared(read_only=True, timeout_seconds=7.0).timeout_seconds == 7.0 + assert ToolPolicy.declared(read_only=True).timeout_seconds is None + + def test_definition_carries_timeout(self): + d = ToolDefinition( + name="x", description="x", parameters=[], handler=lambda: None, + timeout_seconds=3.0, + ) + assert d.timeout_seconds == 3.0 + + def test_decorator_exposes_timeout(self): + reg = ToolRegistry() + + @tool(name="sample", category="data", description="sample", registry=reg, timeout_seconds=4.0) + def sample(): + return 1 + + assert getattr(sample, "_tool_definition", None) is not None + assert sample._tool_definition.timeout_seconds == 4.0 + + +# --------------------------------------------------------------------------- +# 3. Runner wiring +# --------------------------------------------------------------------------- +def _make_tool_call(name, args=None, tc_id="call_1"): + return SimpleNamespace(name=name, arguments=args or {}, id=tc_id) + + +def _register(reg, name, fn, *, category="data", timeout_seconds=None): + reg.register(ToolDefinition( + name=name, description=name, parameters=[], handler=fn, + category=category, timeout_seconds=timeout_seconds, + )) + + +class TestResolvePerToolTimeout: + def test_per_tool_overrides_category(self): + reg = ToolRegistry(category_timeout_map={"data": 30.0}) + _register(reg, "t", lambda: None, category="data", timeout_seconds=2.0) + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg) == 2.0 + + def test_category_used_when_no_per_tool(self): + reg = ToolRegistry(category_timeout_map={"data": 30.0}) + _register(reg, "t", lambda: None, category="data") + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg) == 30.0 + + def test_explicit_per_run_beats_per_tool_and_category(self): + reg = ToolRegistry(category_timeout_map={"data": 30.0}) + _register(reg, "t", lambda: None, category="data", timeout_seconds=2.0) + # first-wins: explicit tool_call_timeout_seconds is the highest priority. + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, 60.0) == 60.0 + + def test_wall_clock_budget_caps_everything(self): + reg = ToolRegistry(category_timeout_map={"data": 30.0}) + _register(reg, "t", lambda: None, category="data", timeout_seconds=2.0) + # remaining budget of 1.0 caps the resolved timeout. + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, 60.0, 1.0) == 1.0 + + def test_no_limits_returns_budget_or_none(self): + reg = ToolRegistry() + _register(reg, "t", lambda: None, category="data") + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg) is None + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, None, 5.0) == 5.0 + + +# --------------------------------------------------------------------------- +# 4. End-to-end _execute_tools +# --------------------------------------------------------------------------- +class TestExecuteToolsTimeout: + def test_per_tool_timeout_fires_and_fail_open(self): + reg = ToolRegistry(category_timeout_map={"data": 0.2}) + + def slow(): + time.sleep(1.0) + return {"ok": True} + + _register(reg, "slow", slow, category="data") + log = [] + results = _execute_tools( + [_make_tool_call("slow")], reg, step=1, + progress_callback=None, tool_calls_log=log, + tool_wait_timeout_seconds=None, + ) + assert len(results) == 1 + parsed = json.loads(results[0]["result_str"]) + assert parsed.get("timeout") is True + assert any(e.get("timeout") is True for e in log) + + def test_no_timeout_when_fast_enough(self): + reg = ToolRegistry(category_timeout_map={"data": 2.0}) + + def fast(): + return {"ok": True} + + _register(reg, "fast", fast, category="data") + log = [] + results = _execute_tools( + [_make_tool_call("fast")], reg, step=1, + progress_callback=None, tool_calls_log=log, + tool_wait_timeout_seconds=None, + ) + assert json.loads(results[0]["result_str"]).get("ok") is True + assert not any(e.get("timeout") for e in log) + + def test_backward_compat_no_limits_runs_inline(self): + # No category map, no global timeout -> tool executes inline, no + # spurious timeout. + reg = ToolRegistry() + + def plain(): + return {"ok": True} + + _register(reg, "plain", plain, category="data") + log = [] + results = _execute_tools( + [_make_tool_call("plain")], reg, step=1, + progress_callback=None, tool_calls_log=log, + tool_wait_timeout_seconds=None, + ) + assert json.loads(results[0]["result_str"]).get("ok") is True + + def test_parallel_per_tool_timeout(self): + reg = ToolRegistry(category_timeout_map={"data": 0.2}) + + def slowA(): + time.sleep(1.0) + return {"a": True} + + def slowB(): + time.sleep(1.0) + return {"b": True} + + _register(reg, "slowA", slowA, category="data") + _register(reg, "slowB", slowB, category="data") + log = [] + results = _execute_tools( + [_make_tool_call("slowA", tc_id="c1"), _make_tool_call("slowB", tc_id="c2")], + reg, step=1, progress_callback=None, tool_calls_log=log, + tool_wait_timeout_seconds=None, + ) + assert len(results) == 2 + assert all(json.loads(r["result_str"]).get("timeout") is True for r in results) + + def test_mixed_fast_and_slow_parallel(self): + """Regression (review): a parallel batch mixing a fast and a slow tool + must let the fast one succeed while the slow one times out at its own + limit — the fast result is not blocked, the batch is bounded by the slow + tool's timeout, and each duration is accurate (the slow entry reports its + own per-tool timeout, not a batch-wide value). + """ + reg = ToolRegistry(category_timeout_map={"data": 0.3}) + + def fast(): + return {"ok": True} + + def slow(): + time.sleep(2.0) + return {"ok": True} + + _register(reg, "fast", fast, category="data") + _register(reg, "slow", slow, category="data") + log = [] + start = time.time() + results = _execute_tools( + [_make_tool_call("fast", tc_id="c1"), _make_tool_call("slow", tc_id="c2")], + reg, step=1, progress_callback=None, tool_calls_log=log, + tool_wait_timeout_seconds=None, + ) + elapsed = time.time() - start + + fast_res = next(r for r in results if r["tc"].name == "fast") + slow_res = next(r for r in results if r["tc"].name == "slow") + assert json.loads(fast_res["result_str"]).get("ok") is True + assert json.loads(slow_res["result_str"]).get("timeout") is True + + fast_entry = next(e for e in log if e["tool"] == "fast") + slow_entry = next(e for e in log if e["tool"] == "slow") + assert fast_entry["success"] is True + assert slow_entry["timeout"] is True + assert slow_entry["duration"] == 0.3 # the slow tool's own limit, accurate + assert elapsed < 1.5 # bounded by 0.3s, not the 2s body + + def test_queued_call_does_not_burn_timeout_before_start(self): + """Review OR-COM-3d6b61f8: with the pool capped at 5, a 6th queued call + must not time out from a deadline that started at *submission*. The fast + 6th tool waits behind five slow (timeout-bound) siblings, then starts and + completes within its own limit. + """ + reg = ToolRegistry(category_timeout_map={"data": 0.2}) + + def slow(): + time.sleep(0.5) + return {"ok": True} + + def fast(): + return {"ok": True} + + _register(reg, "slow", slow, category="data") + _register(reg, "fast", fast, category="data") + tool_calls = [_make_tool_call("slow", tc_id=f"s{i}") for i in range(5)] + [ + _make_tool_call("fast", tc_id="f1") + ] + results = _execute_tools( + tool_calls, reg, step=1, progress_callback=None, tool_calls_log=[], + tool_wait_timeout_seconds=None, + ) + fast_res = next(r for r in results if r["tc"].name == "fast") + # The queued fast tool must run and succeed — never a false timeout that + # started counting before it got a worker. + assert json.loads(fast_res["result_str"]).get("ok") is True + # All 5 slow siblings timed out at their own 0.2s limit. + slow_ress = [r for r in results if r["tc"].name == "slow"] + assert len(slow_ress) == 5 + assert all(json.loads(r["result_str"]).get("timeout") is True for r in slow_ress) + + +# --------------------------------------------------------------------------- +# 5. Config env contract (Issue #1890 test point 1: AGENT_DATA_TOOL_TIMEOUT_S=20) +# --------------------------------------------------------------------------- +class TestConfigEnvContract: + def test_field_names_match_issue_contract(self): + # Issue #1890 specifies env names like AGENT_DATA_TOOL_TIMEOUT_S. + from src.config import Config + + names = {f.name for f in Config.__dataclass_fields__.values()} + assert "agent_data_tool_timeout_s" in names + assert "agent_search_tool_timeout_s" in names + assert "agent_analysis_tool_timeout_s" in names + assert "agent_action_tool_timeout_s" in names + + def test_field_default_is_zero_backward_compatible(self): + # No env set -> 0.0 -> category default disabled -> behaves like today. + from src.config import Config + + default = Config.__dataclass_fields__["agent_data_tool_timeout_s"].default + assert default == 0.0 + + def test_env_value_is_parsed(self, monkeypatch): + import os + + import src.config as cfg + + monkeypatch.setenv("AGENT_DATA_TOOL_TIMEOUT_S", "20") + val = cfg.parse_env_float( + os.getenv("AGENT_DATA_TOOL_TIMEOUT_S"), + 0.0, + field_name="AGENT_DATA_TOOL_TIMEOUT_S", + minimum=0.0, + ) + assert val == 20.0 + + +# --------------------------------------------------------------------------- +# 6. Category coverage (Issue #1890 review: market tools must get a ceiling) +# --------------------------------------------------------------------------- +class TestCategoryTimeoutMap: + def test_market_tools_share_data_timeout(self): + from src.agent.factory import _build_category_timeout_map + + config = SimpleNamespace( + agent_data_tool_timeout_s=15.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + mapping = _build_category_timeout_map(config) + assert mapping["data"] == 15.0 + assert mapping["market"] == 15.0 + assert "search" not in mapping + + +# --------------------------------------------------------------------------- +# 7. Tool registry cache invalidation (Issue #1890 follow-up review) +# --------------------------------------------------------------------------- +class TestToolRegistryCacheInvalidation: + """The module-level ``_TOOL_REGISTRY`` cache must not shadow a reloaded + ``Config`` instance. These tests pin the contract: + + * ``reset_tool_registry`` empties the cache. + * ``get_tool_registry(config)`` rebuilds when the resolved per-category + timeouts differ from the ones the cached registry was built with, and + does so even when the new ``Config`` instance happens to reuse the + collected instance's ``id()``. + * ``SystemConfigService._reload_runtime_singletons`` actually invokes + ``reset_tool_registry`` so API/scheduler/bot entry-points observe the + fresh ``AGENT_*_TOOL_TIMEOUT_S`` values. + """ + + @pytest.fixture(autouse=True) + def _isolate_factory_module_state(self): + """Snapshot/restore ``factory._TOOL_REGISTRY`` so this class cannot + bleed cache state across other tests in the same pytest process. + """ + from src.agent import factory + + saved_registry = factory._TOOL_REGISTRY + saved_timeout_map = factory._CACHED_TIMEOUT_MAP + factory.reset_tool_registry() + try: + yield + finally: + factory._TOOL_REGISTRY = saved_registry + factory._CACHED_TIMEOUT_MAP = saved_timeout_map + + def _empty_tool_lists(self): + """Patch the ALL_*_TOOLS module-level names on ``factory`` so the + rebuild loop is a no-op and the test stays independent of the + full tool-registration chain. + """ + from src.agent import factory + + return ( + patch.object(factory, "ALL_DATA_TOOLS", [], create=True), + patch.object(factory, "ALL_ANALYSIS_TOOLS", [], create=True), + patch.object(factory, "ALL_SEARCH_TOOLS", [], create=True), + patch.object(factory, "ALL_MARKET_TOOLS", [], create=True), + patch.object(factory, "ALL_BACKTEST_TOOLS", [], create=True), + ) + + def test_reset_tool_registry_clears_module_cache(self): + from src.agent import factory + + # Sanity: initial reset empties the module-level cache. + factory.reset_tool_registry() + assert factory._TOOL_REGISTRY is None + assert factory._CACHED_TIMEOUT_MAP is None + + # Idempotent: calling reset on an empty cache is a no-op. + factory.reset_tool_registry() + assert factory._TOOL_REGISTRY is None + assert factory._CACHED_TIMEOUT_MAP is None + + def test_get_tool_registry_rebuilds_when_config_changes(self): + """Two callers passing two *distinct* ``Config`` instances must + observe the matching per-category timeouts on each rebuild — never + a stale view from the first build. + """ + from src.agent import factory + + config_a = SimpleNamespace( + agent_data_tool_timeout_s=10.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + config_b = SimpleNamespace( + agent_data_tool_timeout_s=25.0, + agent_search_tool_timeout_s=5.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + # Force the configs to be distinct objects (SimpleNamespace instances + # already have distinct id() unless they're literal `()`). + assert id(config_a) != id(config_b) + + empty_lists = self._empty_tool_lists() + with patch.object( + factory, "_build_category_timeout_map" + ) as build_map, patch( + "src.agent.tools.registry.ToolRegistry" + ) as registry_cls, empty_lists[0], empty_lists[1], empty_lists[2], empty_lists[3], empty_lists[4]: + registry_instance_a = MagicMock(name="registry_a") + registry_instance_b = MagicMock(name="registry_b") + registry_cls.side_effect = [registry_instance_a, registry_instance_b] + build_map.side_effect = lambda cfg: _build_category_timeout_map(cfg) + + first = factory.get_tool_registry(config_a) + second = factory.get_tool_registry(config_b) + + # Two distinct registry objects — never the cached first build. + assert first is registry_instance_a + assert second is registry_instance_b + assert first is not second + # And the cache now points at the *second* one. + assert factory._TOOL_REGISTRY is registry_instance_b + assert factory._CACHED_TIMEOUT_MAP == { + "data": 25.0, + "search": 5.0, + "market": 25.0, + } + # ToolRegistry was constructed with the *matching* timeout map. + assert registry_cls.call_args_list[0].kwargs["category_timeout_map"] == { + "data": 10.0, + "market": 10.0, + } + assert registry_cls.call_args_list[1].kwargs["category_timeout_map"] == { + "data": 25.0, + "search": 5.0, + "market": 25.0, + } + + def test_get_tool_registry_reuses_cache_for_equivalent_config(self): + """Repeated calls carrying the same effective timeouts must reuse the + cache so the tool-registration cost is paid at most once. + + The final call passes a *different object* with identical values: the + registry only depends on the resolved timeout map, so rebuilding it + would be pure waste on every request. + """ + from src.agent import factory + + config = SimpleNamespace( + agent_data_tool_timeout_s=10.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + equivalent_config = SimpleNamespace( + agent_data_tool_timeout_s=10.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + empty_lists = self._empty_tool_lists() + with patch.object( + factory, "_build_category_timeout_map" + ) as build_map, patch( + "src.agent.tools.registry.ToolRegistry" + ) as registry_cls, empty_lists[0], empty_lists[1], empty_lists[2], empty_lists[3], empty_lists[4]: + registry_instance = MagicMock(name="registry") + registry_cls.return_value = registry_instance + build_map.side_effect = lambda cfg: _build_category_timeout_map(cfg) + + first = factory.get_tool_registry(config) + second = factory.get_tool_registry(config) + third = factory.get_tool_registry(equivalent_config) + + assert first is second is third is registry_instance + # ToolRegistry was instantiated exactly once across three calls. + assert registry_cls.call_count == 1 + + def test_build_agent_executor_forwards_config_to_registry(self): + """Regression for review blockers OR-COM-dd1e8fa7 / OR-COM-bff42110: the + builder must hand its *caller-supplied* ``config`` to + ``get_tool_registry`` so a distinct / updated config actually re-binds + the category timeouts instead of silently reusing the first cached + registry built from a frozen default config. + """ + from src.agent import factory + + config = SimpleNamespace( + agent_data_tool_timeout_s=12.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + agent_arch="single", + ) + with patch.object(factory, "get_tool_registry") as gtr, patch.object( + factory, "resolve_skill_prompt_state" + ) as rsp, patch("src.agent.llm_adapter.LLMToolAdapter"), patch( + "src.agent.executor.AgentExecutor" + ) as ae_cls: + gtr.return_value = MagicMock(name="registry") + rsp.return_value = MagicMock(skill_manager=MagicMock()) + factory.build_agent_executor(config) + # The registry is (re)built from THIS config, not a frozen default. + assert gtr.call_args == ((config,),) + assert ae_cls.call_args.kwargs["tool_registry"] is gtr.return_value + + def test_rebuild_survives_config_instance_id_reuse(self): + """A reloaded ``Config`` must be honoured even when CPython hands the + new instance the *same* ``id()`` as the collected one. + + ``Config.reset_instance()`` drops the only strong reference, so the + allocator frequently reuses that address for the replacement instance. + Keying the cache on ``id(config)`` would read that as "unchanged" and + keep serving the stale registry — exactly the bug this class guards. + """ + from src.agent import factory + + empty_lists = self._empty_tool_lists() + with patch.object( + factory, "_build_category_timeout_map" + ) as build_map, patch( + "src.agent.tools.registry.ToolRegistry" + ) as registry_cls, empty_lists[0], empty_lists[1], empty_lists[2], empty_lists[3], empty_lists[4]: + registry_old = MagicMock(name="registry_old") + registry_new = MagicMock(name="registry_new") + registry_cls.side_effect = [registry_old, registry_new] + build_map.side_effect = lambda cfg: _build_category_timeout_map(cfg) + + config_old = SimpleNamespace( + agent_data_tool_timeout_s=10.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + factory.get_tool_registry(config_old) + assert factory._TOOL_REGISTRY is registry_old + + # Simulate the worst case rather than relying on the allocator: + # a new config whose id() collides with the collected instance. + stale_id = id(config_old) + del config_old + gc.collect() + + config_new = SimpleNamespace( + agent_data_tool_timeout_s=45.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + # Force a deterministic id() collision so this doubles as a real + # regression guard: a future implementation that re-keys the cache + # on ``id(config)`` would read the new config as "unchanged" and + # wrongly reuse ``registry_old``. Patching ``builtins.id`` (not + # ``factory.id``) is required because ``get_tool_registry`` uses the + # builtin directly. + with patch("builtins.id", return_value=stale_id): + rebuilt = factory.get_tool_registry(config_new) + + # Value-keyed cache still notices the change. + assert rebuilt is registry_new + assert registry_cls.call_args_list[-1].kwargs["category_timeout_map"] == { + "data": 45.0, + "market": 45.0, + } + + def test_get_tool_registry_tolerates_partial_config_objects(self): + """Callers hand ``get_tool_registry`` whatever config they hold, which + in tests is routinely a ``MagicMock`` or a stub missing the timeout + attributes. Neither may explode: ``MagicMock() > 0`` raises + ``TypeError`` and a bare stub raises ``AttributeError``, so the + resolver coerces both to "no category limit". + """ + + class _StubConfig: + """No AGENT_*_TOOL_TIMEOUT_S attributes at all.""" + + assert _build_category_timeout_map(MagicMock()) == {} + assert _build_category_timeout_map(_StubConfig()) == {} + + # Garbage values degrade to "no limit" rather than propagating. + noisy = SimpleNamespace( + agent_data_tool_timeout_s="not-a-number", + agent_search_tool_timeout_s=None, + agent_analysis_tool_timeout_s=7.5, + agent_action_tool_timeout_s=-3.0, + ) + assert _build_category_timeout_map(noisy) == {"analysis": 7.5} + + def test_reload_runtime_singletons_drops_registry(self): + """``SystemConfigService._reload_runtime_singletons`` must include + ``reset_tool_registry``; otherwise long-running API/scheduler/bot + processes will keep using the first build's per-category timeouts. + """ + from src.agent import factory + from src.services import system_config_service + + with patch.object(factory, "reset_tool_registry") as reset_mock: + system_config_service.SystemConfigService._reload_runtime_singletons() + + reset_mock.assert_called_once() + + def test_reload_runtime_singletons_actually_invalidates_cached_registry(self): + """End-to-end: with the cache populated against config_old, calling + ``_reload_runtime_singletons`` must clear it so the next + ``get_tool_registry(config_new)`` rebuilds against the new config. + """ + from src.agent import factory + + config_old = SimpleNamespace( + agent_data_tool_timeout_s=10.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + config_new = SimpleNamespace( + agent_data_tool_timeout_s=99.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + + empty_lists = self._empty_tool_lists() + with patch.object( + factory, "_build_category_timeout_map" + ) as build_map, patch( + "src.agent.tools.registry.ToolRegistry" + ) as registry_cls, empty_lists[0], empty_lists[1], empty_lists[2], empty_lists[3], empty_lists[4]: + registry_old = MagicMock(name="registry_old") + registry_new = MagicMock(name="registry_new") + registry_cls.side_effect = [registry_old, registry_new] + build_map.side_effect = lambda cfg: _build_category_timeout_map(cfg) + + from src.services import system_config_service + + # 1) Prime the cache against the old config. + factory.get_tool_registry(config_old) + assert factory._TOOL_REGISTRY is registry_old + + # 2) Simulate a runtime reload — the service-side hook clears + # the cache (real call, not mocked, so we exercise the wiring). + system_config_service.SystemConfigService._reload_runtime_singletons() + assert factory._TOOL_REGISTRY is None + assert factory._CACHED_TIMEOUT_MAP is None + + # 3) Next access, with the new config, must rebuild and observe + # the new timeout. + rebuilt = factory.get_tool_registry(config_new) + assert rebuilt is registry_new + assert factory._TOOL_REGISTRY is registry_new + assert registry_cls.call_args_list[-1].kwargs["category_timeout_map"] == { + "data": 99.0, + "market": 99.0, + } + + +# --------------------------------------------------------------------------- +# 8. Maintainer review blockers (first-wins precedence, non-retriable timeouts, +# finite validation, cache thread-safety) +# --------------------------------------------------------------------------- +class TestFirstWinsTimeoutPrecedence: + """Blocker: first-wins precedence — explicit per-run + ``tool_call_timeout_seconds`` > per-tool declaration > category default > + none — with the remaining wall-clock budget as an unbreakable outer cap. + The previous ``min()`` across all levels let a smaller category default + override an explicit declaration and made an explicit per-run value unable + to relax a stricter per-tool timeout. + """ + + def test_explicit_per_tool_beats_smaller_category(self): + reg = ToolRegistry(category_timeout_map={"data": 10.0}) + _register(reg, "t", lambda: None, category="data", timeout_seconds=30.0) + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg) == 30.0 + + def test_explicit_per_tool_beats_larger_category(self): + reg = ToolRegistry(category_timeout_map={"data": 50.0}) + _register(reg, "t", lambda: None, category="data", timeout_seconds=20.0) + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg) == 20.0 + + def test_explicit_per_run_overrides_per_tool_and_category(self): + reg = ToolRegistry(category_timeout_map={"data": 10.0}) + _register(reg, "t", lambda: None, category="data", timeout_seconds=30.0) + # the caller's explicit per-run value relaxes a stricter per-tool 30s. + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, 60.0) == 60.0 + # and overrides a larger category default too. + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, 15.0) == 15.0 + + def test_wall_clock_budget_is_outer_cap_only(self): + reg = ToolRegistry(category_timeout_map={"data": 10.0}) + _register(reg, "t", lambda: None, category="data", timeout_seconds=30.0) + # base 30 capped to the 15s remaining budget. + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, None, 15.0) == 15.0 + # explicit 60 also capped by the remaining budget. + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, 60.0, 15.0) == 15.0 + # budget disabled -> the first-wins winner stands. + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, None, None) == 30.0 + + def test_explicit_per_run_relaxes_per_tool_end_to_end(self): + reg = ToolRegistry(category_timeout_map={"data": 0.2}) + calls = {"n": 0} + + def slow(): + calls["n"] += 1 + time.sleep(0.8) + return {"ok": True} + + _register(reg, "slow", slow, category="data", timeout_seconds=0.1) + results = _execute_tools( + [_make_tool_call("slow")], reg, step=1, + progress_callback=None, tool_calls_log=[], + tool_call_timeout_seconds=0.6, # explicit override > per-tool 0.1s + tool_wait_timeout_seconds=None, + ) + parsed = json.loads(results[0]["result_str"]) + assert parsed.get("timeout") is True + # The effective timeout is the explicit 0.6s, not the 0.1s per-tool + # declaration — proving the explicit value won the first-wins chain. + match = re.search(r"after ([\d.]+)s", results[0]["result_str"]) + assert match and abs(float(match.group(1)) - 0.6) < 0.05 + + +class TestFilteredRegistryCarriesTimeoutMap: + """Review OR-COM-7f3d3f5b: a filtered registry (``BaseAgent._filtered_registry``) + must retain the source registry's per-category timeout map, otherwise the + category ceilings silently stop applying on tool subsets. + """ + + def test_category_timeout_map_property(self): + reg = ToolRegistry(category_timeout_map={"data": 5.0, "search": 10.0}) + assert reg.category_timeout_map == {"data": 5.0, "search": 10.0} + # returns a copy, not the internal dict + reg.category_timeout_map["data"] = 99.0 + assert reg.category_default_timeout("data") == 5.0 + + def test_filtered_copy_keeps_category_default(self): + reg = ToolRegistry(category_timeout_map={"data": 5.0}) + _register(reg, "t", lambda: None, category="data") + # Mimic BaseAgent._filtered_registry: rebuild with the source map. + from src.agent.tools.registry import ToolRegistry as TR + + filtered = TR(category_timeout_map=reg.category_timeout_map) + for name in reg.list_names(): + filtered.register(reg.get(name)) + # The category ceiling still applies on the filtered subset. + assert _resolve_per_tool_timeout(_make_tool_call("t"), filtered, None, None) == 5.0 + + +class TestTimeoutResultNonRetriable: + """Blocker: a timed-out call must be marked ``retriable: False`` *and* + recorded in ``non_retriable_tool_results`` so an LLM retry of the same call + reuses the cached failure instead of spinning up a second (side-effecting) + execution. Python cannot forcibly cancel an already-started tool thread, so + this is the best-effort guard against duplicate work. + """ + + def test_timeout_result_is_marked_non_retriable(self): + reg = ToolRegistry(category_timeout_map={"data": 0.2}) + + def slow(): + time.sleep(1.0) + return {"ok": True} + + _register(reg, "slow", slow, category="data") + results = _execute_tools( + [_make_tool_call("slow")], reg, step=1, + progress_callback=None, tool_calls_log=[], + tool_wait_timeout_seconds=None, + ) + parsed = json.loads(results[0]["result_str"]) + assert parsed.get("timeout") is True + assert parsed.get("retriable") is False + + def test_timed_out_tool_not_re_executed_on_retry(self): + reg = ToolRegistry(category_timeout_map={"data": 0.2}) + calls = {"n": 0} + + def slow(): + calls["n"] += 1 + time.sleep(1.0) + return {"ok": True} + + _register(reg, "slow", slow, category="data") + shared_non_retriable = {} + + # First call times out and is recorded as non-retriable. + res1 = _execute_tools( + [_make_tool_call("slow", tc_id="c1")], reg, step=1, + progress_callback=None, tool_calls_log=[], + tool_wait_timeout_seconds=None, + non_retriable_tool_results=shared_non_retriable, + ) + assert json.loads(res1[0]["result_str"]).get("timeout") is True + assert _build_tool_cache_key("slow", {}) in shared_non_retriable # keyed by name+args + + first_calls = calls["n"] + # Identical retry: must NOT spin up a second execution. + res2 = _execute_tools( + [_make_tool_call("slow", tc_id="c1")], reg, step=2, + progress_callback=None, tool_calls_log=[], + tool_wait_timeout_seconds=None, + non_retriable_tool_results=shared_non_retriable, + ) + assert json.loads(res2[0]["result_str"]).get("timeout") is True + assert calls["n"] == first_calls # handler never ran again + + def test_timeout_with_non_dict_arguments_does_not_write_none_key(self): + """Regression: a timeout payload built from non-dict ``arguments`` + (cache key ``None``) must not write a shared ``None`` entry into + ``non_retriable_tool_results``, which could alias unrelated no-arg calls. + """ + reg = ToolRegistry(category_timeout_map={"data": 0.2}) + + def slow(): + time.sleep(1.0) + return {"ok": True} + + _register(reg, "slow", slow, category="data") + shared_non_retriable = {} + results = _execute_tools( + [_make_tool_call("slow", args=None)], reg, step=1, + progress_callback=None, tool_calls_log=[], + tool_wait_timeout_seconds=None, + non_retriable_tool_results=shared_non_retriable, + ) + assert json.loads(results[0]["result_str"]).get("timeout") is True + assert None not in shared_non_retriable + + +class TestTimeoutFiniteValidation: + """Blocker: ``inf``/``nan``/negative ``AGENT_*_TOOL_TIMEOUT_S`` must degrade + to "no limit" rather than raising ``OverflowError`` at + ``future.result(timeout=inf)``. + """ + + def test_config_inf_degrades_to_no_limit(self): + from src.agent.factory import _coerce_config_timeout + + cfg = SimpleNamespace(agent_data_tool_timeout_s=float("inf")) + assert _coerce_config_timeout(cfg, "agent_data_tool_timeout_s") == 0.0 + + def test_config_nan_degrades_to_no_limit(self): + from src.agent.factory import _coerce_config_timeout + + cfg = SimpleNamespace(agent_data_tool_timeout_s=float("nan")) + assert _coerce_config_timeout(cfg, "agent_data_tool_timeout_s") == 0.0 + + def test_config_excessive_value_is_clamped(self): + from src.agent.factory import _coerce_config_timeout, _MAX_TOOL_TIMEOUT_S + + cfg = SimpleNamespace(agent_data_tool_timeout_s=999999.0) + assert _coerce_config_timeout(cfg, "agent_data_tool_timeout_s") == _MAX_TOOL_TIMEOUT_S + + def test_per_tool_inf_is_not_a_valid_timeout(self): + reg = ToolRegistry(category_timeout_map={}) + _register(reg, "t", lambda: None, category="data", timeout_seconds=float("inf")) + # inf must not survive into the resolved timeout. + assert _resolve_per_tool_timeout(_make_tool_call("t"), reg, None) is None + + +class TestToolRegistryCacheThreadSafety: + """Blocker: ``get_tool_registry`` / ``reset_tool_registry`` must be safe + under concurrent requests with different ``Config`` objects — no partial + assignment of the shared ``_TOOL_REGISTRY`` / ``_CACHED_TIMEOUT_MAP`` pair. + """ + + @pytest.fixture(autouse=True) + def _isolate_factory_module_state(self): + from src.agent import factory + + saved_registry = factory._TOOL_REGISTRY + saved_timeout_map = factory._CACHED_TIMEOUT_MAP + factory.reset_tool_registry() + try: + yield + finally: + factory._TOOL_REGISTRY = saved_registry + factory._CACHED_TIMEOUT_MAP = saved_timeout_map + + def test_lock_is_present(self): + from src.agent import factory + + assert isinstance(factory._tool_registry_lock, type(threading.Lock())) + + def test_fast_path_returns_cached_registry_consistently(self): + """Review OR-COM-a1e8b0c2: the locked fast path must hand back the same + registry on a cache hit and a valid (non-None) one after a reset — no + check-then-use window against ``reset_tool_registry()``. + """ + from src.agent.factory import get_tool_registry, reset_tool_registry + + config = SimpleNamespace( + agent_data_tool_timeout_s=5.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + r1 = get_tool_registry(config) + r2 = get_tool_registry(config) # cache hit via the locked fast path + assert r1 is r2 + reset_tool_registry() + r3 = get_tool_registry(config) + assert r3 is not None + assert r3.category_default_timeout("data") == 5.0 + + def test_concurrent_builds_do_not_raise(self): + from src.agent import factory + import threading + + config_a = SimpleNamespace( + agent_data_tool_timeout_s=10.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + config_b = SimpleNamespace( + agent_data_tool_timeout_s=25.0, + agent_search_tool_timeout_s=0.0, + agent_analysis_tool_timeout_s=0.0, + agent_action_tool_timeout_s=0.0, + ) + + errors = [] + + def worker(idx): + try: + cfg = config_a if idx % 2 == 0 else config_b + for _ in range(25): + reg = factory.get_tool_registry(cfg) + assert reg is not None + except Exception as exc: # pragma: no cover + errors.append(repr(exc)) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors, f"concurrent get_tool_registry raised: {errors}" + + +# --------------------------------------------------------------------------- +# 9. Maintainer review follow-up: cooperative cancel + local-registry return +# --------------------------------------------------------------------------- +class TestTimeoutCooperativeCancel: + """Blocker: a tool handler may keep running after its timeout fires (Python + cannot forcibly stop an already-started thread). The runner must (a) mark + the result non-retriable (covered by ``TestTimeoutResultNonRetriable``) and + (b) arm a cooperative-cancel signal that opt-in handlers can poll via + ``is_tool_cancellation_requested`` so they can abort early. + """ + + def test_helper_defaults_false(self): + from src.agent.tools.execution import is_tool_cancellation_requested + + assert is_tool_cancellation_requested() is False + + def test_helper_true_when_armed(self): + from src.agent.tools.execution import ( + TOOL_CANCEL_EVENT, + is_tool_cancellation_requested, + ) + + import threading + + event = threading.Event() + token = TOOL_CANCEL_EVENT.set(event) + try: + assert is_tool_cancellation_requested() is False + event.set() + assert is_tool_cancellation_requested() is True + finally: + TOOL_CANCEL_EVENT.reset(token) + + def test_runner_arms_cancel_on_timeout(self): + from src.agent.tools.execution import is_tool_cancellation_requested + + import time as _time + + captured = {"requested": None} + + def slow_handler(**kwargs): + # Poll cancellation the way an opt-in handler would. + deadline = _time.time() + 3.0 + while _time.time() < deadline: + if is_tool_cancellation_requested(): + captured["requested"] = True + return {"ok": True} + _time.sleep(0.02) + captured["requested"] = is_tool_cancellation_requested() + return {"ok": True} + + reg = ToolRegistry() + reg.register(ToolDefinition( + name="slow_tool", description="s", parameters=[], + handler=slow_handler, category="data", + )) + tool_calls = [SimpleNamespace(name="slow_tool", arguments={})] + results = _execute_tools( + tool_calls, reg, 1, None, [], + non_retriable_tool_results={}, + tool_wait_timeout_seconds=0.1, + ) + # Give the still-running background handler time to observe the armed signal. + _time.sleep(0.8) + assert len(results) == 1 + parsed = json.loads(results[0]["result_str"]) + assert parsed.get("timeout") is True + assert parsed.get("retriable") is False + assert captured["requested"] is True + + def test_check_tool_execution_honors_runner_cancel(self): + """Regression: ``check_tool_execution()`` — the checkpoint real + data/backtest/tool-surface tools already call — must observe the + runner-armed ``TOOL_CANCEL_EVENT`` and abort a timed-out handler early, + instead of letting it run to completion in the background thread. + """ + from src.agent.tools.execution import check_tool_execution + + import time as _time + + captured = {"completed": False} + + def slow_handler(**kwargs): + # Poll the checkpoint the way data_tools / backtest_tools do. + deadline = _time.time() + 3.0 + while _time.time() < deadline: + check_tool_execution() + _time.sleep(0.02) + captured["completed"] = True + return {"ok": True} + + reg = ToolRegistry() + reg.register(ToolDefinition( + name="checkpoint_tool", description="s", parameters=[], + handler=slow_handler, category="data", + )) + tool_calls = [SimpleNamespace(name="checkpoint_tool", arguments={})] + results = _execute_tools( + tool_calls, reg, 1, None, [], + non_retriable_tool_results={}, + tool_wait_timeout_seconds=0.1, + ) + # The handler aborts at its next checkpoint shortly after the 0.1s + # timeout; it must NOT run its full 3s body to completion. + _time.sleep(0.8) + assert len(results) == 1 + parsed = json.loads(results[0]["result_str"]) + assert parsed.get("timeout") is True + assert captured["completed"] is False + + +class TestGetToolRegistryReturnsLocalRegistry: + """Blocker: ``get_tool_registry`` must return the registry it just built for + the caller's timeout map (not the shared global), so a concurrent rebuild + for a different ``Config`` cannot leak a mismatched registry into this call. + """ + + @pytest.fixture(autouse=True) + def _isolate(self): + from src.agent import factory + + saved_registry = factory._TOOL_REGISTRY + saved_timeout_map = factory._CACHED_TIMEOUT_MAP + factory.reset_tool_registry() + try: + yield + finally: + factory._TOOL_REGISTRY = saved_registry + factory._CACHED_TIMEOUT_MAP = saved_timeout_map + + def test_returns_locally_built_registry(self, monkeypatch): + from src.agent import factory + + sentinel = ToolRegistry({"data": 11}) + monkeypatch.setattr(factory, "_build_tool_registry", lambda m: sentinel) + cfg = SimpleNamespace(agent_data_tool_timeout_s=11) + assert factory.get_tool_registry(cfg) is sentinel