wehub-resource-sync
222 lines
7.1 KiB
Python
222 lines
7.1 KiB
Python
"""Tool call retry logic with exponential backoff and timeout.
|
|
|
|
Wraps tool invocations with configurable timeout enforcement and
|
|
retry on transient failures. Errors are returned as strings to the
|
|
LLM (not raised) so the agent can decide how to proceed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import concurrent.futures
|
|
import logging
|
|
import time
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
from omnigent.spec.types import RetryPolicy, ToolsConfig
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def resolve_tool_timeout(
|
|
tool_name: str,
|
|
tools_config: ToolsConfig,
|
|
per_tool_timeout: int | None,
|
|
) -> int:
|
|
"""
|
|
Resolve the effective timeout for a tool call.
|
|
|
|
:param tool_name: Tool name for logging, e.g.
|
|
``"github.list_issues"``.
|
|
:param tools_config: The agent's global tools config.
|
|
:param per_tool_timeout: Per-tool timeout override. ``None``
|
|
inherits ``tools_config.timeout``.
|
|
:returns: Effective timeout in seconds.
|
|
"""
|
|
if per_tool_timeout is not None:
|
|
return per_tool_timeout
|
|
return tools_config.timeout
|
|
|
|
|
|
def resolve_tool_retry(
|
|
tool_name: str,
|
|
tools_config: ToolsConfig,
|
|
per_tool_retry: RetryPolicy | None,
|
|
) -> RetryPolicy:
|
|
"""
|
|
Resolve the effective retry config for a tool call.
|
|
|
|
:param tool_name: Tool name for logging, e.g.
|
|
``"github.list_issues"``.
|
|
:param tools_config: The agent's global tools config.
|
|
:param per_tool_retry: Per-tool retry override. ``None``
|
|
inherits ``tools_config.retry``.
|
|
:returns: Effective retry config.
|
|
"""
|
|
if per_tool_retry is not None:
|
|
return per_tool_retry
|
|
return tools_config.retry
|
|
|
|
|
|
def call_tool_with_timeout(
|
|
call_fn: Callable[[], str],
|
|
timeout: int,
|
|
cancel_fn: Callable[[], None] | None = None,
|
|
) -> str:
|
|
"""
|
|
Execute a tool call with a wall-clock timeout.
|
|
|
|
Uses a thread pool to enforce the timeout. If the tool
|
|
exceeds the deadline, calls ``cancel_fn`` (if provided)
|
|
to kill the underlying process, then raises ``TimeoutError``.
|
|
|
|
:param call_fn: Zero-argument callable that executes the
|
|
tool and returns a string result.
|
|
:param timeout: Timeout in seconds, e.g. ``60``.
|
|
:param cancel_fn: Optional callable to invoke on timeout,
|
|
e.g. ``tool.cancel`` which sends SIGKILL to a
|
|
subprocess. ``None`` means no cancellation action.
|
|
:returns: The tool's string result.
|
|
:raises TimeoutError: If the tool does not complete within
|
|
the timeout.
|
|
"""
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
|
|
future = pool.submit(call_fn)
|
|
try:
|
|
return future.result(timeout=timeout)
|
|
except concurrent.futures.TimeoutError:
|
|
if cancel_fn is not None:
|
|
cancel_fn()
|
|
raise TimeoutError(f"Tool execution timed out after {timeout}s") from None
|
|
|
|
|
|
def execute_tool_with_retry(
|
|
tool_name: str,
|
|
call_fn: Callable[[], str],
|
|
timeout: int,
|
|
retry_config: RetryPolicy,
|
|
on_event: Callable[[dict[str, Any]], None],
|
|
cancel_fn: Callable[[], None] | None = None,
|
|
) -> str:
|
|
"""
|
|
Execute a tool call with timeout and retry.
|
|
|
|
On transient failure (timeout), retries with exponential
|
|
backoff. When all retries are exhausted, returns an error
|
|
string (not raised) so the LLM can decide how to proceed.
|
|
|
|
:param tool_name: The tool name for error messages and SSE
|
|
events, e.g. ``"github.list_issues"``.
|
|
:param call_fn: Zero-argument callable that executes the
|
|
tool and returns a string result.
|
|
:param timeout: Per-call timeout in seconds, e.g. ``60``.
|
|
:param retry_config: Retry policy for this tool.
|
|
:param on_event: Callback to emit SSE events (retry and
|
|
error). Called with the event dict.
|
|
:param cancel_fn: Optional callable to invoke on timeout,
|
|
e.g. ``tool.cancel`` which sends SIGKILL to a
|
|
subprocess. ``None`` means no cancellation action.
|
|
:returns: The tool's string result, or an error string if
|
|
all retries are exhausted.
|
|
"""
|
|
last_error: str | None = None
|
|
total_tries = retry_config.max_retries + 1
|
|
|
|
for attempt in range(total_tries):
|
|
try:
|
|
return call_tool_with_timeout(call_fn, timeout, cancel_fn)
|
|
except TimeoutError as exc:
|
|
last_error = str(exc)
|
|
if attempt + 1 < total_tries:
|
|
_emit_tool_retry(
|
|
tool_name,
|
|
attempt,
|
|
retry_config,
|
|
"timeout",
|
|
last_error,
|
|
on_event,
|
|
)
|
|
except Exception as exc:
|
|
last_error = f"Tool '{tool_name}' raised {type(exc).__name__}: {exc}"
|
|
# Non-timeout exceptions are not retried.
|
|
break
|
|
|
|
_emit_tool_error(tool_name, total_tries, last_error, on_event)
|
|
return f"Error: {last_error} ({total_tries} attempts)"
|
|
|
|
|
|
def _emit_tool_retry(
|
|
tool_name: str,
|
|
attempt: int,
|
|
retry_config: RetryPolicy,
|
|
error_code: str,
|
|
error_message: str,
|
|
on_event: Callable[[dict[str, Any]], None],
|
|
) -> None:
|
|
"""
|
|
Emit a ``response.retry`` SSE event for a tool and sleep.
|
|
|
|
:param tool_name: The tool name, e.g. ``"github.list_issues"``.
|
|
:param attempt: Current zero-based attempt index.
|
|
:param retry_config: Retry policy with backoff parameters.
|
|
:param error_code: Error code string, e.g. ``"timeout"``.
|
|
:param error_message: Human-readable error description.
|
|
:param on_event: Callback to emit the SSE event.
|
|
"""
|
|
delay = retry_config.compute_backoff_delay(retry_index=attempt + 1)
|
|
total_tries = retry_config.max_retries + 1
|
|
event: dict[str, Any] = {
|
|
"type": "response.retry",
|
|
"source": "tool",
|
|
"tool_name": tool_name,
|
|
"attempt": attempt + 2,
|
|
"max_attempts": total_tries,
|
|
"delay_seconds": round(delay, 2),
|
|
"error": {
|
|
"code": error_code,
|
|
"message": error_message,
|
|
"detail": None,
|
|
},
|
|
}
|
|
on_event(event)
|
|
_logger.info(
|
|
"Tool %r retry %d/%d after %.1fs: %s",
|
|
tool_name,
|
|
attempt + 2,
|
|
total_tries,
|
|
delay,
|
|
error_code,
|
|
)
|
|
time.sleep(delay)
|
|
|
|
|
|
def _emit_tool_error(
|
|
tool_name: str,
|
|
max_attempts: int,
|
|
error_message: str | None,
|
|
on_event: Callable[[dict[str, Any]], None],
|
|
) -> None:
|
|
"""
|
|
Emit a ``response.error`` SSE event for a terminal tool failure.
|
|
|
|
:param tool_name: The tool name, e.g. ``"github.list_issues"``.
|
|
:param max_attempts: Total attempts configured.
|
|
:param error_message: Human-readable error description.
|
|
:param on_event: Callback to emit the SSE event.
|
|
"""
|
|
# Defensive fallback — error_message should always be set by the
|
|
# retry loop, but "Unknown error" is safe if somehow None.
|
|
msg = error_message or "Unknown error"
|
|
on_event(
|
|
{
|
|
"type": "response.error",
|
|
"source": "tool",
|
|
"tool_name": tool_name,
|
|
"error": {
|
|
"code": "timeout",
|
|
"message": f"{msg} ({max_attempts} attempts exhausted)",
|
|
"detail": None,
|
|
},
|
|
}
|
|
)
|