wehub-resource-sync
152 lines
4.9 KiB
Python
152 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine import Agent, AgentConfig, ToolResult
|
|
from opensquilla.engine.types import RunHeartbeatEvent, ToolCall, ToolResultEvent
|
|
from opensquilla.provider import ChatConfig, Message, ToolDefinition, ToolInputSchema
|
|
from opensquilla.provider import DoneEvent as ProviderDone
|
|
from opensquilla.provider import ToolUseEndEvent as ProviderToolUseEnd
|
|
from opensquilla.provider import ToolUseStartEvent as ProviderToolUseStart
|
|
|
|
|
|
class _OneToolProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(self, tool_name: str = "slow_tool") -> None:
|
|
self.tool_name = tool_name
|
|
self.calls: list[list[Message]] = []
|
|
|
|
def chat(
|
|
self,
|
|
messages: list[Message],
|
|
tools: list[Any] | None = None,
|
|
config: ChatConfig | None = None,
|
|
) -> AsyncIterator[Any]:
|
|
self.calls.append(messages)
|
|
return self._stream()
|
|
|
|
async def _stream(self) -> AsyncIterator[Any]:
|
|
yield ProviderToolUseStart(tool_use_id="tool-1", tool_name=self.tool_name)
|
|
yield ProviderToolUseEnd(
|
|
tool_use_id="tool-1",
|
|
tool_name=self.tool_name,
|
|
arguments={},
|
|
)
|
|
yield ProviderDone(stop_reason="tool_use", input_tokens=1, output_tokens=1)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
def _tool_def(name: str = "slow_tool") -> ToolDefinition:
|
|
return ToolDefinition(
|
|
name=name,
|
|
description=f"Mock tool {name}",
|
|
input_schema=ToolInputSchema(properties={}, required=[]),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_long_active_tool_emits_run_heartbeat_before_tool_result() -> None:
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
await asyncio.sleep(0.08)
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="ok",
|
|
)
|
|
|
|
agent = Agent(
|
|
provider=_OneToolProvider(),
|
|
config=AgentConfig(
|
|
max_iterations=1,
|
|
tool_timeout=1.0,
|
|
metadata={"tool_activity_heartbeat_interval": 0.02},
|
|
),
|
|
tool_definitions=[_tool_def()],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("run")]
|
|
heartbeat_index = next(
|
|
index for index, event in enumerate(events) if isinstance(event, RunHeartbeatEvent)
|
|
)
|
|
result_index = next(
|
|
index for index, event in enumerate(events) if isinstance(event, ToolResultEvent)
|
|
)
|
|
|
|
heartbeat = events[heartbeat_index]
|
|
assert isinstance(heartbeat, RunHeartbeatEvent)
|
|
assert heartbeat.phase == "tool"
|
|
assert heartbeat_index < result_index
|
|
result = events[result_index]
|
|
assert isinstance(result, ToolResultEvent)
|
|
assert result.result == "ok"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_activity_heartbeat_does_not_extend_tool_timeout() -> None:
|
|
cancelled = asyncio.Event()
|
|
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
try:
|
|
await asyncio.sleep(2.0)
|
|
except asyncio.CancelledError:
|
|
cancelled.set()
|
|
raise
|
|
return ToolResult(
|
|
tool_use_id=tc.tool_use_id,
|
|
tool_name=tc.tool_name,
|
|
content="late",
|
|
)
|
|
|
|
agent = Agent(
|
|
provider=_OneToolProvider(),
|
|
config=AgentConfig(
|
|
max_iterations=1,
|
|
tool_timeout=0.06,
|
|
metadata={"tool_activity_heartbeat_interval": 0.02},
|
|
),
|
|
tool_definitions=[_tool_def()],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("run")]
|
|
result = next(event for event in events if isinstance(event, ToolResultEvent))
|
|
|
|
assert any(isinstance(event, RunHeartbeatEvent) for event in events)
|
|
assert cancelled.is_set()
|
|
assert result.is_error
|
|
assert result.execution_status is not None
|
|
assert result.execution_status["status"] == "timeout"
|
|
assert result.execution_status["reason"] == "runtime_timeout"
|
|
assert result.execution_status["timed_out"] is True
|
|
assert result.result.startswith("Tool 'slow_tool' timed out after ")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_task_cancellation_becomes_tool_error_without_cancelling_turn() -> None:
|
|
async def _handler(tc: ToolCall) -> ToolResult:
|
|
raise asyncio.CancelledError
|
|
|
|
agent = Agent(
|
|
provider=_OneToolProvider(),
|
|
config=AgentConfig(max_iterations=1),
|
|
tool_definitions=[_tool_def()],
|
|
tool_handler=_handler,
|
|
)
|
|
|
|
events = [event async for event in agent.run_turn("run")]
|
|
result = next(event for event in events if isinstance(event, ToolResultEvent))
|
|
|
|
assert result.is_error
|
|
assert result.execution_status is not None
|
|
assert result.execution_status["status"] == "cancelled"
|
|
assert result.execution_status["reason"] == "cancelled"
|
|
assert result.result == "Tool 'slow_tool' was cancelled"
|