264 lines
9.0 KiB
Python
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)]
|