Files
2026-07-13 13:12:33 +08:00

592 lines
21 KiB
Python

from __future__ import annotations
import asyncio
from types import SimpleNamespace
from typing import Any
import pytest
from opensquilla.engine.types import DoneEvent, TextDeltaEvent
from opensquilla.gateway.background_completion import BackgroundCompletionManager
from opensquilla.gateway.boot import GatewayServer
from opensquilla.gateway.routing import ReplyTarget, RouteEnvelope, SourceKind
from opensquilla.session.models import AgentTaskStatus
PARENT = "agent:main:channel:parent"
PARENT_TASK = "task-parent"
class _SessionManager:
def __init__(self, *, parent: Any | None = None, transcript: list[Any] | None = None) -> None:
self.parent = parent
self.transcript = list(transcript or [])
async def read_transcript(self, session_key: str):
return list(self.transcript)
async def get_session(self, session_key: str):
return self.parent
class _Adapter:
def __init__(self, *, fail: bool = False) -> None:
self.fail = fail
self.sent: list[Any] = []
async def send(self, message: Any) -> None:
if self.fail:
raise RuntimeError("channel down")
self.sent.append(message)
class _ChannelManager:
def __init__(self, adapter: _Adapter | None) -> None:
self.adapter = adapter
self.requested_names: list[str] = []
def get(self, channel_name: str):
self.requested_names.append(channel_name)
return self.adapter
class _TaskRuntime:
def __init__(self) -> None:
self.parent_released = asyncio.Event()
self.synthesis_released = asyncio.Event()
self.synthesis_status = AgentTaskStatus.SUCCEEDED
self.sent: list[tuple[str, str, dict[str, Any] | None]] = []
self.stream_event_sink = None
self._tasks: dict[str, Any] = {}
async def send(
self,
session_key: str,
message: str,
provenance: dict[str, Any] | None = None,
stream_event_sink=None,
):
self.sent.append((session_key, message, provenance))
self.stream_event_sink = stream_event_sink
return SimpleNamespace(task_id="task-synthesis")
async def wait(self, task_id: str):
if task_id == PARENT_TASK:
await self.parent_released.wait()
return SimpleNamespace(task_id=task_id, status=AgentTaskStatus.SUCCEEDED)
if task_id == "task-synthesis":
await self.synthesis_released.wait()
return SimpleNamespace(task_id=task_id, status=self.synthesis_status)
raise KeyError(task_id)
async def emit_text_delta(self, text: str) -> None:
assert self.stream_event_sink is not None
await self.stream_event_sink(TextDeltaEvent(text=text))
async def emit_done_text(self, text: str) -> None:
assert self.stream_event_sink is not None
await self.stream_event_sink(DoneEvent(text=text))
async def _wait_until(predicate, *, timeout: float = 1.0) -> None:
deadline = asyncio.get_running_loop().time() + timeout
while not predicate():
if asyncio.get_running_loop().time() >= deadline:
raise AssertionError("condition was not reached")
await asyncio.sleep(0.01)
@pytest.mark.asyncio
async def test_parent_wake_waits_out_of_band_and_delivers_final_channel_text() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
adapter = _Adapter()
session_manager = _SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C123", last_thread_id="T456"),
transcript=[SimpleNamespace(role="assistant", content="yield placeholder")],
)
manager = BackgroundCompletionManager(
session_manager=session_manager,
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(adapter),
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
assert runtime.sent == []
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
await runtime.emit_text_delta("final answer")
runtime.synthesis_released.set()
await _wait_until(lambda: any(event == "session.event.task_group.done" for event, _ in events))
assert adapter.sent[0].content == "final answer"
assert adapter.sent[0].reply_to == "T456"
assert adapter.sent[0].metadata == {"channel": "C123"}
assert [event for event, _ in events] == [
"session.event.task_group.synthesizing",
"session.event.task_group.done",
]
done = events[-1][1]
assert done["group_id"] == f"subagent:{PARENT}:{PARENT_TASK}"
assert done["parent_session_key"] == PARENT
assert done["parent_task_id"] == PARENT_TASK
assert done["synthesis_task_id"] == "task-synthesis"
assert done["delivery_status"] == "sent"
@pytest.mark.asyncio
async def test_synthesis_done_text_delivers_when_no_text_delta_emitted() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
adapter = _Adapter()
session_manager = _SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C123", last_thread_id="T456"),
transcript=[
SimpleNamespace(role="assistant", content="yield placeholder"),
SimpleNamespace(role="assistant", content="unrelated transcript text"),
],
)
manager = BackgroundCompletionManager(
session_manager=session_manager,
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(adapter),
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
await runtime.emit_done_text("done-only final answer")
runtime.synthesis_released.set()
await _wait_until(lambda: any(event == "session.event.task_group.done" for event, _ in events))
assert adapter.sent[0].content == "done-only final answer"
assert events[-1][1]["delivery_status"] == "sent"
@pytest.mark.asyncio
async def test_parent_wake_uses_parent_task_route_when_session_route_changes() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
runtime._tasks[PARENT_TASK] = SimpleNamespace(
envelope=RouteEnvelope(
source_kind=SourceKind.CHANNEL,
source_name="slack",
agent_id="main",
session_key=PARENT,
channel_name="slack",
channel_id="C-old",
thread_id="T-old",
reply_target=ReplyTarget(
kind="channel",
channel_name="slack",
to="C-old",
thread_id="T-old",
),
)
)
adapter = _Adapter()
channel_manager = _ChannelManager(adapter)
session_manager = _SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C-new", last_thread_id="T-new"),
)
manager = BackgroundCompletionManager(
session_manager=session_manager,
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: channel_manager,
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
await runtime.emit_text_delta("final answer")
runtime.synthesis_released.set()
await _wait_until(lambda: any(event == "session.event.task_group.done" for event, _ in events))
assert channel_manager.requested_names == ["slack"]
assert adapter.sent[0].content == "final answer"
assert adapter.sent[0].reply_to == "T-old"
assert adapter.sent[0].metadata == {"channel": "C-old"}
@pytest.mark.asyncio
async def test_parent_wake_freezes_session_route_before_synthesis_finishes() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
adapter = _Adapter()
session_manager = _SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C-old", last_thread_id="T-old"),
)
manager = BackgroundCompletionManager(
session_manager=session_manager,
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(adapter),
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
session_manager.parent.last_to = "C-new"
session_manager.parent.last_thread_id = "T-new"
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
await runtime.emit_text_delta("final answer")
runtime.synthesis_released.set()
await _wait_until(lambda: any(event == "session.event.task_group.done" for event, _ in events))
assert adapter.sent[0].reply_to == "T-old"
assert adapter.sent[0].metadata == {"channel": "C-old"}
@pytest.mark.asyncio
async def test_parent_wake_uses_target_captured_before_parent_task_eviction() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
runtime._tasks[PARENT_TASK] = SimpleNamespace(
envelope=RouteEnvelope(
source_kind=SourceKind.CHANNEL,
source_name="slack",
agent_id="main",
session_key=PARENT,
channel_name="slack",
channel_id="C-old",
thread_id="T-old",
reply_target=ReplyTarget(
kind="channel",
channel_name="slack",
to="C-old",
thread_id="T-old",
),
)
)
adapter = _Adapter()
session_manager = _SessionManager(
parent=SimpleNamespace(
last_channel="slack",
last_to="C-original",
last_thread_id="T-original",
),
)
manager = BackgroundCompletionManager(
session_manager=session_manager,
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(adapter),
)
await manager.capture_delivery_target(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
task_runtime=runtime,
)
runtime._tasks.clear()
session_manager.parent.last_to = "C-new"
session_manager.parent.last_thread_id = "T-new"
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
await runtime.emit_text_delta("final answer")
runtime.synthesis_released.set()
await _wait_until(lambda: any(event == "session.event.task_group.done" for event, _ in events))
assert adapter.sent[0].reply_to == "T-old"
assert adapter.sent[0].metadata == {"channel": "C-old"}
@pytest.mark.asyncio
async def test_unrelated_post_watermark_assistant_text_is_not_applicable() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
adapter = _Adapter()
session_manager = _SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C123", last_thread_id="T456"),
transcript=[SimpleNamespace(role="assistant", content="yield placeholder")],
)
manager = BackgroundCompletionManager(
session_manager=session_manager,
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(adapter),
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
session_manager.transcript.append(SimpleNamespace(role="assistant", content="unrelated answer"))
runtime.synthesis_released.set()
await _wait_until(lambda: any(event == "session.event.task_group.done" for event, _ in events))
assert adapter.sent == []
assert events[-1][1]["delivery_status"] == "not_applicable"
@pytest.mark.asyncio
async def test_channel_delivery_failure_is_reported_without_failing_group() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
session_manager = _SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C123", last_thread_id="T456"),
transcript=[SimpleNamespace(role="assistant", content="yield placeholder")],
)
manager = BackgroundCompletionManager(
session_manager=session_manager,
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(_Adapter(fail=True)),
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
await runtime.emit_text_delta("final answer")
runtime.synthesis_released.set()
await _wait_until(lambda: any(event == "session.event.task_group.done" for event, _ in events))
assert events[-1][0] == "session.event.task_group.done"
assert events[-1][1]["delivery_status"] == "failed"
assert events[-1][1]["delivery_error_class"] == "RuntimeError"
@pytest.mark.asyncio
async def test_synthesis_failure_emits_group_failed() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
runtime.synthesis_status = AgentTaskStatus.FAILED
session_manager = _SessionManager(parent=SimpleNamespace(last_channel=None, last_to=None))
manager = BackgroundCompletionManager(
session_manager=session_manager,
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(None),
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
runtime.synthesis_released.set()
await _wait_until(
lambda: any(event == "session.event.task_group.failed" for event, _ in events)
)
failed = events[-1][1]
assert failed["status"] == "failed"
assert failed["synthesis_status"] == "failed"
assert failed["delivery_status"] == "not_applicable"
@pytest.mark.asyncio
async def test_waiting_event_is_not_reemitted_after_wake_starts() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
manager = BackgroundCompletionManager(
session_manager=_SessionManager(parent=SimpleNamespace(last_channel=None, last_to=None)),
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(None),
)
await manager.emit_waiting(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
pending_count=0,
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
await manager.emit_waiting(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
pending_count=0,
)
assert [event for event, _ in events] == ["session.event.task_group.waiting"]
@pytest.mark.asyncio
async def test_background_completion_drain_waits_for_detached_watcher() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
adapter = _Adapter()
manager = BackgroundCompletionManager(
session_manager=_SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C123", last_thread_id="T456")
),
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(adapter),
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
drain_task = asyncio.create_task(manager.drain(timeout=1.0))
await asyncio.sleep(0)
assert not drain_task.done()
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
await runtime.emit_text_delta("final answer")
runtime.synthesis_released.set()
await drain_task
assert adapter.sent[0].content == "final answer"
assert events[-1][0] == "session.event.task_group.done"
@pytest.mark.asyncio
async def test_background_completion_drain_timeout_does_not_cancel_watcher() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
adapter = _Adapter()
manager = BackgroundCompletionManager(
session_manager=_SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C123", last_thread_id="T456")
),
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(adapter),
)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
await manager.drain(timeout=0.01)
runtime.parent_released.set()
await _wait_until(lambda: len(runtime.sent) == 1)
await runtime.emit_text_delta("final answer")
runtime.synthesis_released.set()
await _wait_until(lambda: any(event == "session.event.task_group.done" for event, _ in events))
assert adapter.sent[0].content == "final answer"
@pytest.mark.asyncio
async def test_background_completion_close_rejects_new_wake_registration() -> None:
events: list[tuple[str, dict[str, Any]]] = []
runtime = _TaskRuntime()
manager = BackgroundCompletionManager(
session_manager=_SessionManager(
parent=SimpleNamespace(last_channel="slack", last_to="C123", last_thread_id="T456")
),
event_emitter=lambda _session, event, payload: _record(events, event, payload),
channel_manager_ref=lambda: _ChannelManager(_Adapter()),
)
await manager.close(timeout=0.1)
await manager.send_parent_wake(
parent_session_key=PARENT,
parent_task_id=PARENT_TASK,
payloads=[{"child_session_key": "child"}],
task_runtime=runtime,
message="wake",
provenance={"kind": "internal_system"},
)
await asyncio.sleep(0)
assert runtime.sent == []
assert events == []
@pytest.mark.asyncio
async def test_gateway_close_drains_background_completion_before_stopping_channels() -> None:
order: list[str] = []
class _Runtime:
async def shutdown(self, **_kwargs) -> None:
order.append("runtime")
class _Services:
task_runtime = _Runtime()
async def close(self) -> None:
order.append("services")
class _Background:
async def close(self, **_kwargs) -> None:
order.append("background")
class _Channels:
async def stop_all(self) -> None:
order.append("channels")
server = GatewayServer(app=SimpleNamespace(), config=SimpleNamespace())
server._services = _Services()
server._background_completion_manager = _Background()
server._channel_manager = _Channels()
await server.close()
assert order.index("runtime") < order.index("background") < order.index("channels")
assert order[-1] == "services"
async def _record(
events: list[tuple[str, dict[str, Any]]],
event: str,
payload: dict[str, Any],
) -> None:
events.append((event, payload))