wehub-resource-sync
1329 lines
51 KiB
Python
1329 lines
51 KiB
Python
"""
|
||
Tunnel three-layer integration test: Omnigent → WS tunnel → runner → harness.
|
||
|
||
Production-shaped variant of ``test_sessions_three_layer.py``. The
|
||
prior test wired the AP's ``_runner_client`` directly to the runner
|
||
app via ``httpx.ASGITransport``, which short-circuits the WebSocket
|
||
tunnel that ships in production. This test instead drives the real
|
||
``/v1/runners/{runner_id}/tunnel`` route with an
|
||
``asgiref.testing.ApplicationCommunicator`` and pumps
|
||
``RequestFrame`` traffic through ``dispatch_via_asgi`` into a real
|
||
``create_runner_app``. The AP's ``RunnerRouter`` ends up serving a
|
||
``WSTunnelTransport``-backed client (the same code path the real
|
||
server uses), so every native-sessions HTTP hop crosses the live
|
||
tunnel registry.
|
||
|
||
Catches bugs that the ASGI-only three-layer test misses: frame
|
||
encode/decode regressions, tunnel registry leak/registration races,
|
||
``WSTunnelTransport`` error mapping, and any AP-side code that
|
||
reaches for ``app.state.tunnel_registry`` / ``get_runner_router()``.
|
||
|
||
Layers:
|
||
Omnigent server ──HTTP──> RunnerRouter ──WSTunnelTransport──>
|
||
TunnelRegistry ──tunnel WS route──> ApplicationCommunicator ──>
|
||
forwarder task ──> create_runner_app ──> FakeProcessManager ──>
|
||
EchoHarness ASGI app
|
||
|
||
The TODO at the end of the module flags ``FakeProcessManager`` and
|
||
``_build_harness_agent_bundle`` as candidates to factor out of
|
||
``test_sessions_three_layer.py`` once a second consumer (this file)
|
||
lands.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import contextlib
|
||
import io
|
||
import json
|
||
import tarfile
|
||
from collections.abc import AsyncIterator, Iterator
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import httpx
|
||
import pytest
|
||
import pytest_asyncio
|
||
import yaml
|
||
from asgiref.testing import ApplicationCommunicator
|
||
from fastapi import FastAPI
|
||
from omnigent_client._events import ResponseCompleted, ResponseFailed
|
||
from omnigent_client._files import FilesNamespace
|
||
from omnigent_client._sessions import SessionsNamespace
|
||
|
||
from omnigent.repl._repl import _SessionsChatReplAdapter
|
||
from omnigent.runner.app import create_runner_app
|
||
from omnigent.runner.transports.ws_tunnel.frames import (
|
||
HelloFrame,
|
||
RequestFrame,
|
||
decode_frame,
|
||
encode_frame,
|
||
)
|
||
from omnigent.runner.transports.ws_tunnel.serve import dispatch_via_asgi
|
||
from omnigent.runner.transports.ws_tunnel.transport import WSTunnelTransport
|
||
from omnigent.runtime import (
|
||
init as init_runtime,
|
||
)
|
||
from omnigent.runtime import (
|
||
set_harness_process_manager,
|
||
set_runner_id,
|
||
set_runner_router,
|
||
)
|
||
from omnigent.runtime.agent_cache import AgentCache
|
||
from omnigent.runtime.harnesses import _HARNESS_MODULES
|
||
from omnigent.server.app import create_app
|
||
from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore
|
||
from omnigent.stores.artifact_store.local import LocalArtifactStore
|
||
from omnigent.stores.conversation_store.sqlalchemy_store import (
|
||
SqlAlchemyConversationStore,
|
||
)
|
||
from omnigent.stores.file_store.sqlalchemy_store import SqlAlchemyFileStore
|
||
from tests.runner.helpers import NullServerClient
|
||
from tests.runtime.harnesses._test_scaffold_harnesses import _EchoHarness
|
||
|
||
# ``claude-sdk`` is one of the harness names the agent spec validator
|
||
# accepts. The fixture rewrites its registry entry to point at the
|
||
# echo-harness scaffold module for the duration of the test (then
|
||
# restores) so ``_create_executor`` picks the runner client path
|
||
# without trying to import the real claude-sdk module.
|
||
_TEST_HARNESS_NAME = "claude-sdk"
|
||
_RUNNER_ID = "runner-tunnel-three-layer-test"
|
||
|
||
|
||
class FakeProcessManager:
|
||
"""In-process replacement for :class:`HarnessProcessManager`.
|
||
|
||
Returns a per-conversation ``httpx.AsyncClient`` whose transport
|
||
is an ``ASGITransport`` to a freshly-built :class:`_EchoHarness`
|
||
FastAPI app with ``app.state.conversation_id`` pinned to that
|
||
conversation. Mirrors the parts of the production process
|
||
manager that the runner and ``the harness HTTP client`` actually
|
||
call.
|
||
|
||
Duplicates :class:`tests.server.integration.test_sessions_three_layer.FakeProcessManager`;
|
||
factor into a shared helper module once a third consumer needs it.
|
||
"""
|
||
|
||
handles_tool_dispatch = False
|
||
|
||
def __init__(self) -> None:
|
||
self._apps: dict[str, FastAPI] = {}
|
||
self._clients: dict[str, httpx.AsyncClient] = {}
|
||
self._in_flight: dict[str, str] = {}
|
||
|
||
async def start(self) -> None: # pragma: no cover — protocol shim
|
||
pass
|
||
|
||
async def shutdown(self) -> None:
|
||
for client in self._clients.values():
|
||
await client.aclose()
|
||
self._apps.clear()
|
||
self._clients.clear()
|
||
|
||
async def get_client(
|
||
self,
|
||
conversation_id: str,
|
||
harness: str,
|
||
env: dict[str, str] | None = None,
|
||
) -> httpx.AsyncClient:
|
||
if conversation_id not in self._clients:
|
||
app = _EchoHarness().build()
|
||
app.state.conversation_id = conversation_id
|
||
self._apps[conversation_id] = app
|
||
self._clients[conversation_id] = httpx.AsyncClient(
|
||
transport=httpx.ASGITransport(app=app),
|
||
base_url="http://harness",
|
||
timeout=httpx.Timeout(5.0, read=None),
|
||
)
|
||
return self._clients[conversation_id]
|
||
|
||
def has_session(self, conversation_id: str) -> bool:
|
||
return conversation_id in self._clients
|
||
|
||
def has_active_turn(self, conversation_id: str) -> bool:
|
||
return conversation_id in self._in_flight
|
||
|
||
def mark_in_flight(self, conversation_id: str, response_id: str) -> None:
|
||
# Mirror the real manager: the runner registers the live turn on
|
||
# response.created so the idle reaper spares it.
|
||
self._in_flight[conversation_id] = response_id
|
||
|
||
def clear_in_flight(self, conversation_id: str) -> None:
|
||
self._in_flight.pop(conversation_id, None)
|
||
|
||
async def release(self, conversation_id: str) -> None:
|
||
client = self._clients.pop(conversation_id, None)
|
||
if client is not None:
|
||
await client.aclose()
|
||
self._apps.pop(conversation_id, None)
|
||
self._in_flight.pop(conversation_id, None)
|
||
|
||
async def forward_cancel(self, conversation_id: str) -> None:
|
||
# EchoHarness completes synchronously; cancellation is not
|
||
# exercised by the happy-path test. A future tunneled
|
||
# cancellation test would extend EchoHarness to block on
|
||
# an event the test can release after the cancel frame.
|
||
del conversation_id
|
||
|
||
|
||
def _build_harness_agent_bundle() -> bytes:
|
||
"""Build an agent bundle that routes through the harness path.
|
||
|
||
Uses ``executor.type='omnigent'`` with
|
||
``executor.harness=<TEST_HARNESS>`` so ``_create_executor`` in
|
||
``runtime/workflow.py`` routes to ``the harness HTTP client`` and
|
||
thus through the runner → harness HTTP chain.
|
||
|
||
:returns: tar.gz bytes with a single ``config.yaml``.
|
||
"""
|
||
config: dict[str, Any] = {
|
||
"spec_version": 1,
|
||
"name": "echo-tunnel-test",
|
||
"executor": {
|
||
"type": "omnigent",
|
||
"config": {"harness": _TEST_HARNESS_NAME},
|
||
"model": "test-model",
|
||
},
|
||
}
|
||
buf = io.BytesIO()
|
||
with tarfile.open(fileobj=buf, mode="w:gz") as tf:
|
||
data = yaml.dump(config).encode()
|
||
info = tarfile.TarInfo(name="config.yaml")
|
||
info.size = len(data)
|
||
tf.addfile(info, io.BytesIO(data))
|
||
return buf.getvalue()
|
||
|
||
|
||
# ── Tunnel handshake / forwarder helpers ──────────────────
|
||
|
||
|
||
def _websocket_scope(
|
||
path: str,
|
||
*,
|
||
client_host: str = "127.0.0.1",
|
||
) -> dict[str, object]:
|
||
"""Build a minimal ASGI WebSocket scope for the tunnel route."""
|
||
return {
|
||
"type": "websocket",
|
||
"asgi": {"version": "3.0"},
|
||
"scheme": "ws",
|
||
"path": path,
|
||
"raw_path": path.encode("ascii"),
|
||
"query_string": b"",
|
||
"headers": [],
|
||
"client": (client_host, 50000),
|
||
"server": ("testserver", 80),
|
||
"subprotocols": [],
|
||
}
|
||
|
||
|
||
async def _connect_runner_tunnel(
|
||
app: FastAPI,
|
||
runner_id: str,
|
||
) -> ApplicationCommunicator:
|
||
"""Open an ASGI WebSocket against the runner tunnel route.
|
||
|
||
:returns: A communicator already past ``websocket.accept``.
|
||
"""
|
||
communicator = ApplicationCommunicator(
|
||
app,
|
||
_websocket_scope(f"/v1/runners/{runner_id}/tunnel"),
|
||
)
|
||
await communicator.send_input({"type": "websocket.connect"})
|
||
accepted = await communicator.receive_output(timeout=2.0)
|
||
assert accepted["type"] == "websocket.accept", (
|
||
f"Tunnel route did not accept the WS handshake; got {accepted!r}"
|
||
)
|
||
return communicator
|
||
|
||
|
||
async def _send_hello_and_wait(
|
||
communicator: ApplicationCommunicator,
|
||
app: FastAPI,
|
||
runner_id: str,
|
||
*,
|
||
harnesses: list[str],
|
||
) -> None:
|
||
"""Send a HelloFrame and wait until the registry lists the runner."""
|
||
hello = HelloFrame(
|
||
runner_version="0.1.0-test",
|
||
frame_protocol_version=1,
|
||
harnesses=list(harnesses),
|
||
envs=["os_sandbox"],
|
||
)
|
||
await communicator.send_input(
|
||
{"type": "websocket.receive", "text": encode_frame(hello)},
|
||
)
|
||
registry = app.state.tunnel_registry
|
||
|
||
async def _registered() -> None:
|
||
while registry.get(runner_id) is None:
|
||
await asyncio.sleep(0.01)
|
||
|
||
await asyncio.wait_for(_registered(), timeout=2.0)
|
||
|
||
|
||
async def _forward_requests_to_runner(
|
||
communicator: ApplicationCommunicator,
|
||
runner_app: FastAPI,
|
||
) -> None:
|
||
"""Pump frames from the tunnel into the runner ASGI app.
|
||
|
||
Reads ``websocket.send`` outputs from the server side and decodes
|
||
each ``RequestFrame``. Each request is dispatched in its own
|
||
asyncio task so streaming responses (SSE relays) don't block
|
||
concurrent requests — the production runner WS loop has the same
|
||
shape. Frame writes back to the communicator are serialised
|
||
behind ``_write_lock`` because ``ApplicationCommunicator.send_input``
|
||
is not safe to call from multiple concurrent tasks.
|
||
"""
|
||
pending: set[asyncio.Task[None]] = set()
|
||
write_lock = asyncio.Lock()
|
||
|
||
async def _send_text(payload: str) -> None:
|
||
async with write_lock:
|
||
await communicator.send_input(
|
||
{"type": "websocket.receive", "text": payload},
|
||
)
|
||
|
||
try:
|
||
while True:
|
||
# ApplicationCommunicator.receive_output() defaults to a
|
||
# 1-second timeout. Under xdist scheduling pressure in CI
|
||
# the >=1s gap between WS frames is common (one xdist
|
||
# worker holds the event loop while EchoHarness runs), so
|
||
# the default triggers asyncio.TimeoutError which kills
|
||
# the forwarder and tears down the tunnel mid-turn. That
|
||
# severs the AP-side SSE relay before it has persisted the
|
||
# assistant message, and the test fails with "no assistant
|
||
# text in session snapshot". 60s is well above any
|
||
# plausible per-frame interval but still bounded so a
|
||
# genuinely stuck test fails rather than hangs.
|
||
output = await communicator.receive_output(timeout=60.0)
|
||
if output["type"] == "websocket.close":
|
||
return
|
||
if output["type"] != "websocket.send":
|
||
continue
|
||
|
||
text = output.get("text")
|
||
if not isinstance(text, str):
|
||
continue
|
||
|
||
frame = decode_frame(text)
|
||
if not isinstance(frame, RequestFrame):
|
||
continue
|
||
|
||
task = asyncio.create_task(
|
||
dispatch_via_asgi(runner_app, frame, _send_text),
|
||
name=f"tunnel-dispatch-{frame.id}",
|
||
)
|
||
pending.add(task)
|
||
task.add_done_callback(pending.discard)
|
||
finally:
|
||
for task in list(pending):
|
||
task.cancel()
|
||
for task in list(pending):
|
||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||
await task
|
||
|
||
|
||
# ── Fixture: Omnigent app + tunnel WS + runner app + EchoHarness ────
|
||
|
||
|
||
@dataclass
|
||
class _TunnelStack:
|
||
ap_client: httpx.AsyncClient
|
||
ap_app: FastAPI
|
||
fake_pm: FakeProcessManager
|
||
|
||
|
||
@pytest_asyncio.fixture()
|
||
async def tunnel_three_layer_stack(tmp_path: Path) -> AsyncIterator[_TunnelStack]:
|
||
"""Wire Omnigent server + WS-tunneled runner + EchoHarness in-process.
|
||
|
||
Lifecycle: build stores, init runtime, override the test harness
|
||
module entry, build Omnigent app + runner app, open the WS tunnel via
|
||
``ApplicationCommunicator``, send hello, start the forwarder
|
||
task, register the runner in the AP-side ``set_runner_router``
|
||
so resource paths can resolve it, then yield an
|
||
``httpx.AsyncClient`` to the AP.
|
||
|
||
Teardown shuts down in reverse order (httpx → forwarder →
|
||
communicator → fake PM → runtime globals → DBOS) so DBOS
|
||
background threads observe a clean exit before pytest-asyncio
|
||
closes the loop.
|
||
"""
|
||
db_uri = f"sqlite:///{tmp_path / 'test.db'}"
|
||
|
||
agent_store = SqlAlchemyAgentStore(db_uri)
|
||
conv_store = SqlAlchemyConversationStore(db_uri)
|
||
file_store = SqlAlchemyFileStore(db_uri)
|
||
artifact_store = LocalArtifactStore(str(tmp_path / "artifacts"))
|
||
agent_cache = AgentCache(
|
||
artifact_store=artifact_store,
|
||
cache_dir=tmp_path / "cache",
|
||
)
|
||
|
||
init_runtime(
|
||
conversation_store=conv_store,
|
||
agent_store=agent_store,
|
||
agent_cache=agent_cache,
|
||
file_store=file_store,
|
||
artifact_store=artifact_store,
|
||
)
|
||
|
||
# Override the harness module entry so ``_create_executor`` routes
|
||
# ``claude-sdk`` to the EchoHarness scaffold for this test only.
|
||
_saved_harness_module = _HARNESS_MODULES.get(_TEST_HARNESS_NAME)
|
||
_HARNESS_MODULES[_TEST_HARNESS_NAME] = "tests.runtime.harnesses._test_scaffold_harnesses"
|
||
|
||
fake_pm = FakeProcessManager()
|
||
set_harness_process_manager(fake_pm)
|
||
set_runner_id(_RUNNER_ID)
|
||
|
||
async def _resolve_spec(agent_id: str, session_id: str | None = None) -> Any:
|
||
record = agent_store.get(agent_id)
|
||
if record is None:
|
||
return None
|
||
return agent_cache.load(agent_id, record.bundle_location).spec
|
||
|
||
runner_app = create_runner_app(
|
||
process_manager=fake_pm,
|
||
spec_resolver=_resolve_spec,
|
||
server_client=NullServerClient(), # type: ignore[arg-type]
|
||
)
|
||
|
||
# Build the Omnigent server. ``create_app`` constructs the
|
||
# ``TunnelRegistry`` + ``RunnerRouter`` synchronously (before
|
||
# lifespan) and threads ``runner_router`` into every router as a
|
||
# closure, so the WS route + the sessions route share the same
|
||
# registry without us needing to drive the lifespan context
|
||
# manager. We still publish the router into the runtime globals
|
||
# so paths like ``_get_runner_client_for_resource_access`` (which
|
||
# reads ``get_runner_router()``) see the same instance.
|
||
ap_app = create_app(
|
||
agent_store=agent_store,
|
||
file_store=file_store,
|
||
conversation_store=conv_store,
|
||
artifact_store=artifact_store,
|
||
agent_cache=agent_cache,
|
||
)
|
||
set_runner_router(ap_app.state.runner_router)
|
||
|
||
communicator = await _connect_runner_tunnel(ap_app, _RUNNER_ID)
|
||
await _send_hello_and_wait(
|
||
communicator,
|
||
ap_app,
|
||
_RUNNER_ID,
|
||
harnesses=[_TEST_HARNESS_NAME],
|
||
)
|
||
|
||
forwarder_task = asyncio.create_task(
|
||
_forward_requests_to_runner(communicator, runner_app),
|
||
name="tunnel-three-layer-forwarder",
|
||
)
|
||
|
||
ap_client = httpx.AsyncClient(
|
||
transport=httpx.ASGITransport(app=ap_app),
|
||
base_url="http://ap",
|
||
timeout=httpx.Timeout(10.0, read=None),
|
||
)
|
||
|
||
try:
|
||
yield _TunnelStack(ap_client=ap_client, ap_app=ap_app, fake_pm=fake_pm)
|
||
finally:
|
||
# Cancel any AP-side background tasks (SSE relays, etc.) the
|
||
# production lifespan would have owned. ``create_app`` does not
|
||
# track them itself, so we sweep the loop for tasks whose
|
||
# ``get_name()`` advertises a runner-relay coroutine and cancel
|
||
# them before tearing down the tunnel. Without this, an
|
||
# in-flight ``runner-relay-*`` task races the WS disconnect and
|
||
# surfaces ``ConnectionError: tunnel closed`` as an uncaught
|
||
# task exception during teardown.
|
||
loop_tasks = []
|
||
for task in asyncio.all_tasks():
|
||
if task is asyncio.current_task():
|
||
continue
|
||
coro_qualname = getattr(task.get_coro(), "__qualname__", "")
|
||
if (
|
||
task.get_name().startswith("runner-relay-")
|
||
or "_relay_runner_stream" in coro_qualname
|
||
):
|
||
loop_tasks.append(task)
|
||
for t in loop_tasks:
|
||
t.cancel()
|
||
for t in loop_tasks:
|
||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||
await t
|
||
|
||
await ap_client.aclose()
|
||
|
||
forwarder_task.cancel()
|
||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||
await forwarder_task
|
||
|
||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||
await communicator.send_input(
|
||
{"type": "websocket.disconnect", "code": 1000},
|
||
)
|
||
with contextlib.suppress(asyncio.TimeoutError, Exception):
|
||
await communicator.wait(timeout=2.0)
|
||
|
||
# Close any cached WSTunnelTransport-backed clients before
|
||
# we tear the registry down so they can flush in-flight
|
||
# response futures cleanly.
|
||
with contextlib.suppress(Exception):
|
||
await ap_app.state.runner_router.aclose()
|
||
|
||
await fake_pm.shutdown()
|
||
set_runner_router(None)
|
||
set_runner_id(None)
|
||
set_harness_process_manager(None)
|
||
if _saved_harness_module is None:
|
||
_HARNESS_MODULES.pop(_TEST_HARNESS_NAME, None)
|
||
else:
|
||
_HARNESS_MODULES[_TEST_HARNESS_NAME] = _saved_harness_module
|
||
|
||
|
||
def _fake_client(ap_client: httpx.AsyncClient) -> object:
|
||
"""Build the SDK namespace shim expected by the REPL adapter.
|
||
|
||
The sessions adapter only uses these namespace attributes; using
|
||
the real namespace classes keeps request/response parsing on the
|
||
production client path while avoiding an actual network client.
|
||
|
||
:param ap_client: ASGI-backed client pointed at the Omnigent app.
|
||
:returns: Duck-typed client with the namespaces the adapter uses.
|
||
"""
|
||
base = str(ap_client.base_url).rstrip("/")
|
||
|
||
class _FakeClient:
|
||
"""Duck-typed Omnigent client backed by the tunnel-stack httpx client."""
|
||
|
||
def __init__(self) -> None:
|
||
self.sessions = SessionsNamespace(ap_client, base)
|
||
self.files = FilesNamespace(ap_client, base)
|
||
|
||
return _FakeClient()
|
||
|
||
|
||
def _new_repl_adapter(
|
||
ap_client: httpx.AsyncClient,
|
||
*,
|
||
agent_name: str,
|
||
session_id: str | None = None,
|
||
runner_id: str = _RUNNER_ID,
|
||
runner_recover: Any | None = None,
|
||
) -> _SessionsChatReplAdapter:
|
||
"""Create a sessions REPL adapter over the tunneled Omnigent stack.
|
||
|
||
:param ap_client: ASGI-backed client pointed at the Omnigent app.
|
||
:param agent_name: Human-readable agent display name.
|
||
:param session_id: Optional existing session to resume.
|
||
:param runner_id: Runner id the adapter should bind before send.
|
||
:param runner_recover: Optional callback returning replacement runner id.
|
||
:returns: Configured :class:`_SessionsChatReplAdapter`.
|
||
"""
|
||
return _SessionsChatReplAdapter(
|
||
client=_fake_client(ap_client), # type: ignore[arg-type]
|
||
agent_name=agent_name,
|
||
session_id=session_id,
|
||
session_bundle=None if session_id is not None else _build_harness_agent_bundle(),
|
||
runner_id=runner_id,
|
||
runner_recover=runner_recover,
|
||
)
|
||
|
||
|
||
async def _collect_adapter_turn(
|
||
adapter: _SessionsChatReplAdapter,
|
||
text: str,
|
||
) -> list[object]:
|
||
"""Send one adapter turn and collect terminal events.
|
||
|
||
``_SessionsChatReplAdapter.send`` intentionally yields only the
|
||
terminal SDK-shaped event; streaming deltas render through the
|
||
background callback path and are covered by the smoke test in this
|
||
file plus the legacy adapter file.
|
||
|
||
:param adapter: Adapter under test.
|
||
:param text: User input for :meth:`_SessionsChatReplAdapter.send`.
|
||
:returns: Events yielded by the ``send()`` generator.
|
||
"""
|
||
send_events: list[object] = []
|
||
async for ev in adapter.send(text):
|
||
send_events.append(ev)
|
||
if isinstance(ev, (ResponseCompleted, ResponseFailed)):
|
||
break
|
||
return send_events
|
||
|
||
|
||
# ── Tests ─────────────────────────────────────────────────
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_native_session_happy_path_via_ws_tunnel(
|
||
tunnel_three_layer_stack: _TunnelStack,
|
||
) -> None:
|
||
"""End-to-end native sessions flow through the live WS tunnel.
|
||
|
||
Mirrors ``test_session_happy_path_through_three_layers`` but
|
||
forces traffic across the production tunnel route. Failure here
|
||
points at either the AP↔tunnel wire (frame encoding, registry
|
||
handshake) or the runner ASGI dispatcher — both invisible to the
|
||
pure-ASGI three-layer test.
|
||
"""
|
||
ap_client = tunnel_three_layer_stack.ap_client
|
||
bundle = _build_harness_agent_bundle()
|
||
|
||
create_resp = await ap_client.post(
|
||
"/v1/sessions",
|
||
data={"metadata": json.dumps({})},
|
||
files={"bundle": ("agent.tar.gz", bundle, "application/gzip")},
|
||
)
|
||
assert create_resp.status_code == 201, (
|
||
f"session create failed: {create_resp.status_code} {create_resp.text}"
|
||
)
|
||
session_id = create_resp.json()["session_id"]
|
||
|
||
bind_resp = await ap_client.patch(
|
||
f"/v1/sessions/{session_id}",
|
||
json={"runner_id": _RUNNER_ID},
|
||
)
|
||
assert bind_resp.status_code == 200, (
|
||
f"runner bind failed: {bind_resp.status_code} {bind_resp.text}"
|
||
)
|
||
|
||
event_resp = await ap_client.post(
|
||
f"/v1/sessions/{session_id}/events",
|
||
json={
|
||
"type": "message",
|
||
"data": {
|
||
"role": "user",
|
||
"content": [{"type": "input_text", "text": "ping"}],
|
||
},
|
||
},
|
||
)
|
||
assert event_resp.status_code == 202, (
|
||
f"message event failed: {event_resp.status_code} {event_resp.text}"
|
||
)
|
||
|
||
# Poll until the relay has persisted the assistant message and
|
||
# the cache reflects a terminal status. Checking status alone is
|
||
# not enough: ``_get_session_snapshot`` defaults to ``"idle"`` on
|
||
# a cache miss when the legacy ``get_runner_client()`` singleton
|
||
# is unset (which is the case in any router-only deployment,
|
||
# including this fixture). On a cold-cache run the relay can
|
||
# still be in its initial transport handshake when the test
|
||
# polls, so the first GET returns the default ``"idle"`` and the
|
||
# old loop exited before the harness had emitted a single
|
||
# delta, turning into a flaky "no assistant text" failure.
|
||
# Waiting for the persisted assistant item makes us cancel-race
|
||
# resistant too: the fixture's ``failed`` to ``idle`` rewrite
|
||
# can no longer let us bail on a turn the runner never finished.
|
||
# 600 × 100 ms = 60 s budget, well above the WS-frame round-trip
|
||
# overhead but generous enough to absorb CI scheduler jitter.
|
||
def _assistant_texts(snapshot: dict[str, Any]) -> list[str]:
|
||
return [
|
||
chunk.get("text", "")
|
||
for item in snapshot["items"]
|
||
if item["type"] == "message" and item["data"]["role"] == "assistant"
|
||
for chunk in item["data"]["content"]
|
||
if chunk.get("type") == "output_text"
|
||
]
|
||
|
||
snap: dict[str, Any] | None = None
|
||
final_status: str | None = None
|
||
for _ in range(600):
|
||
snap_resp = await ap_client.get(f"/v1/sessions/{session_id}")
|
||
snap = snap_resp.json()
|
||
final_status = snap["status"]
|
||
if final_status == "failed":
|
||
break
|
||
if final_status == "idle" and _assistant_texts(snap):
|
||
break
|
||
await asyncio.sleep(0.1)
|
||
|
||
assert snap is not None
|
||
assert final_status == "idle", (
|
||
f"session did not reach idle (got {final_status!r}); snapshot={json.dumps(snap, indent=2)}"
|
||
)
|
||
|
||
assistant_texts = _assistant_texts(snap)
|
||
assert assistant_texts, (
|
||
f"no assistant text in session snapshot; items={json.dumps(snap['items'], indent=2)}"
|
||
)
|
||
joined = " ".join(assistant_texts)
|
||
assert "ping" in joined, f"echoed input missing from assistant text: {joined!r}"
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_runner_router_uses_ws_tunnel_for_session_resources(
|
||
tunnel_three_layer_stack: _TunnelStack,
|
||
) -> None:
|
||
"""Resource-access lookups go through the same WS-tunneled client.
|
||
|
||
Lightweight REPL-adapter-shaped smoke: after binding, the AP's
|
||
runner router must hand out a ``WSTunnelTransport``-backed client
|
||
(not a fallback ASGI one) for the session. This catches the
|
||
regression class where ``_get_runner_client_for_resource_access``
|
||
falls back to ``get_runner_client()`` because the router lookup
|
||
silently returned ``None``.
|
||
"""
|
||
ap_client = tunnel_three_layer_stack.ap_client
|
||
ap_app = tunnel_three_layer_stack.ap_app
|
||
|
||
create_resp = await ap_client.post(
|
||
"/v1/sessions",
|
||
data={"metadata": json.dumps({})},
|
||
files={
|
||
"bundle": (
|
||
"agent.tar.gz",
|
||
_build_harness_agent_bundle(),
|
||
"application/gzip",
|
||
),
|
||
},
|
||
)
|
||
assert create_resp.status_code == 201
|
||
session_id = create_resp.json()["session_id"]
|
||
|
||
bind_resp = await ap_client.patch(
|
||
f"/v1/sessions/{session_id}",
|
||
json={"runner_id": _RUNNER_ID},
|
||
)
|
||
assert bind_resp.status_code == 200
|
||
|
||
routed = ap_app.state.runner_router.client_for_session_resources(session_id)
|
||
assert routed.runner_id == _RUNNER_ID, (
|
||
f"router did not pin to {_RUNNER_ID!r}; got {routed.runner_id!r}"
|
||
)
|
||
|
||
transport = routed.client._transport
|
||
assert isinstance(transport, WSTunnelTransport), (
|
||
"router handed out a non-tunnel transport; tunnel wiring "
|
||
f"regressed. Got {type(transport).__name__}."
|
||
)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_repl_adapter_smoke_via_ws_tunnel(
|
||
tunnel_three_layer_stack: _TunnelStack,
|
||
) -> None:
|
||
"""REPL adapter ``send()`` rides the same tunnel and yields terminal.
|
||
|
||
Lightweight smoke that the legacy adapter still completes its
|
||
end-to-end loop when the AP↔runner hop is the production WS
|
||
tunnel rather than ASGITransport. Builds the same client-namespace
|
||
shim ``test_sessions_repl_adapter.py`` uses (typed-duck, not the
|
||
real SDK ``Client``) so we don't pull a network dependency and
|
||
don't import a private helper across test files.
|
||
|
||
Asserts a terminal completion arrives — the EchoHarness emits a
|
||
``response.output_text.delta`` then completes. A failed terminal is
|
||
not accepted here because this test is intended to catch tunnel
|
||
regressions, not merely prove that the adapter can surface them.
|
||
"""
|
||
adapter = _new_repl_adapter(
|
||
tunnel_three_layer_stack.ap_client,
|
||
agent_name="tunnel-smoke-agent",
|
||
)
|
||
|
||
try:
|
||
send_events = await _collect_adapter_turn(adapter, "ping")
|
||
|
||
assert not any(isinstance(e, ResponseFailed) for e in send_events), (
|
||
"adapter.send failed over the tunnel; "
|
||
f"events={[type(e).__name__ for e in send_events]}"
|
||
)
|
||
assert any(isinstance(e, ResponseCompleted) for e in send_events), (
|
||
"adapter.send did not yield ResponseCompleted over the tunnel; "
|
||
f"events={[type(e).__name__ for e in send_events]}"
|
||
)
|
||
finally:
|
||
with contextlib.suppress(Exception):
|
||
await adapter.aclose()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_repl_adapter_resume_rebinds_via_ws_tunnel(
|
||
tunnel_three_layer_stack: _TunnelStack,
|
||
) -> None:
|
||
"""Resumed REPL adapter rebinding works on the native tunnel path.
|
||
|
||
Starts from an existing session bound to the live runner, then
|
||
resumes through an adapter whose local ``_bound_runner_id`` cache
|
||
is stale. The send must PATCH the session before dispatching over
|
||
the tunnel; otherwise the user message would not reach the runner
|
||
or the adapter would skip the resume binding path entirely.
|
||
"""
|
||
ap_client = tunnel_three_layer_stack.ap_client
|
||
|
||
create_resp = await ap_client.post(
|
||
"/v1/sessions",
|
||
data={"metadata": json.dumps({})},
|
||
files={
|
||
"bundle": (
|
||
"agent.tar.gz",
|
||
_build_harness_agent_bundle(),
|
||
"application/gzip",
|
||
),
|
||
},
|
||
)
|
||
assert create_resp.status_code == 201
|
||
session_id = create_resp.json()["session_id"]
|
||
|
||
bind_resp = await ap_client.patch(
|
||
f"/v1/sessions/{session_id}",
|
||
json={"runner_id": _RUNNER_ID},
|
||
)
|
||
assert bind_resp.status_code == 200
|
||
|
||
adapter = _new_repl_adapter(
|
||
ap_client,
|
||
agent_name="tunnel-resume-agent",
|
||
session_id=session_id,
|
||
runner_id=_RUNNER_ID,
|
||
)
|
||
adapter._bound_runner_id = "stale-runner-cache"
|
||
|
||
try:
|
||
send_events = await _collect_adapter_turn(adapter, "resumed turn")
|
||
assert not any(isinstance(e, ResponseFailed) for e in send_events), (
|
||
f"resumed turn failed; events={[type(e).__name__ for e in send_events]}"
|
||
)
|
||
assert any(isinstance(e, ResponseCompleted) for e in send_events), (
|
||
f"resumed turn did not complete; events={[type(e).__name__ for e in send_events]}"
|
||
)
|
||
assert adapter.session_id == session_id
|
||
assert adapter._bound_runner_id == _RUNNER_ID
|
||
|
||
snap_resp = await ap_client.get(f"/v1/sessions/{session_id}")
|
||
assert snap_resp.status_code == 200
|
||
snap = snap_resp.json()
|
||
assert snap["runner_id"] == _RUNNER_ID
|
||
user_texts = [
|
||
block["text"]
|
||
for item in snap["items"]
|
||
if item.get("type") == "message"
|
||
and isinstance(item.get("data"), dict)
|
||
and item["data"].get("role") == "user"
|
||
for block in item["data"].get("content", [])
|
||
if isinstance(block, dict) and block.get("type") == "input_text"
|
||
]
|
||
assert "resumed turn" in user_texts, (
|
||
f"resumed input missing from snapshot: {json.dumps(snap['items'], indent=2)}"
|
||
)
|
||
finally:
|
||
with contextlib.suppress(Exception):
|
||
await adapter.aclose()
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_on_runner_connect_restarts_relay_via_router(
|
||
tunnel_three_layer_stack: _TunnelStack,
|
||
) -> None:
|
||
"""Reconnect hook restarts relays via the router.
|
||
|
||
Pre-fix, ``_on_runner_connect`` short-circuited on
|
||
``get_runner_client() is None`` — always true in multi-runner
|
||
deployments where only ``set_runner_router`` is wired — so every
|
||
reconnect was a silent no-op.
|
||
|
||
Drives a real WS tunnel disconnect/reconnect and asserts the hook
|
||
(a) routes through ``client_for_session_resources`` for the bound
|
||
session, (b) reaches ``_ensure_runner_relay`` with the exact
|
||
routed client, and (c) leaves a sibling session bound to a
|
||
different runner alone (per-conv ``runner_id`` filter).
|
||
"""
|
||
from omnigent.runner.routing import RoutedRunner
|
||
from omnigent.server.routes import sessions as sessions_routes
|
||
from omnigent.server.routes.sessions import _runner_relay_tasks
|
||
|
||
ap_client = tunnel_three_layer_stack.ap_client
|
||
ap_app = tunnel_three_layer_stack.ap_app
|
||
fake_pm = tunnel_three_layer_stack.fake_pm
|
||
|
||
create_resp = await ap_client.post(
|
||
"/v1/sessions",
|
||
data={"metadata": json.dumps({})},
|
||
files={
|
||
"bundle": (
|
||
"agent.tar.gz",
|
||
_build_harness_agent_bundle(),
|
||
"application/gzip",
|
||
),
|
||
},
|
||
)
|
||
assert create_resp.status_code == 201, create_resp.text
|
||
session_id = create_resp.json()["session_id"]
|
||
|
||
bind_resp = await ap_client.patch(
|
||
f"/v1/sessions/{session_id}",
|
||
json={"runner_id": _RUNNER_ID},
|
||
)
|
||
assert bind_resp.status_code == 200, bind_resp.text
|
||
|
||
# Sibling session bound to a different runner so the per-conv
|
||
# ``runner_id`` filter has something to skip. Written via the
|
||
# store directly because PATCH-bind would spawn a relay against
|
||
# a runner that has no WS pump, hanging teardown.
|
||
from omnigent.runtime import get_conversation_store
|
||
|
||
other_runner_id = "runner-other-irrelevant"
|
||
other_create_resp = await ap_client.post(
|
||
"/v1/sessions",
|
||
data={"metadata": json.dumps({})},
|
||
files={
|
||
"bundle": (
|
||
"agent.tar.gz",
|
||
_build_harness_agent_bundle(),
|
||
"application/gzip",
|
||
),
|
||
},
|
||
)
|
||
assert other_create_resp.status_code == 201, other_create_resp.text
|
||
other_session_id = other_create_resp.json()["session_id"]
|
||
get_conversation_store().replace_runner_id(other_session_id, other_runner_id)
|
||
|
||
initial_handle = _runner_relay_tasks.get(session_id)
|
||
assert initial_handle is not None, (
|
||
"PATCH-bind did not spawn a relay; the reconnect assertion "
|
||
"below would be vacuous without a baseline relay to replace."
|
||
)
|
||
initial_task = initial_handle.task
|
||
|
||
# Spies on the router resolver and ``_ensure_runner_relay``.
|
||
# The stub ``.post`` short-circuits because ``WSTunnelTransport``
|
||
# ignores httpx ``timeout=`` and would hang. The
|
||
# ensure spy does not chain through — the real helper would spawn
|
||
# an SSE task that errors immediately against the stub client.
|
||
router = ap_app.state.runner_router
|
||
real_resolver = router.client_for_session_resources
|
||
routed_calls: list[str] = []
|
||
routed_clients: dict[str, Any] = {}
|
||
|
||
class _StubResponse:
|
||
status_code = 200
|
||
|
||
def raise_for_status(self) -> None:
|
||
return None
|
||
|
||
class _StubClient:
|
||
async def post(self, *args: Any, **kwargs: Any) -> _StubResponse:
|
||
return _StubResponse()
|
||
|
||
def _spy_resolver(conv_id: str): # type: ignore[no-untyped-def]
|
||
routed_calls.append(conv_id)
|
||
real_routed = real_resolver(conv_id)
|
||
fake = RoutedRunner(runner_id=real_routed.runner_id, client=_StubClient()) # type: ignore[arg-type]
|
||
routed_clients[conv_id] = fake.client
|
||
return fake
|
||
|
||
router.client_for_session_resources = _spy_resolver # type: ignore[method-assign]
|
||
|
||
real_ensure = sessions_routes._ensure_runner_relay
|
||
ensure_calls: list[tuple[str, str | None, Any]] = []
|
||
|
||
def _spy_ensure(sid, rid, client, store=None): # type: ignore[no-untyped-def]
|
||
ensure_calls.append((sid, rid, client))
|
||
|
||
sessions_routes._ensure_runner_relay = _spy_ensure # type: ignore[assignment]
|
||
|
||
new_communicator: ApplicationCommunicator | None = None
|
||
new_forwarder_task: asyncio.Task[None] | None = None
|
||
try:
|
||
# Deregister aborts the in-flight relay with ConnectionError;
|
||
# the relay's done-callback then clears its slot.
|
||
ap_app.state.tunnel_registry.deregister(_RUNNER_ID)
|
||
|
||
async def _relay_slot_cleared() -> None:
|
||
while True:
|
||
handle = _runner_relay_tasks.get(session_id)
|
||
if handle is None or handle.task.done():
|
||
return
|
||
await asyncio.sleep(0.01)
|
||
|
||
await asyncio.wait_for(_relay_slot_cleared(), timeout=2.0)
|
||
|
||
# Fresh WS + hello re-registers and fires _on_runner_connect.
|
||
new_communicator = await _connect_runner_tunnel(ap_app, _RUNNER_ID)
|
||
await _send_hello_and_wait(
|
||
new_communicator,
|
||
ap_app,
|
||
_RUNNER_ID,
|
||
harnesses=[_TEST_HARNESS_NAME],
|
||
)
|
||
|
||
async def _resolve_spec(agent_id: str) -> Any:
|
||
return None
|
||
|
||
new_runner_app = create_runner_app(
|
||
process_manager=fake_pm,
|
||
spec_resolver=_resolve_spec,
|
||
server_client=NullServerClient(), # type: ignore[arg-type]
|
||
)
|
||
new_forwarder_task = asyncio.create_task(
|
||
_forward_requests_to_runner(new_communicator, new_runner_app),
|
||
name="tunnel-reconnect-forwarder",
|
||
)
|
||
|
||
async def _hook_did_its_job() -> None:
|
||
while True:
|
||
if session_id in routed_calls and any(
|
||
sid == session_id for sid, _, _ in ensure_calls
|
||
):
|
||
return
|
||
await asyncio.sleep(0.02)
|
||
|
||
try:
|
||
await asyncio.wait_for(_hook_did_its_job(), timeout=5.0)
|
||
except asyncio.TimeoutError:
|
||
registry_entry = ap_app.state.tunnel_registry.get(_RUNNER_ID)
|
||
raise AssertionError(
|
||
"Reconnect did not drive the connect hook to completion within 5s.\n"
|
||
f" routed_calls = {routed_calls}\n"
|
||
f" ensure_calls = {ensure_calls}\n"
|
||
f" registry has runner = {registry_entry is not None}\n"
|
||
f" initial_task.done = {initial_task.done()}"
|
||
) from None
|
||
|
||
assert session_id in routed_calls, (
|
||
"_on_runner_connect did not call "
|
||
"runner_router.client_for_session_resources for the bound "
|
||
"session — legacy get_runner_client() path is still in use "
|
||
"(B2 regression)."
|
||
)
|
||
ensure_for_session = [
|
||
(sid, rid, client) for sid, rid, client in ensure_calls if sid == session_id
|
||
]
|
||
assert ensure_for_session, (
|
||
"_on_runner_connect resolved the router but never reached "
|
||
"_ensure_runner_relay for the bound session."
|
||
)
|
||
_last_sid, last_rid, last_client = ensure_for_session[-1]
|
||
assert last_rid == _RUNNER_ID, (
|
||
f"_ensure_runner_relay was called with runner_id "
|
||
f"{last_rid!r}, expected {_RUNNER_ID!r}."
|
||
)
|
||
assert last_client is routed_clients[session_id], (
|
||
"_ensure_runner_relay was called with a client that is not the "
|
||
"one returned by runner_router.client_for_session_resources — "
|
||
"the routed client is not threading through the hook."
|
||
)
|
||
|
||
# Sibling session bound to a different runner must not be touched
|
||
# — the per-conv runner_id filter is the only cross-runner guard.
|
||
assert other_session_id not in routed_calls, (
|
||
f"Connect hook resolved router for sibling session "
|
||
f"{other_session_id!r} bound to {other_runner_id!r} — "
|
||
f"per-conv runner_id filter regressed."
|
||
)
|
||
assert all(sid != other_session_id for sid, _, _ in ensure_calls), (
|
||
f"Connect hook called _ensure_runner_relay for sibling session "
|
||
f"{other_session_id!r} bound to {other_runner_id!r}."
|
||
)
|
||
finally:
|
||
router.client_for_session_resources = real_resolver # type: ignore[method-assign]
|
||
sessions_routes._ensure_runner_relay = real_ensure # type: ignore[assignment]
|
||
if new_forwarder_task is not None:
|
||
new_forwarder_task.cancel()
|
||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||
await new_forwarder_task
|
||
if new_communicator is not None:
|
||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||
await new_communicator.send_input(
|
||
{"type": "websocket.disconnect", "code": 1000},
|
||
)
|
||
with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError, Exception):
|
||
await new_communicator.wait(timeout=2.0)
|
||
|
||
|
||
@contextlib.asynccontextmanager
|
||
async def _reconnect_fires_connect_hook(
|
||
ap_app: FastAPI,
|
||
fake_pm: Any,
|
||
*,
|
||
wait_for_recover: str,
|
||
) -> AsyncIterator[list[str]]:
|
||
"""Drive a real tunnel disconnect/reconnect so ``_on_runner_connect`` fires.
|
||
|
||
Deregisters ``_RUNNER_ID`` then opens a fresh WS + hello, which is
|
||
the production trigger for the connect hook. Stubs the router
|
||
resolver + ``_ensure_runner_relay`` so the hook does not spawn a real
|
||
SSE relay against a stub client, and wraps the REAL
|
||
``_publish_runner_recovered_status`` with a spy that records
|
||
completion (the closure re-imports it from the module on each call,
|
||
so the patch is picked up). Waits until the recovery helper has run
|
||
to completion for ``wait_for_recover`` before yielding, so callers
|
||
can assert on the post-recovery state.
|
||
|
||
:yields: The list of session ids the recovery helper ran for.
|
||
"""
|
||
from omnigent.runner.routing import RoutedRunner
|
||
from omnigent.server.routes import sessions as sessions_routes
|
||
|
||
router = ap_app.state.runner_router
|
||
real_resolver = router.client_for_session_resources
|
||
|
||
class _StubResponse:
|
||
status_code = 200
|
||
|
||
def raise_for_status(self) -> None:
|
||
return None
|
||
|
||
class _StubClient:
|
||
async def post(self, *args: Any, **kwargs: Any) -> _StubResponse:
|
||
return _StubResponse()
|
||
|
||
def _spy_resolver(conv_id: str): # type: ignore[no-untyped-def]
|
||
real_routed = real_resolver(conv_id)
|
||
return RoutedRunner(runner_id=real_routed.runner_id, client=_StubClient()) # type: ignore[arg-type]
|
||
|
||
router.client_for_session_resources = _spy_resolver # type: ignore[method-assign]
|
||
|
||
real_ensure = sessions_routes._ensure_runner_relay
|
||
|
||
def _stub_ensure(sid, rid, client, store=None): # type: ignore[no-untyped-def]
|
||
return None
|
||
|
||
sessions_routes._ensure_runner_relay = _stub_ensure # type: ignore[assignment]
|
||
|
||
# Wrap (not replace) the real recovery helper so the narrowed
|
||
# disconnect-vs-failure guard is exercised, and record completion so
|
||
# the test can wait for a deterministic signal rather than sleeping.
|
||
real_recover = sessions_routes._publish_runner_recovered_status
|
||
recovered_calls: list[str] = []
|
||
|
||
async def _spy_recover(sid, store, **kwargs): # type: ignore[no-untyped-def]
|
||
await real_recover(sid, store, **kwargs)
|
||
recovered_calls.append(sid)
|
||
|
||
sessions_routes._publish_runner_recovered_status = _spy_recover # type: ignore[assignment]
|
||
|
||
communicator: ApplicationCommunicator | None = None
|
||
forwarder_task: asyncio.Task[None] | None = None
|
||
try:
|
||
ap_app.state.tunnel_registry.deregister(_RUNNER_ID)
|
||
|
||
communicator = await _connect_runner_tunnel(ap_app, _RUNNER_ID)
|
||
await _send_hello_and_wait(
|
||
communicator,
|
||
ap_app,
|
||
_RUNNER_ID,
|
||
harnesses=[_TEST_HARNESS_NAME],
|
||
)
|
||
|
||
async def _resolve_spec(agent_id: str) -> Any:
|
||
return None
|
||
|
||
new_runner_app = create_runner_app(
|
||
process_manager=fake_pm,
|
||
spec_resolver=_resolve_spec,
|
||
server_client=NullServerClient(), # type: ignore[arg-type]
|
||
)
|
||
forwarder_task = asyncio.create_task(
|
||
_forward_requests_to_runner(communicator, new_runner_app),
|
||
name="tunnel-recover-reconnect-forwarder",
|
||
)
|
||
|
||
async def _recovered() -> None:
|
||
while wait_for_recover not in recovered_calls:
|
||
await asyncio.sleep(0.02)
|
||
|
||
try:
|
||
await asyncio.wait_for(_recovered(), timeout=5.0)
|
||
except asyncio.TimeoutError:
|
||
raise AssertionError(
|
||
"Reconnect did not drive _on_runner_connect to call "
|
||
f"_publish_runner_recovered_status for {wait_for_recover!r} "
|
||
f"within 5s. recovered_calls={recovered_calls}"
|
||
) from None
|
||
|
||
yield recovered_calls
|
||
finally:
|
||
router.client_for_session_resources = real_resolver # type: ignore[method-assign]
|
||
sessions_routes._ensure_runner_relay = real_ensure # type: ignore[assignment]
|
||
sessions_routes._publish_runner_recovered_status = real_recover # type: ignore[assignment]
|
||
if forwarder_task is not None:
|
||
forwarder_task.cancel()
|
||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||
await forwarder_task
|
||
if communicator is not None:
|
||
with contextlib.suppress(asyncio.CancelledError, Exception):
|
||
await communicator.send_input(
|
||
{"type": "websocket.disconnect", "code": 1000},
|
||
)
|
||
with contextlib.suppress(asyncio.CancelledError, asyncio.TimeoutError, Exception):
|
||
await communicator.wait(timeout=2.0)
|
||
|
||
|
||
async def _bind_failed_session(
|
||
ap_client: httpx.AsyncClient,
|
||
*,
|
||
error_code: str,
|
||
error_message: str,
|
||
) -> str:
|
||
"""Create a session, bind it to ``_RUNNER_ID`` (no relay), mark it failed.
|
||
|
||
Binds via the store directly (like the sibling session in
|
||
``test_on_runner_connect_restarts_relay_via_router``) because a
|
||
PATCH-bind would spawn a relay whose teardown races the reconnect.
|
||
Then stamps the ``failed`` cache status + persisted
|
||
``last_task_error`` labels a real disconnect / task failure leaves
|
||
behind, so the reconnect hook sees an authentic pre-recovery state.
|
||
"""
|
||
from omnigent.runtime import get_conversation_store
|
||
from omnigent.server.routes import sessions as sessions_module
|
||
from omnigent.server.schemas import ErrorDetail
|
||
|
||
create_resp = await ap_client.post(
|
||
"/v1/sessions",
|
||
data={"metadata": json.dumps({})},
|
||
files={
|
||
"bundle": (
|
||
"agent.tar.gz",
|
||
_build_harness_agent_bundle(),
|
||
"application/gzip",
|
||
),
|
||
},
|
||
)
|
||
assert create_resp.status_code == 201, create_resp.text
|
||
session_id = create_resp.json()["session_id"]
|
||
|
||
store = get_conversation_store()
|
||
store.replace_runner_id(session_id, _RUNNER_ID)
|
||
sessions_module._session_status_cache[session_id] = "failed"
|
||
await sessions_module._persist_session_status_error_labels(
|
||
session_id,
|
||
ErrorDetail(code=error_code, message=error_message),
|
||
store,
|
||
)
|
||
return session_id
|
||
|
||
|
||
@pytest.fixture()
|
||
def _isolated_session_status_cache() -> Iterator[None]:
|
||
"""Snapshot + clear + restore the module-global session-status cache.
|
||
|
||
``_session_status_cache`` lives at module scope in the sessions
|
||
route, so it survives across tests while each test rebuilds its own
|
||
conversation store on a fresh ``tmp_path`` DB. A prior test can leave
|
||
entries behind (and per-test session ids can collide), which makes
|
||
the reconnect-recovery assertions non-deterministic when the whole
|
||
suite runs. Start each guarded test from an empty cache and restore
|
||
the pre-test contents afterward so nothing leaks in either direction.
|
||
"""
|
||
from omnigent.server.routes import sessions as sessions_module
|
||
|
||
saved = dict(sessions_module._session_status_cache)
|
||
sessions_module._session_status_cache.clear()
|
||
try:
|
||
yield
|
||
finally:
|
||
sessions_module._session_status_cache.clear()
|
||
sessions_module._session_status_cache.update(saved)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.usefixtures("_isolated_session_status_cache")
|
||
async def test_on_runner_connect_clears_disconnect_failure_on_idle_reconnect(
|
||
tunnel_three_layer_stack: _TunnelStack,
|
||
) -> None:
|
||
"""Reconnect-to-idle clears a persisted ``runner_disconnected`` failure.
|
||
|
||
A transient WS blip drops and reconnects a runner whose process
|
||
survived, with no new turn. The disconnect left the session marked
|
||
``failed`` with persisted ``runner_disconnected`` labels; before this
|
||
wiring nothing cleared them until the next user message, so the
|
||
Subagents panel kept the grey "Disconnected" dot forever. Drives a
|
||
real tunnel reconnect (which fires ``_on_runner_connect``) and asserts
|
||
the session leaves ``failed`` and the disconnect labels are cleared —
|
||
WITHOUT sending any message.
|
||
"""
|
||
from omnigent.runtime import get_conversation_store
|
||
from omnigent.server.routes import sessions as sessions_module
|
||
|
||
ap_client = tunnel_three_layer_stack.ap_client
|
||
ap_app = tunnel_three_layer_stack.ap_app
|
||
fake_pm = tunnel_three_layer_stack.fake_pm
|
||
|
||
session_id = await _bind_failed_session(
|
||
ap_client,
|
||
error_code="runner_disconnected",
|
||
error_message="Runner disconnected unexpectedly.",
|
||
)
|
||
store = get_conversation_store()
|
||
|
||
# Sanity: this is the exact state a real tunnel drop leaves behind.
|
||
conv = store.get_conversation(session_id)
|
||
assert conv is not None
|
||
assert sessions_module._last_task_error_from_labels(conv.labels) == {
|
||
"code": "runner_disconnected",
|
||
"message": "Runner disconnected unexpectedly.",
|
||
}
|
||
assert sessions_module._session_status_cache.get(session_id) == "failed"
|
||
|
||
try:
|
||
# Assert INSIDE the block: the context manager's teardown closes
|
||
# the tunnel, which re-fires ``_on_runner_disconnect`` and would
|
||
# re-mark the session failed. We want the reconnect-recovery state.
|
||
async with _reconnect_fires_connect_hook(ap_app, fake_pm, wait_for_recover=session_id):
|
||
# No message was sent — the reconnect alone cleared the stale
|
||
# disconnect failure.
|
||
assert sessions_module._session_status_cache.get(session_id) != "failed"
|
||
cleared = store.get_conversation(session_id)
|
||
assert cleared is not None
|
||
assert sessions_module._last_task_error_from_labels(cleared.labels) is None
|
||
finally:
|
||
sessions_module._session_status_cache.pop(session_id, None)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
@pytest.mark.usefixtures("_isolated_session_status_cache")
|
||
async def test_on_runner_connect_preserves_genuine_failure_on_reconnect(
|
||
tunnel_three_layer_stack: _TunnelStack,
|
||
) -> None:
|
||
"""Reconnect must NOT erase a genuine (non-disconnect) task failure.
|
||
|
||
A runner reconnect proves the tunnel is healthy again, which
|
||
invalidates a benign disconnect failure — but says nothing about a
|
||
real task error. A session that genuinely failed (a
|
||
``last_task_error`` code other than ``runner_disconnected``) must keep
|
||
its ``failed`` state and error labels across the reconnect. Drives the
|
||
real ``_on_runner_connect`` and asserts the failure survives.
|
||
"""
|
||
from omnigent.runtime import get_conversation_store
|
||
from omnigent.server.routes import sessions as sessions_module
|
||
|
||
ap_client = tunnel_three_layer_stack.ap_client
|
||
ap_app = tunnel_three_layer_stack.ap_app
|
||
fake_pm = tunnel_three_layer_stack.fake_pm
|
||
|
||
session_id = await _bind_failed_session(
|
||
ap_client,
|
||
error_code="runner_error",
|
||
error_message="Turn failed: agent raised.",
|
||
)
|
||
store = get_conversation_store()
|
||
|
||
try:
|
||
async with _reconnect_fires_connect_hook(ap_app, fake_pm, wait_for_recover=session_id):
|
||
# The recovery helper ran (the hook called it) but self-guarded:
|
||
# a genuine failure is not a stale disconnect, so status + labels
|
||
# are untouched.
|
||
assert sessions_module._session_status_cache.get(session_id) == "failed"
|
||
survived = store.get_conversation(session_id)
|
||
assert survived is not None
|
||
assert sessions_module._last_task_error_from_labels(survived.labels) == {
|
||
"code": "runner_error",
|
||
"message": "Turn failed: agent raised.",
|
||
}
|
||
finally:
|
||
sessions_module._session_status_cache.pop(session_id, None)
|
||
|
||
|
||
# TODO: factor ``FakeProcessManager`` and ``_build_harness_agent_bundle``
|
||
# out of ``test_sessions_three_layer.py`` + this file into a shared
|
||
# ``_three_layer_helpers.py`` module once a third caller arrives. Kept
|
||
# duplicated for now to keep this PR scoped to fixture + tests only.
|