mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-08-08 05:18:49 +03:00
fix(webui): reconcile threads after browser resume
This commit is contained in:
@@ -7,8 +7,11 @@ from nanobot.agent.loop import AgentLoop
|
||||
from nanobot.bus.events import OutboundMessage
|
||||
from nanobot.bus.outbound_events import GoalStatusEvent
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.channels.websocket.runtime import WebSocketChannel
|
||||
from nanobot.providers.base import GenerationSettings, LLMResponse
|
||||
from nanobot.session.webui_turns import WebuiTurnCoordinator
|
||||
from nanobot.session import webui_turns as wth
|
||||
from nanobot.session.webui_turns import WebuiTurnCoordinator, WebuiTurnRoutePolicy
|
||||
from nanobot.webui.metadata import WEBSOCKET_TURN_OWNER_METADATA_KEY
|
||||
|
||||
|
||||
def _make_loop(tmp_path):
|
||||
@@ -32,6 +35,7 @@ def _make_loop(tmp_path):
|
||||
sessions=loop.sessions,
|
||||
schedule_background=lambda coro: loop._schedule_background(coro),
|
||||
).subscribe(loop.runtime_events)
|
||||
loop.turn_delivery_factory.route_policy = WebuiTurnRoutePolicy(loop.sessions)
|
||||
loop.tools.get_definitions = MagicMock(return_value=[])
|
||||
return loop
|
||||
|
||||
@@ -39,29 +43,51 @@ def _make_loop(tmp_path):
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_direct_websocket_clears_run_status(tmp_path) -> None:
|
||||
loop = _make_loop(tmp_path)
|
||||
|
||||
response = await loop.process_direct(
|
||||
"deliver reminder",
|
||||
session_key="cron:reminder-1",
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
gateway = MagicMock()
|
||||
channel = WebSocketChannel(
|
||||
{"enabled": True, "allowFrom": ["*"]},
|
||||
loop.bus,
|
||||
gateway=gateway,
|
||||
)
|
||||
|
||||
assert response is not None
|
||||
assert response.content == "done"
|
||||
try:
|
||||
response = await loop.process_direct(
|
||||
"deliver reminder",
|
||||
session_key="cron:reminder-1",
|
||||
channel="websocket",
|
||||
chat_id="chat-1",
|
||||
)
|
||||
|
||||
events = []
|
||||
while loop.bus.outbound_size:
|
||||
events.append(await loop.bus.consume_outbound())
|
||||
assert response is not None
|
||||
assert response.content == "done"
|
||||
|
||||
statuses = [
|
||||
event.event
|
||||
for event in events
|
||||
if isinstance(event.event, GoalStatusEvent)
|
||||
]
|
||||
assert [status.status for status in statuses] == ["running", "idle"]
|
||||
assert isinstance(statuses[0].started_at, float)
|
||||
assert statuses[1].started_at is None
|
||||
events = []
|
||||
while loop.bus.outbound_size:
|
||||
event = await loop.bus.consume_outbound()
|
||||
events.append(event)
|
||||
await channel.send(event)
|
||||
|
||||
status_messages = [
|
||||
event
|
||||
for event in events
|
||||
if isinstance(event.event, GoalStatusEvent)
|
||||
]
|
||||
statuses = [event.event for event in status_messages]
|
||||
assert [status.status for status in statuses] == ["running", "idle"]
|
||||
assert isinstance(statuses[0].started_at, float)
|
||||
assert statuses[1].started_at is None
|
||||
owners = {
|
||||
event.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
|
||||
for event in status_messages
|
||||
}
|
||||
assert len(owners) == 1
|
||||
assert wth.websocket_turn_wall_started_at("chat-1") is None
|
||||
assert "chat-1" not in wth._WEBSOCKET_ACTIVE_TURNS
|
||||
finally:
|
||||
wth._WEBSOCKET_ACTIVE_TURNS.clear()
|
||||
wth._WEBSOCKET_TURN_WALL_STARTED_AT.clear()
|
||||
wth._WEBSOCKET_TURN_IDS.clear()
|
||||
wth._WEBSOCKET_TURN_OWNERS.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -28,7 +28,10 @@ from nanobot.utils.progress_events import (
|
||||
invoke_file_edit_progress,
|
||||
on_progress_accepts_file_edit_events,
|
||||
)
|
||||
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
|
||||
from nanobot.webui.metadata import (
|
||||
WEBSOCKET_TURN_OWNER_METADATA_KEY,
|
||||
WEBUI_TURN_METADATA_KEY,
|
||||
)
|
||||
|
||||
|
||||
def _make_loop(tmp_path: Path) -> AgentLoop:
|
||||
@@ -903,6 +906,12 @@ class TestToolEventProgress:
|
||||
turn_id = turn_ids.pop()
|
||||
assert isinstance(turn_id, str)
|
||||
assert turn_id.startswith("subagent:")
|
||||
owners = {
|
||||
message.metadata.get(WEBSOCKET_TURN_OWNER_METADATA_KEY)
|
||||
for message in visible_events
|
||||
}
|
||||
assert len(owners) == 1
|
||||
assert isinstance(owners.pop(), str)
|
||||
assert all(
|
||||
(message.channel, message.chat_id) == ("websocket", "chat-a")
|
||||
and message.metadata.get("webui") is True
|
||||
@@ -910,6 +919,7 @@ class TestToolEventProgress:
|
||||
and set(message.metadata) <= {
|
||||
"webui",
|
||||
"_wants_stream",
|
||||
WEBSOCKET_TURN_OWNER_METADATA_KEY,
|
||||
WEBUI_TURN_METADATA_KEY,
|
||||
"latency_ms",
|
||||
}
|
||||
|
||||
@@ -1,12 +1,147 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from nanobot.agent.turn_delivery import TurnDeliveryFactory
|
||||
from nanobot.bus.events import InboundMessage
|
||||
from nanobot.bus.queue import MessageBus
|
||||
from nanobot.bus.runtime_events import RuntimeEventBus
|
||||
from nanobot.session.manager import SessionManager
|
||||
from nanobot.session.webui_turns import WebuiTurnRoutePolicy
|
||||
from nanobot.webui.metadata import WEBUI_TURN_METADATA_KEY
|
||||
from nanobot.webui.metadata import (
|
||||
WEBSOCKET_TURN_OWNER_METADATA_KEY,
|
||||
WEBUI_TURN_METADATA_KEY,
|
||||
)
|
||||
|
||||
|
||||
def test_websocket_lifecycles_get_distinct_internal_owners(tmp_path: Path) -> None:
|
||||
factory = TurnDeliveryFactory(
|
||||
MessageBus(),
|
||||
RuntimeEventBus(),
|
||||
route_policy=WebuiTurnRoutePolicy(SessionManager(tmp_path / "sessions")),
|
||||
)
|
||||
first_msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user",
|
||||
chat_id="chat-a",
|
||||
content="first",
|
||||
metadata={WEBSOCKET_TURN_OWNER_METADATA_KEY: "attacker-reused-owner"},
|
||||
)
|
||||
second_msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user",
|
||||
chat_id="chat-a",
|
||||
content="second",
|
||||
metadata={WEBSOCKET_TURN_OWNER_METADATA_KEY: "attacker-reused-owner"},
|
||||
)
|
||||
|
||||
first = factory.create(first_msg, first_msg.session_key)
|
||||
second = factory.create(second_msg, second_msg.session_key)
|
||||
first_owner = first.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
|
||||
second_owner = second.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
|
||||
|
||||
assert first_owner == first.delivery_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
|
||||
assert first_owner != second_owner
|
||||
assert first_owner != "attacker-reused-owner"
|
||||
assert second_owner != "attacker-reused-owner"
|
||||
assert WEBUI_TURN_METADATA_KEY not in first.lifecycle_message.metadata
|
||||
assert first_msg.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] == first_owner
|
||||
assert second_msg.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] == second_owner
|
||||
|
||||
|
||||
def test_websocket_lifecycle_reuses_registered_ingress_owner(tmp_path: Path) -> None:
|
||||
from nanobot.session import webui_turns as wth
|
||||
|
||||
owner = wth.register_queued_websocket_turn_if_idle("chat-queued", "turn-queued")
|
||||
assert owner is not None
|
||||
msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user",
|
||||
chat_id="chat-queued",
|
||||
content="queued",
|
||||
metadata={
|
||||
WEBSOCKET_TURN_OWNER_METADATA_KEY: owner,
|
||||
WEBUI_TURN_METADATA_KEY: "turn-queued",
|
||||
},
|
||||
)
|
||||
factory = TurnDeliveryFactory(
|
||||
MessageBus(),
|
||||
RuntimeEventBus(),
|
||||
route_policy=WebuiTurnRoutePolicy(SessionManager(tmp_path / "sessions")),
|
||||
)
|
||||
|
||||
try:
|
||||
delivery = factory.create(msg, msg.session_key)
|
||||
|
||||
assert delivery.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] == owner
|
||||
assert msg.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY] == owner
|
||||
finally:
|
||||
wth.clear_websocket_turn_if_current("chat-queued", owner)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_chat_different_sessions_restore_previous_active_projection(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from nanobot.session import webui_turns as wth
|
||||
|
||||
factory = TurnDeliveryFactory(
|
||||
MessageBus(),
|
||||
RuntimeEventBus(),
|
||||
route_policy=WebuiTurnRoutePolicy(SessionManager(tmp_path / "sessions")),
|
||||
)
|
||||
first_msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user",
|
||||
chat_id="shared-chat",
|
||||
content="first",
|
||||
metadata={WEBUI_TURN_METADATA_KEY: "turn-first"},
|
||||
session_key_override="websocket:session-first",
|
||||
)
|
||||
second_msg = InboundMessage(
|
||||
channel="websocket",
|
||||
sender_id="user",
|
||||
chat_id="shared-chat",
|
||||
content="second",
|
||||
metadata={WEBUI_TURN_METADATA_KEY: "turn-second"},
|
||||
session_key_override="websocket:session-second",
|
||||
)
|
||||
first = factory.create(first_msg, first_msg.session_key)
|
||||
second = factory.create(second_msg, second_msg.session_key)
|
||||
first_owner = first.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
|
||||
second_owner = second.lifecycle_message.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
|
||||
bus = MagicMock()
|
||||
bus.publish_outbound = AsyncMock()
|
||||
|
||||
try:
|
||||
await wth.publish_turn_run_status(
|
||||
bus,
|
||||
first.lifecycle_message,
|
||||
"running",
|
||||
started_at=100.0,
|
||||
)
|
||||
await wth.publish_turn_run_status(
|
||||
bus,
|
||||
second.lifecycle_message,
|
||||
"running",
|
||||
started_at=200.0,
|
||||
)
|
||||
|
||||
assert wth.websocket_turn_wall_started_at("shared-chat") == 200.0
|
||||
assert wth.websocket_turn_id("shared-chat") == "turn-second"
|
||||
assert wth.clear_websocket_turn_if_current("shared-chat", second_owner) is True
|
||||
assert wth.websocket_turn_wall_started_at("shared-chat") == 100.0
|
||||
assert wth.websocket_turn_id("shared-chat") == "turn-first"
|
||||
assert wth._WEBSOCKET_TURN_OWNERS["shared-chat"] == first_owner
|
||||
assert wth.clear_websocket_turn_if_current("shared-chat", first_owner) is True
|
||||
assert wth.websocket_turn_wall_started_at("shared-chat") is None
|
||||
finally:
|
||||
wth._WEBSOCKET_ACTIVE_TURNS.pop("shared-chat", None)
|
||||
wth._WEBSOCKET_TURN_WALL_STARTED_AT.pop("shared-chat", None)
|
||||
wth._WEBSOCKET_TURN_IDS.pop("shared-chat", None)
|
||||
wth._WEBSOCKET_TURN_OWNERS.pop("shared-chat", None)
|
||||
|
||||
|
||||
def test_late_subagent_route_requires_webui_owned_session(tmp_path: Path) -> None:
|
||||
@@ -45,6 +180,7 @@ def test_late_subagent_route_requires_webui_owned_session(tmp_path: Path) -> Non
|
||||
assert set(first_visible_route.metadata) == {
|
||||
"webui",
|
||||
"_wants_stream",
|
||||
WEBSOCKET_TURN_OWNER_METADATA_KEY,
|
||||
WEBUI_TURN_METADATA_KEY,
|
||||
}
|
||||
assert first_visible_route.metadata["webui"] is True
|
||||
@@ -54,6 +190,10 @@ def test_late_subagent_route_requires_webui_owned_session(tmp_path: Path) -> Non
|
||||
assert first_turn_id.startswith("subagent:")
|
||||
assert second_turn_id.startswith("subagent:")
|
||||
assert first_turn_id != second_turn_id
|
||||
assert (
|
||||
first_visible_route.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
|
||||
!= second_visible_route.metadata[WEBSOCKET_TURN_OWNER_METADATA_KEY]
|
||||
)
|
||||
assert msg.metadata == {
|
||||
"injected_event": "subagent_result",
|
||||
"subagent_task_id": "sub-1",
|
||||
|
||||
Reference in New Issue
Block a user