Files
opensquilla--opensquilla/tests/test_engine/test_runtime_agent_iteration_timeout.py
2026-07-13 13:12:33 +08:00

264 lines
9.0 KiB
Python

from __future__ import annotations
import asyncio
from collections.abc import AsyncIterator
from types import SimpleNamespace
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from opensquilla.engine.agent import Agent, _IterationStreamTimeoutError
from opensquilla.engine.runtime import TurnRunner
from opensquilla.engine.types import AgentConfig, DoneEvent
from opensquilla.gateway.config import GatewayConfig
class _SessionConfigManager:
def __init__(self, config: object | None) -> None:
self.config = config
def get_session_config(self, session_key: str) -> object | None:
return self.config
def test_resolve_agent_iteration_timeout_prefers_explicit_value(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("OPENSQUILLA_AGENT_ITERATION_TIMEOUT", "222")
runner = TurnRunner(
provider_selector=None,
session_manager=_SessionConfigManager(
SimpleNamespace(agent_iteration_timeout_seconds=111.0)
),
config=GatewayConfig(agent_iteration_timeout_seconds=333.0),
)
assert runner._resolve_agent_iteration_timeout("agent:main:test", 444.0) == 444.0
def test_resolve_agent_iteration_timeout_prefers_session_config(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("OPENSQUILLA_AGENT_ITERATION_TIMEOUT", "222")
runner = TurnRunner(
provider_selector=None,
session_manager=_SessionConfigManager(
SimpleNamespace(agent_iteration_timeout_seconds=111.0)
),
config=GatewayConfig(agent_iteration_timeout_seconds=333.0),
)
assert runner._resolve_agent_iteration_timeout("agent:main:test") == 111.0
def test_resolve_agent_iteration_timeout_prefers_env_over_gateway_config(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("OPENSQUILLA_AGENT_ITERATION_TIMEOUT", "222")
runner = TurnRunner(
provider_selector=None,
session_manager=_SessionConfigManager(None),
config=GatewayConfig(agent_iteration_timeout_seconds=333.0),
)
assert runner._resolve_agent_iteration_timeout("agent:main:test") == 222.0
def test_resolve_agent_iteration_timeout_uses_gateway_config(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("OPENSQUILLA_AGENT_ITERATION_TIMEOUT", raising=False)
runner = TurnRunner(
provider_selector=None,
config=GatewayConfig(agent_iteration_timeout_seconds=333.0),
)
assert runner._resolve_agent_iteration_timeout("agent:main:test") == 333.0
def test_resolve_agent_iteration_timeout_uses_agent_default_without_config(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("OPENSQUILLA_AGENT_ITERATION_TIMEOUT", raising=False)
runner = TurnRunner(provider_selector=None, config=None)
assert (
runner._resolve_agent_iteration_timeout("agent:main:test")
== AgentConfig().iteration_timeout
)
def test_resolve_agent_iteration_timeout_invalid_env_falls_through(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("OPENSQUILLA_AGENT_ITERATION_TIMEOUT", "not-a-float")
runner = TurnRunner(
provider_selector=None,
session_manager=_SessionConfigManager(None),
config=GatewayConfig(agent_iteration_timeout_seconds=333.0),
)
assert runner._resolve_agent_iteration_timeout("agent:main:test") == 333.0
def test_resolve_agent_iteration_timeout_floors_to_5400_in_coding_mode(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("OPENSQUILLA_AGENT_ITERATION_TIMEOUT", raising=False)
cfg = GatewayConfig(agent_iteration_timeout_seconds=600.0)
cfg.skills.coding_mode = True # coding mode waits on code-task up to 90 min
runner = TurnRunner(provider_selector=None, config=cfg)
# the per-iteration watchdog is floored so a long process(wait) is not clamped
assert runner._resolve_agent_iteration_timeout("agent:main:test") == 5400.0
def test_resolve_agent_iteration_timeout_no_floor_when_coding_off(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("OPENSQUILLA_AGENT_ITERATION_TIMEOUT", raising=False)
cfg = GatewayConfig(agent_iteration_timeout_seconds=600.0)
runner = TurnRunner(provider_selector=None, config=cfg)
# coding mode OFF -> the configured small value is preserved (no floor)
assert runner._resolve_agent_iteration_timeout("agent:main:test") == 600.0
def test_resolve_agent_iteration_timeout_rejects_invalid_explicit_value() -> None:
runner = TurnRunner(provider_selector=None, config=GatewayConfig())
with pytest.raises(ValueError, match="iteration_timeout"):
runner._resolve_agent_iteration_timeout("agent:main:test", -1.0)
@pytest.mark.asyncio
async def test_run_threads_iteration_timeout_into_agent_config(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Regression: runner.run(iteration_timeout=X) must reach AgentConfig.
iteration_timeout was previously declared on TurnRunner.run() and
referenced inside _run_turn() at the resolver call site, but never
plumbed through _run_turn()'s signature or the two run() -> _run_turn()
call sites. Every turn would hit NameError before reaching the resolver.
The existing isolation tests above exercise the resolver directly and
so would not have caught the threading gap.
"""
from opensquilla.tools.types import ToolContext
seen_kwargs: list[dict[str, Any]] = []
real_agent_config = AgentConfig
def recording_agent_config(**kwargs: Any) -> AgentConfig:
seen_kwargs.append(kwargs)
return real_agent_config(**kwargs)
monkeypatch.setattr("opensquilla.engine.types.AgentConfig", recording_agent_config)
provider = MagicMock()
provider.provider_name = "stub"
async def _chat(*args: Any, **kwargs: Any) -> AsyncIterator[Any]:
yield DoneEvent()
provider.chat = _chat
selector = MagicMock()
selector.resolve.return_value = provider
selector.clone.return_value = selector
selector.current_config = MagicMock(model="stub-model")
session_manager = MagicMock()
session_manager.get = AsyncMock(return_value=None)
session_manager.append_message = AsyncMock(return_value=None)
session_manager.update = AsyncMock(return_value=None)
session_manager.get_compaction_summary = AsyncMock(return_value=None)
session_manager.get_transcript = AsyncMock(return_value=[])
runner = TurnRunner(
provider_selector=selector,
session_manager=session_manager,
)
tool_ctx = ToolContext(session_key="agent:main:iter-thread-test")
async for _ in runner.run(
message="hi",
session_key="agent:main:iter-thread-test",
tool_context=tool_ctx,
iteration_timeout=444.0,
):
pass
assert any(kw.get("iteration_timeout") == 444.0 for kw in seen_kwargs), (
f"AgentConfig never received iteration_timeout=444.0; saw {seen_kwargs!r}"
)
@pytest.mark.asyncio
async def test_stream_iteration_timeout_does_not_double_close_provider_stream(
monkeypatch: pytest.MonkeyPatch,
) -> None:
agent = Agent.__new__(Agent)
agent.config = MagicMock(timeout=1.0, iteration_timeout=0.01)
close_calls = 0
async def provider_stream() -> AsyncIterator[dict[str, str]]:
try:
await asyncio.sleep(1.0)
yield {"type": "chunk", "data": "late"}
finally:
await asyncio.sleep(0)
async def record_close(_stream_iter: AsyncIterator[Any]) -> None:
nonlocal close_calls
close_calls += 1
monkeypatch.setattr(agent, "_close_provider_stream", record_close)
loop = asyncio.get_running_loop()
with pytest.raises(_IterationStreamTimeoutError):
async for _event in agent._stream_provider_events_with_deadline(
provider_stream(),
loop=loop,
total_deadline=None,
):
pass
assert close_calls == 0
@pytest.mark.asyncio
async def test_total_deadline_limited_wait_stays_total_timeout_on_early_wake(
monkeypatch: pytest.MonkeyPatch,
) -> None:
agent = Agent.__new__(Agent)
agent.config = SimpleNamespace(timeout=0.05, iteration_timeout=30.0)
observed_timeouts: list[float | None] = []
async def provider_stream() -> AsyncIterator[dict[str, str]]:
await asyncio.sleep(30.0)
yield {"type": "chunk", "data": "late"}
async def return_before_deadline(
futures: set[asyncio.Future[Any]],
*,
timeout: float | None = None,
) -> tuple[set[asyncio.Future[Any]], set[asyncio.Future[Any]]]:
observed_timeouts.append(timeout)
return set(), futures
monkeypatch.setattr(asyncio, "wait", return_before_deadline)
fake_loop: Any = SimpleNamespace(time=lambda: 10.0)
with pytest.raises(TimeoutError, match="total timeout") as exc_info:
async for _event in agent._stream_provider_events_with_deadline(
provider_stream(),
loop=fake_loop,
total_deadline=10.05,
):
pass
assert type(exc_info.value) is TimeoutError
assert observed_timeouts == [pytest.approx(0.05)]