fix(agent): add LLM request timeout to prevent session lock starvation

This commit is contained in:
yorkhellen
2026-04-25 03:40:34 +08:00
committed by Xubin Ren
parent e52fe2a8e2
commit 076e4166d7
2 changed files with 59 additions and 4 deletions
+30 -4
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
import asyncio import asyncio
from dataclasses import dataclass, field from dataclasses import dataclass, field
import inspect import inspect
import os
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@@ -13,7 +14,7 @@ from loguru import logger
from nanobot.agent.hook import AgentHook, AgentHookContext from nanobot.agent.hook import AgentHook, AgentHookContext
from nanobot.utils.prompt_templates import render_template from nanobot.utils.prompt_templates import render_template
from nanobot.agent.tools.registry import ToolRegistry 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 ( from nanobot.utils.helpers import (
build_assistant_message, build_assistant_message,
estimate_message_tokens, estimate_message_tokens,
@@ -74,6 +75,7 @@ class AgentRunSpec:
retry_wait_callback: Any | None = None retry_wait_callback: Any | None = None
checkpoint_callback: Any | None = None checkpoint_callback: Any | None = None
injection_callback: Any | None = None injection_callback: Any | None = None
llm_timeout_s: float | None = None
@dataclass(slots=True) @dataclass(slots=True)
@@ -570,6 +572,19 @@ class AgentRunner:
hook: AgentHook, hook: AgentHook,
context: AgentHookContext, 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( kwargs = self._build_request_kwargs(
spec, spec,
messages, messages,
@@ -579,11 +594,23 @@ class AgentRunner:
async def _stream(delta: str) -> None: async def _stream(delta: str) -> None:
await hook.on_stream(context, delta) await hook.on_stream(context, delta)
return await self.provider.chat_stream_with_retry( coro = self.provider.chat_stream_with_retry(
**kwargs, **kwargs,
on_content_delta=_stream, 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( async def _request_finalization_retry(
self, self,
@@ -984,4 +1011,3 @@ class AgentRunner:
if current: if current:
batches.append(current) batches.append(current)
return batches return batches
+29
View File
@@ -252,6 +252,35 @@ async def test_runner_returns_max_iterations_fallback():
assert result.messages[-1]["role"] == "assistant" assert result.messages[-1]["role"] == "assistant"
assert result.messages[-1]["content"] == result.final_content 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 @pytest.mark.asyncio
async def test_runner_returns_structured_tool_error(): async def test_runner_returns_structured_tool_error():
from nanobot.agent.runner import AgentRunSpec, AgentRunner from nanobot.agent.runner import AgentRunSpec, AgentRunner