mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-09-04 02:01:48 +03:00
fix(agent): observe session reply timeout task failures
This commit is contained in:
@@ -13,6 +13,8 @@ from dataclasses import dataclass
|
|||||||
from typing import Any, Protocol
|
from typing import Any, Protocol
|
||||||
from uuid import uuid4
|
from uuid import uuid4
|
||||||
|
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
from nanobot.agent.tools.base import Tool, ToolResult, tool_parameters
|
from nanobot.agent.tools.base import Tool, ToolResult, tool_parameters
|
||||||
from nanobot.agent.tools.context import RequestContext, ToolContext, current_request_context
|
from nanobot.agent.tools.context import RequestContext, ToolContext, current_request_context
|
||||||
from nanobot.agent.tools.schema import (
|
from nanobot.agent.tools.schema import (
|
||||||
@@ -325,11 +327,20 @@ class SendSessionMessageTool(Tool):
|
|||||||
def expire() -> None:
|
def expire() -> None:
|
||||||
task = asyncio.create_task(self._expire_pending_reply(key, pending))
|
task = asyncio.create_task(self._expire_pending_reply(key, pending))
|
||||||
self._expiry_tasks.add(task)
|
self._expiry_tasks.add(task)
|
||||||
task.add_done_callback(self._expiry_tasks.discard)
|
task.add_done_callback(self._on_expiry_task_done)
|
||||||
|
|
||||||
schedule = self._schedule_later or asyncio.get_running_loop().call_later
|
schedule = self._schedule_later or asyncio.get_running_loop().call_later
|
||||||
pending.timer = schedule(float(timeout_seconds), expire)
|
pending.timer = schedule(float(timeout_seconds), expire)
|
||||||
|
|
||||||
|
def _on_expiry_task_done(self, task: asyncio.Task[None]) -> None:
|
||||||
|
self._expiry_tasks.discard(task)
|
||||||
|
if task.cancelled():
|
||||||
|
return
|
||||||
|
try:
|
||||||
|
task.result()
|
||||||
|
except Exception:
|
||||||
|
logger.exception("Session reply timeout delivery failed")
|
||||||
|
|
||||||
async def _expire_pending_reply(
|
async def _expire_pending_reply(
|
||||||
self,
|
self,
|
||||||
key: tuple[str, str],
|
key: tuple[str, str],
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ import asyncio
|
|||||||
import json
|
import json
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Callable
|
from typing import Callable
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
@@ -277,6 +278,45 @@ async def test_reply_timeout_injects_a_user_input_back_into_the_source(
|
|||||||
assert timeout.content == f"No reply from @{target.name} after 5 seconds."
|
assert timeout.content == f"No reply from @{target.name} after 5 seconds."
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_reply_timeout_observes_background_delivery_failure(
|
||||||
|
tmp_path: Path,
|
||||||
|
) -> None:
|
||||||
|
sessions = SessionManager(tmp_path)
|
||||||
|
_persist(sessions, "websocket:source", "websocket:target")
|
||||||
|
bus = MessageBus()
|
||||||
|
scheduler = _Scheduler()
|
||||||
|
tool = SendSessionMessageTool(
|
||||||
|
sessions=sessions,
|
||||||
|
bus=bus,
|
||||||
|
schedule_later=scheduler,
|
||||||
|
)
|
||||||
|
target = _handle(sessions, "websocket:target")
|
||||||
|
|
||||||
|
await tool.enqueue(
|
||||||
|
source_session_key="websocket:source",
|
||||||
|
target_handle=target.name,
|
||||||
|
content="Question",
|
||||||
|
expect_reply=True,
|
||||||
|
reply_timeout_seconds=5,
|
||||||
|
)
|
||||||
|
await bus.consume_inbound()
|
||||||
|
|
||||||
|
async def fail_publish(_message) -> None:
|
||||||
|
raise RuntimeError("queue unavailable")
|
||||||
|
|
||||||
|
bus.publish_inbound = fail_publish
|
||||||
|
with patch("nanobot.agent.tools.session_messages.logger") as logger:
|
||||||
|
scheduler.calls[0][1].fire()
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
await asyncio.sleep(0)
|
||||||
|
|
||||||
|
assert tool._expiry_tasks == set()
|
||||||
|
logger.exception.assert_called_once_with(
|
||||||
|
"Session reply timeout delivery failed",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_reverse_message_cancels_the_pending_reply_timeout(
|
async def test_reverse_message_cancels_the_pending_reply_timeout(
|
||||||
tmp_path: Path,
|
tmp_path: Path,
|
||||||
|
|||||||
Reference in New Issue
Block a user