mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-06 17:38:35 +00:00
fix(providers): validate stream idle timeout config
This commit is contained in:
parent
7bec0f6e01
commit
846410f936
@ -3,7 +3,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import os
|
|
||||||
import re
|
import re
|
||||||
import secrets
|
import secrets
|
||||||
import string
|
import string
|
||||||
@ -14,6 +13,7 @@ from nanobot.providers.base import (
|
|||||||
LLMProvider,
|
LLMProvider,
|
||||||
LLMResponse,
|
LLMResponse,
|
||||||
ToolCallRequest,
|
ToolCallRequest,
|
||||||
|
resolve_stream_idle_timeout_s,
|
||||||
tool_arguments_object_for_replay,
|
tool_arguments_object_for_replay,
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -613,7 +613,7 @@ class AnthropicProvider(LLMProvider):
|
|||||||
messages, tools, model, max_tokens, temperature,
|
messages, tools, model, max_tokens, temperature,
|
||||||
reasoning_effort, tool_choice,
|
reasoning_effort, tool_choice,
|
||||||
)
|
)
|
||||||
idle_timeout_s = int(os.environ.get("NANOBOT_STREAM_IDLE_TIMEOUT_S", "90"))
|
idle_timeout_s = resolve_stream_idle_timeout_s()
|
||||||
try:
|
try:
|
||||||
async with self._client.messages.stream(**kwargs) as stream:
|
async with self._client.messages.stream(**kwargs) as stream:
|
||||||
if on_content_delta or on_thinking_delta or on_tool_call_delta:
|
if on_content_delta or on_thinking_delta or on_tool_call_delta:
|
||||||
@ -682,7 +682,7 @@ class AnthropicProvider(LLMProvider):
|
|||||||
return LLMResponse(
|
return LLMResponse(
|
||||||
content=(
|
content=(
|
||||||
f"Error calling LLM: stream stalled for more than "
|
f"Error calling LLM: stream stalled for more than "
|
||||||
f"{idle_timeout_s} seconds"
|
f"{idle_timeout_s:g} seconds"
|
||||||
),
|
),
|
||||||
finish_reason="error",
|
finish_reason="error",
|
||||||
error_kind="timeout",
|
error_kind="timeout",
|
||||||
|
|||||||
@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
import json
|
import json
|
||||||
|
import os
|
||||||
import re
|
import re
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
@ -16,6 +17,34 @@ from loguru import logger
|
|||||||
|
|
||||||
from nanobot.utils.helpers import image_placeholder_text
|
from nanobot.utils.helpers import image_placeholder_text
|
||||||
|
|
||||||
|
STREAM_IDLE_TIMEOUT_ENV = "NANOBOT_STREAM_IDLE_TIMEOUT_S"
|
||||||
|
DEFAULT_STREAM_IDLE_TIMEOUT_S = 90.0
|
||||||
|
MAX_STREAM_IDLE_TIMEOUT_S = 3600.0
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_stream_idle_timeout_s(
|
||||||
|
*,
|
||||||
|
env_value: str | None = None,
|
||||||
|
default: float = DEFAULT_STREAM_IDLE_TIMEOUT_S,
|
||||||
|
maximum: float = MAX_STREAM_IDLE_TIMEOUT_S,
|
||||||
|
) -> float:
|
||||||
|
"""Return a safe streaming idle timeout from env/config text."""
|
||||||
|
raw = os.environ.get(STREAM_IDLE_TIMEOUT_ENV) if env_value is None else env_value
|
||||||
|
if raw is None or not raw.strip():
|
||||||
|
return default
|
||||||
|
try:
|
||||||
|
value = float(raw)
|
||||||
|
except (TypeError, ValueError):
|
||||||
|
logger.warning("Ignoring invalid {}={!r}; using {}", STREAM_IDLE_TIMEOUT_ENV, raw, default)
|
||||||
|
return default
|
||||||
|
if value <= 0:
|
||||||
|
logger.warning("Ignoring non-positive {}={!r}; using {}", STREAM_IDLE_TIMEOUT_ENV, raw, default)
|
||||||
|
return default
|
||||||
|
if value > maximum:
|
||||||
|
logger.warning("Clamping {}={!r} to {}", STREAM_IDLE_TIMEOUT_ENV, raw, maximum)
|
||||||
|
return maximum
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class ToolCallRequest:
|
class ToolCallRequest:
|
||||||
|
|||||||
@ -15,6 +15,7 @@ from nanobot.providers.base import (
|
|||||||
LLMResponse,
|
LLMResponse,
|
||||||
ToolCallRequest,
|
ToolCallRequest,
|
||||||
parse_tool_arguments,
|
parse_tool_arguments,
|
||||||
|
resolve_stream_idle_timeout_s,
|
||||||
tool_arguments_object_for_replay,
|
tool_arguments_object_for_replay,
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -701,7 +702,7 @@ class BedrockProvider(LLMProvider):
|
|||||||
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
||||||
) -> LLMResponse:
|
) -> LLMResponse:
|
||||||
_ = on_thinking_delta, on_tool_call_delta
|
_ = on_thinking_delta, on_tool_call_delta
|
||||||
idle_timeout_s = int(os.environ.get("NANOBOT_STREAM_IDLE_TIMEOUT_S", "90"))
|
idle_timeout_s = resolve_stream_idle_timeout_s()
|
||||||
content_parts: list[str] = []
|
content_parts: list[str] = []
|
||||||
reasoning_parts: list[str] = []
|
reasoning_parts: list[str] = []
|
||||||
thinking_blocks: list[dict[str, Any]] = []
|
thinking_blocks: list[dict[str, Any]] = []
|
||||||
@ -742,7 +743,7 @@ class BedrockProvider(LLMProvider):
|
|||||||
return LLMResponse(
|
return LLMResponse(
|
||||||
content=(
|
content=(
|
||||||
f"Error calling LLM: stream stalled for more than "
|
f"Error calling LLM: stream stalled for more than "
|
||||||
f"{idle_timeout_s} seconds"
|
f"{idle_timeout_s:g} seconds"
|
||||||
),
|
),
|
||||||
finish_reason="error",
|
finish_reason="error",
|
||||||
error_kind="timeout",
|
error_kind="timeout",
|
||||||
|
|||||||
@ -5,7 +5,6 @@ from __future__ import annotations
|
|||||||
import asyncio
|
import asyncio
|
||||||
import hashlib
|
import hashlib
|
||||||
import json
|
import json
|
||||||
import os
|
|
||||||
from collections.abc import Awaitable, Callable
|
from collections.abc import Awaitable, Callable
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@ -13,7 +12,12 @@ import httpx
|
|||||||
from loguru import logger
|
from loguru import logger
|
||||||
from oauth_cli_kit import get_token as get_codex_token
|
from oauth_cli_kit import get_token as get_codex_token
|
||||||
|
|
||||||
from nanobot.providers.base import LLMProvider, LLMResponse, ToolCallRequest
|
from nanobot.providers.base import (
|
||||||
|
LLMProvider,
|
||||||
|
LLMResponse,
|
||||||
|
ToolCallRequest,
|
||||||
|
resolve_stream_idle_timeout_s,
|
||||||
|
)
|
||||||
from nanobot.providers.openai_responses import (
|
from nanobot.providers.openai_responses import (
|
||||||
consume_sse_with_reasoning,
|
consume_sse_with_reasoning,
|
||||||
convert_messages,
|
convert_messages,
|
||||||
@ -199,7 +203,7 @@ async def _request_codex(
|
|||||||
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
on_thinking_delta: Callable[[str], Awaitable[None]] | None = None,
|
||||||
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
||||||
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
|
) -> tuple[str, list[ToolCallRequest], str, dict[str, int], str | None]:
|
||||||
idle_timeout_s = int(os.environ.get("NANOBOT_STREAM_IDLE_TIMEOUT_S", "90"))
|
idle_timeout_s = resolve_stream_idle_timeout_s()
|
||||||
async with httpx.AsyncClient(timeout=idle_timeout_s, verify=verify) as client:
|
async with httpx.AsyncClient(timeout=idle_timeout_s, verify=verify) as client:
|
||||||
async with client.stream("POST", url, headers=headers, json=body) as response:
|
async with client.stream("POST", url, headers=headers, json=body) as response:
|
||||||
if response.status_code != 200:
|
if response.status_code != 200:
|
||||||
|
|||||||
@ -25,6 +25,7 @@ from nanobot.providers.base import (
|
|||||||
LLMResponse,
|
LLMResponse,
|
||||||
ToolCallRequest,
|
ToolCallRequest,
|
||||||
parse_tool_arguments,
|
parse_tool_arguments,
|
||||||
|
resolve_stream_idle_timeout_s,
|
||||||
tool_arguments_json_for_replay,
|
tool_arguments_json_for_replay,
|
||||||
)
|
)
|
||||||
from nanobot.providers.openai_responses import (
|
from nanobot.providers.openai_responses import (
|
||||||
@ -1386,7 +1387,7 @@ class OpenAICompatProvider(LLMProvider):
|
|||||||
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
on_tool_call_delta: Callable[[dict[str, Any]], Awaitable[None]] | None = None,
|
||||||
) -> LLMResponse:
|
) -> LLMResponse:
|
||||||
await self._ensure_client()
|
await self._ensure_client()
|
||||||
idle_timeout_s = int(os.environ.get("NANOBOT_STREAM_IDLE_TIMEOUT_S", "90"))
|
idle_timeout_s = resolve_stream_idle_timeout_s()
|
||||||
try:
|
try:
|
||||||
if self._should_use_responses_api(model, reasoning_effort):
|
if self._should_use_responses_api(model, reasoning_effort):
|
||||||
try:
|
try:
|
||||||
@ -1503,7 +1504,7 @@ class OpenAICompatProvider(LLMProvider):
|
|||||||
return LLMResponse(
|
return LLMResponse(
|
||||||
content=(
|
content=(
|
||||||
f"Error calling LLM: stream stalled for more than "
|
f"Error calling LLM: stream stalled for more than "
|
||||||
f"{idle_timeout_s} seconds"
|
f"{idle_timeout_s:g} seconds"
|
||||||
),
|
),
|
||||||
finish_reason="error",
|
finish_reason="error",
|
||||||
error_kind="timeout",
|
error_kind="timeout",
|
||||||
|
|||||||
145
tests/providers/test_stream_idle_timeout_config.py
Normal file
145
tests/providers/test_stream_idle_timeout_config.py
Normal file
@ -0,0 +1,145 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from types import SimpleNamespace
|
||||||
|
from typing import Any
|
||||||
|
from unittest.mock import AsyncMock, MagicMock
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
import nanobot.providers.openai_codex_provider as codex_provider
|
||||||
|
from nanobot.providers.anthropic_provider import AnthropicProvider
|
||||||
|
from nanobot.providers.base import (
|
||||||
|
DEFAULT_STREAM_IDLE_TIMEOUT_S,
|
||||||
|
MAX_STREAM_IDLE_TIMEOUT_S,
|
||||||
|
resolve_stream_idle_timeout_s,
|
||||||
|
)
|
||||||
|
from nanobot.providers.bedrock_provider import BedrockProvider
|
||||||
|
from nanobot.providers.openai_compat_provider import OpenAICompatProvider
|
||||||
|
|
||||||
|
|
||||||
|
class _AsyncStream:
|
||||||
|
def __init__(self, chunks: list[Any]) -> None:
|
||||||
|
self._chunks = chunks
|
||||||
|
self._idx = 0
|
||||||
|
|
||||||
|
def __aiter__(self) -> _AsyncStream:
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __anext__(self) -> Any:
|
||||||
|
if self._idx >= len(self._chunks):
|
||||||
|
raise StopAsyncIteration
|
||||||
|
chunk = self._chunks[self._idx]
|
||||||
|
self._idx += 1
|
||||||
|
return chunk
|
||||||
|
|
||||||
|
|
||||||
|
class _AnthropicStream(_AsyncStream):
|
||||||
|
def __init__(self, chunks: list[Any]) -> None:
|
||||||
|
super().__init__(chunks)
|
||||||
|
self.get_final_message = AsyncMock(return_value=SimpleNamespace(
|
||||||
|
content=[SimpleNamespace(type="text", text="ok")],
|
||||||
|
stop_reason="end_turn",
|
||||||
|
usage=SimpleNamespace(input_tokens=1, output_tokens=1),
|
||||||
|
))
|
||||||
|
|
||||||
|
async def __aenter__(self) -> _AnthropicStream:
|
||||||
|
return self
|
||||||
|
|
||||||
|
async def __aexit__(self, *_exc: object) -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class _BedrockClient:
|
||||||
|
def converse_stream(self, **_kwargs: Any) -> dict[str, Any]:
|
||||||
|
return {"stream": iter([
|
||||||
|
{"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "ok"}}},
|
||||||
|
{"messageStop": {"stopReason": "end_turn"}},
|
||||||
|
])}
|
||||||
|
|
||||||
|
|
||||||
|
def test_stream_idle_timeout_parser_rejects_invalid_values() -> None:
|
||||||
|
assert resolve_stream_idle_timeout_s(env_value="abc") == DEFAULT_STREAM_IDLE_TIMEOUT_S
|
||||||
|
assert resolve_stream_idle_timeout_s(env_value="-1") == DEFAULT_STREAM_IDLE_TIMEOUT_S
|
||||||
|
assert resolve_stream_idle_timeout_s(env_value="0") == DEFAULT_STREAM_IDLE_TIMEOUT_S
|
||||||
|
|
||||||
|
|
||||||
|
def test_stream_idle_timeout_parser_accepts_and_clamps_numeric_values() -> None:
|
||||||
|
assert resolve_stream_idle_timeout_s(env_value="1.5") == 1.5
|
||||||
|
assert resolve_stream_idle_timeout_s(env_value="7200") == MAX_STREAM_IDLE_TIMEOUT_S
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_openai_compat_stream_ignores_invalid_idle_timeout_env(monkeypatch) -> None:
|
||||||
|
monkeypatch.setenv("NANOBOT_STREAM_IDLE_TIMEOUT_S", "abc")
|
||||||
|
provider = OpenAICompatProvider(api_key="sk-test", api_base="https://example.com/v1")
|
||||||
|
|
||||||
|
chunk = SimpleNamespace(
|
||||||
|
choices=[SimpleNamespace(
|
||||||
|
delta=SimpleNamespace(
|
||||||
|
content="ok",
|
||||||
|
reasoning_content=None,
|
||||||
|
reasoning=None,
|
||||||
|
tool_calls=None,
|
||||||
|
function_call=None,
|
||||||
|
),
|
||||||
|
finish_reason="stop",
|
||||||
|
)],
|
||||||
|
usage=None,
|
||||||
|
)
|
||||||
|
provider._client = SimpleNamespace(
|
||||||
|
chat=SimpleNamespace(completions=SimpleNamespace(
|
||||||
|
create=AsyncMock(return_value=_AsyncStream([chunk])),
|
||||||
|
)),
|
||||||
|
)
|
||||||
|
|
||||||
|
result = await provider.chat_stream(messages=[{"role": "user", "content": "hi"}])
|
||||||
|
|
||||||
|
assert result.content == "ok"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_anthropic_stream_ignores_invalid_idle_timeout_env(monkeypatch) -> None:
|
||||||
|
monkeypatch.setenv("NANOBOT_STREAM_IDLE_TIMEOUT_S", "abc")
|
||||||
|
provider = AnthropicProvider(api_key="sk-test")
|
||||||
|
provider._client = MagicMock()
|
||||||
|
provider._client.messages.stream = MagicMock(return_value=_AnthropicStream([]))
|
||||||
|
|
||||||
|
result = await provider.chat_stream(messages=[{"role": "user", "content": "hi"}])
|
||||||
|
|
||||||
|
assert result.content == "ok"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_bedrock_stream_ignores_invalid_idle_timeout_env(monkeypatch) -> None:
|
||||||
|
monkeypatch.setenv("NANOBOT_STREAM_IDLE_TIMEOUT_S", "abc")
|
||||||
|
provider = BedrockProvider(region="us-east-1", client=_BedrockClient())
|
||||||
|
|
||||||
|
result = await provider.chat_stream(messages=[{"role": "user", "content": "hi"}])
|
||||||
|
|
||||||
|
assert result.content == "ok"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_codex_stream_ignores_invalid_idle_timeout_env(monkeypatch) -> None:
|
||||||
|
monkeypatch.setenv("NANOBOT_STREAM_IDLE_TIMEOUT_S", "abc")
|
||||||
|
original_client = httpx.AsyncClient
|
||||||
|
seen: dict[str, float] = {}
|
||||||
|
|
||||||
|
def handler(request: httpx.Request) -> httpx.Response:
|
||||||
|
return httpx.Response(200, request=request)
|
||||||
|
|
||||||
|
def fake_client(*, timeout: float, verify: bool) -> httpx.AsyncClient:
|
||||||
|
seen["timeout"] = timeout
|
||||||
|
return original_client(transport=httpx.MockTransport(handler), timeout=timeout)
|
||||||
|
|
||||||
|
monkeypatch.setattr(codex_provider.httpx, "AsyncClient", fake_client)
|
||||||
|
|
||||||
|
await codex_provider._request_codex(
|
||||||
|
"https://codex.example/responses",
|
||||||
|
{},
|
||||||
|
{"input": []},
|
||||||
|
verify=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert seen["timeout"] == DEFAULT_STREAM_IDLE_TIMEOUT_S
|
||||||
Loading…
x
Reference in New Issue
Block a user