fix(agent): add LLM request timeout to prevent session lock starvation
This commit is contained in:
+30
-4
@@ -5,6 +5,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
from dataclasses import dataclass, field
|
||||
import inspect
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
@@ -13,7 +14,7 @@ from loguru import logger
|
||||
from nanobot.agent.hook import AgentHook, AgentHookContext
|
||||
from nanobot.utils.prompt_templates import render_template
|
||||
from nanobot.agent.tools.registry import ToolRegistry
|
||||
from nanobot.providers.base import LLMProvider, ToolCallRequest
|
||||
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
|
||||
from nanobot.utils.helpers import (
|
||||
build_assistant_message,
|
||||
estimate_message_tokens,
|
||||
@@ -74,6 +75,7 @@ class AgentRunSpec:
|
||||
retry_wait_callback: Any | None = None
|
||||
checkpoint_callback: Any | None = None
|
||||
injection_callback: Any | None = None
|
||||
llm_timeout_s: float | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -570,6 +572,19 @@ class AgentRunner:
|
||||
hook: AgentHook,
|
||||
context: AgentHookContext,
|
||||
):
|
||||
timeout_s: float | None = spec.llm_timeout_s
|
||||
if timeout_s is None:
|
||||
# Default to a finite timeout to avoid per-session lock starvation when an LLM
|
||||
# request hangs indefinitely (e.g. gateway/network stall).
|
||||
# Set NANOBOT_LLM_TIMEOUT_S=0 to disable.
|
||||
raw = os.environ.get("NANOBOT_LLM_TIMEOUT_S", "300").strip()
|
||||
try:
|
||||
timeout_s = float(raw)
|
||||
except (TypeError, ValueError):
|
||||
timeout_s = 300.0
|
||||
if timeout_s is not None and timeout_s <= 0:
|
||||
timeout_s = None
|
||||
|
||||
kwargs = self._build_request_kwargs(
|
||||
spec,
|
||||
messages,
|
||||
@@ -579,11 +594,23 @@ class AgentRunner:
|
||||
async def _stream(delta: str) -> None:
|
||||
await hook.on_stream(context, delta)
|
||||
|
||||
return await self.provider.chat_stream_with_retry(
|
||||
coro = self.provider.chat_stream_with_retry(
|
||||
**kwargs,
|
||||
on_content_delta=_stream,
|
||||
)
|
||||
return await self.provider.chat_with_retry(**kwargs)
|
||||
else:
|
||||
coro = self.provider.chat_with_retry(**kwargs)
|
||||
|
||||
if timeout_s is None:
|
||||
return await coro
|
||||
try:
|
||||
return await asyncio.wait_for(coro, timeout=timeout_s)
|
||||
except asyncio.TimeoutError:
|
||||
return LLMResponse(
|
||||
content=f"Error calling LLM: timed out after {timeout_s:g}s",
|
||||
finish_reason="error",
|
||||
error_kind="timeout",
|
||||
)
|
||||
|
||||
async def _request_finalization_retry(
|
||||
self,
|
||||
@@ -984,4 +1011,3 @@ class AgentRunner:
|
||||
if current:
|
||||
batches.append(current)
|
||||
return batches
|
||||
|
||||
|
||||
@@ -252,6 +252,35 @@ async def test_runner_returns_max_iterations_fallback():
|
||||
assert result.messages[-1]["role"] == "assistant"
|
||||
assert result.messages[-1]["content"] == result.final_content
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_times_out_hung_llm_request():
|
||||
from nanobot.agent.runner import AgentRunSpec, AgentRunner
|
||||
|
||||
provider = MagicMock()
|
||||
|
||||
async def chat_with_retry(**kwargs):
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
provider.chat_with_retry = chat_with_retry
|
||||
tools = MagicMock()
|
||||
tools.get_definitions.return_value = []
|
||||
|
||||
runner = AgentRunner(provider)
|
||||
started = time.monotonic()
|
||||
result = await runner.run(AgentRunSpec(
|
||||
initial_messages=[{"role": "user", "content": "hello"}],
|
||||
tools=tools,
|
||||
model="test-model",
|
||||
max_iterations=1,
|
||||
max_tool_result_chars=_MAX_TOOL_RESULT_CHARS,
|
||||
llm_timeout_s=0.05,
|
||||
))
|
||||
|
||||
assert (time.monotonic() - started) < 1.0
|
||||
assert result.stop_reason == "error"
|
||||
assert "timed out" in (result.final_content or "").lower()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_runner_returns_structured_tool_error():
|
||||
from nanobot.agent.runner import AgentRunSpec, AgentRunner
|
||||
|
||||
Reference in New Issue
Block a user