refactor: use rpc

tests/integration_tests/box/test_box_integration.py

@@ -12,21 +12,22 @@ CI pipeline.  Run them locally with::
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import shutil
 import socket
 import subprocess
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
-from aiohttp.test_utils import TestServer
 
 from langbot.pkg.box.backend import BaseSandboxBackend
-from langbot.pkg.box.client import RemoteBoxRuntimeClient
+from langbot.pkg.box.client import ActionRPCBoxClient
 from langbot.pkg.box.errors import BoxBackendUnavailableError, BoxRuntimeUnavailableError
 from langbot.pkg.box.models import BoxExecutionStatus, BoxNetworkMode, BoxSpec
 from langbot.pkg.box.runtime import BoxRuntime
-from langbot.pkg.box.server import create_app as create_server_app
+from langbot.pkg.box.server import BoxServerHandler
 from langbot.pkg.box.service import BoxService
 
 import langbot_plugin.api.entities.builtin.pipeline.query as pipeline_query
@@ -77,23 +78,61 @@ requires_socket = pytest.mark.skipif(
 )
 
 
+# ── Helpers ──────────────────────────────────────────────────────────
+
+
+class _QueueConnection:
+    """In-process Connection backed by asyncio Queues — no real IO."""
+
+    def __init__(self, rx: asyncio.Queue[str], tx: asyncio.Queue[str]):
+        self._rx = rx
+        self._tx = tx
+
+    async def send(self, message: str) -> None:
+        await self._tx.put(message)
+
+    async def receive(self) -> str:
+        return await self._rx.get()
+
+    async def close(self) -> None:
+        pass
+
+
+async def _make_rpc_pair(runtime: BoxRuntime):
+    """Create an in-process (ActionRPCBoxClient, server_task, client_task) connected via queues."""
+    from langbot_plugin.runtime.io.handler import Handler
+
+    c2s: asyncio.Queue[str] = asyncio.Queue()
+    s2c: asyncio.Queue[str] = asyncio.Queue()
+    client_conn = _QueueConnection(rx=s2c, tx=c2s)
+    server_conn = _QueueConnection(rx=c2s, tx=s2c)
+
+    server_handler = BoxServerHandler(server_conn, runtime)
+    server_task = asyncio.create_task(server_handler.run())
+
+    client_handler = Handler.__new__(Handler)
+    Handler.__init__(client_handler, client_conn)
+    client_task = asyncio.create_task(client_handler.run())
+
+    client = ActionRPCBoxClient(logger=_logger)
+    client.set_handler(client_handler)
+
+    return client, server_task, client_task
+
+
 # ── Fixtures ──────────────────────────────────────────────────────────
 
 
 @pytest.fixture
 async def box_client():
-    """Yield a RemoteBoxRuntimeClient backed by a real BoxRuntime HTTP server."""
+    """Yield an ActionRPCBoxClient backed by a real BoxRuntime via in-process RPC."""
     runtime = BoxRuntime(logger=_logger)
-    app = create_server_app(runtime)
-    server = TestServer(app)
-    await server.start_server()
-    client = RemoteBoxRuntimeClient(
-        base_url=str(server.make_url('')),
-        logger=_logger,
-    )
+    await runtime.initialize()
+    client, server_task, client_task = await _make_rpc_pair(runtime)
     yield client
-    await client.shutdown()
-    await server.close()
+    server_task.cancel()
+    client_task.cancel()
+    await runtime.shutdown()
 
 
 # ── 1. Simple command execution ───────────────────────────────────────
@@ -102,7 +141,7 @@ async def box_client():
 @requires_container
 @requires_socket
 @pytest.mark.asyncio
-async def test_exec_simple_command(box_client: RemoteBoxRuntimeClient):
+async def test_exec_simple_command(box_client: ActionRPCBoxClient):
     """Box starts a simple command and returns stdout."""
     spec = BoxSpec(
         cmd='echo hello-box',
@@ -123,7 +162,7 @@ async def test_exec_simple_command(box_client: RemoteBoxRuntimeClient):
 @requires_container
 @requires_socket
 @pytest.mark.asyncio
-async def test_session_persists_files(box_client: RemoteBoxRuntimeClient):
+async def test_session_persists_files(box_client: ActionRPCBoxClient):
     """Write a file in one exec, read it back in a second exec on the same session."""
     sid = 'int-persist'
 
@@ -151,7 +190,7 @@ async def test_session_persists_files(box_client: RemoteBoxRuntimeClient):
 @requires_container
 @requires_socket
 @pytest.mark.asyncio
-async def test_timeout_kills_command(box_client: RemoteBoxRuntimeClient):
+async def test_timeout_kills_command(box_client: ActionRPCBoxClient):
     """A long-running command is killed after timeout_sec."""
     session_id = 'int-timeout'
     spec = BoxSpec(
@@ -176,7 +215,7 @@ async def test_timeout_kills_command(box_client: RemoteBoxRuntimeClient):
 @requires_container
 @requires_socket
 @pytest.mark.asyncio
-async def test_offline_cannot_reach_network(box_client: RemoteBoxRuntimeClient):
+async def test_offline_cannot_reach_network(box_client: ActionRPCBoxClient):
     """With network=OFF the sandbox cannot reach the internet."""
     spec = BoxSpec(
         cmd='wget -q -O /dev/null --timeout=3 http://1.1.1.1 2>&1; exit $?',
@@ -217,16 +256,11 @@ class _UnavailableBackend(BaseSandboxBackend):
 @requires_socket
 @pytest.mark.asyncio
 async def test_backend_unavailable_returns_error():
-    """When no backend is available the full HTTP path returns BoxBackendUnavailableError."""
+    """When no backend is available the full RPC path returns BoxBackendUnavailableError."""
     runtime = BoxRuntime(logger=_logger, backends=[_UnavailableBackend()])
-    app = create_server_app(runtime)
-    server = TestServer(app)
-    await server.start_server()
+    await runtime.initialize()
+    client, server_task, client_task = await _make_rpc_pair(runtime)
     try:
-        client = RemoteBoxRuntimeClient(
-            base_url=str(server.make_url('')),
-            logger=_logger,
-        )
         spec = BoxSpec(
             cmd='echo hello',
             session_id='int-no-backend',
@@ -234,46 +268,24 @@ async def test_backend_unavailable_returns_error():
         )
         with pytest.raises(BoxBackendUnavailableError):
             await client.execute(spec)
-        await client.shutdown()
     finally:
-        await server.close()
+        server_task.cancel()
+        client_task.cancel()
+        await runtime.shutdown()
 
 
-# ── 6. Runtime unreachable ────────────────────────────────────────────
-
-
-@requires_socket
-@pytest.mark.asyncio
-async def test_runtime_unreachable_returns_error():
-    """Connecting to a non-existent runtime raises BoxRuntimeUnavailableError."""
-    client = RemoteBoxRuntimeClient(
-        base_url='http://127.0.0.1:19999',
-        logger=_logger,
-    )
-    try:
-        with pytest.raises(BoxRuntimeUnavailableError):
-            await client.initialize()
-    finally:
-        await client.shutdown()
-
-
-# ── 7. Full service-to-runtime path ──────────────────────────────────
+# ── 6. Full service-to-runtime path ──────────────────────────────────
 
 
 @requires_container
 @requires_socket
 @pytest.mark.asyncio
 async def test_full_service_to_remote_runtime(tmp_path):
-    """BoxService -> RemoteBoxRuntimeClient -> HTTP -> BoxRuntime -> real backend."""
+    """BoxService -> ActionRPCBoxClient -> RPC -> BoxRuntime -> real backend."""
     runtime = BoxRuntime(logger=_logger)
-    app = create_server_app(runtime)
-    server = TestServer(app)
-    await server.start_server()
+    await runtime.initialize()
+    client, server_task, client_task = await _make_rpc_pair(runtime)
     try:
-        client = RemoteBoxRuntimeClient(
-            base_url=str(server.make_url('')),
-            logger=_logger,
-        )
         host_dir = tmp_path / 'workspace'
         host_dir.mkdir()
 
@@ -303,6 +315,7 @@ async def test_full_service_to_remote_runtime(tmp_path):
         assert result['status'] == 'completed'
         assert 'service-path' in result['stdout']
         assert result['session_id'] == '42'
-        await client.shutdown()
     finally:
-        await server.close()
+        server_task.cancel()
+        client_task.cancel()
+        await runtime.shutdown()

tests/integration_tests/box/test_box_mcp_integration.py

@@ -20,13 +20,14 @@ import subprocess
 
 import aiohttp
 import pytest
+from aiohttp import web
 from aiohttp.test_utils import TestServer
 
-from langbot.pkg.box.client import RemoteBoxRuntimeClient
+from langbot.pkg.box.client import ActionRPCBoxClient
 from langbot.pkg.box.errors import BoxSessionNotFoundError
 from langbot.pkg.box.models import BoxManagedProcessSpec, BoxManagedProcessStatus, BoxSpec
 from langbot.pkg.box.runtime import BoxRuntime
-from langbot.pkg.box.server import create_app as create_server_app
+from langbot.pkg.box.server import BoxServerHandler, create_ws_relay_app
 
 _logger = logging.getLogger('test.box.mcp_integration')
 
@@ -69,23 +70,71 @@ requires_socket = pytest.mark.skipif(
 )
 
 
+# ── Helpers ──────────────────────────────────────────────────────────
+
+
+class _QueueConnection:
+    """In-process Connection backed by asyncio Queues — no real IO."""
+
+    def __init__(self, rx: asyncio.Queue[str], tx: asyncio.Queue[str]):
+        self._rx = rx
+        self._tx = tx
+
+    async def send(self, message: str) -> None:
+        await self._tx.put(message)
+
+    async def receive(self) -> str:
+        return await self._rx.get()
+
+    async def close(self) -> None:
+        pass
+
+
+async def _make_rpc_pair(runtime: BoxRuntime):
+    """Create an in-process RPC pair connected via queues."""
+    from langbot_plugin.runtime.io.handler import Handler
+
+    c2s: asyncio.Queue[str] = asyncio.Queue()
+    s2c: asyncio.Queue[str] = asyncio.Queue()
+    client_conn = _QueueConnection(rx=s2c, tx=c2s)
+    server_conn = _QueueConnection(rx=c2s, tx=s2c)
+
+    server_handler = BoxServerHandler(server_conn, runtime)
+    server_task = asyncio.create_task(server_handler.run())
+
+    client_handler = Handler.__new__(Handler)
+    Handler.__init__(client_handler, client_conn)
+    client_task = asyncio.create_task(client_handler.run())
+
+    client = ActionRPCBoxClient(logger=_logger)
+    client.set_handler(client_handler)
+
+    return client, server_task, client_task
+
+
 # ── Fixtures ──────────────────────────────────────────────────────────
 
 
 @pytest.fixture
 async def box_server():
-    """Yield a (TestServer, RemoteBoxRuntimeClient) backed by a real BoxRuntime."""
+    """Yield a (ws_relay_url, ActionRPCBoxClient) backed by a real BoxRuntime."""
     runtime = BoxRuntime(logger=_logger)
-    app = create_server_app(runtime)
-    server = TestServer(app)
-    await server.start_server()
-    client = RemoteBoxRuntimeClient(
-        base_url=str(server.make_url('')),
-        logger=_logger,
-    )
-    yield server, client
-    await client.shutdown()
-    await server.close()
+    await runtime.initialize()
+
+    # Start ws relay for managed process attach
+    ws_app = create_ws_relay_app(runtime)
+    ws_server = TestServer(ws_app)
+    await ws_server.start_server()
+
+    client, server_task, client_task = await _make_rpc_pair(runtime)
+
+    ws_relay_url = str(ws_server.make_url(''))
+    yield ws_relay_url, client
+
+    server_task.cancel()
+    client_task.cancel()
+    await runtime.shutdown()
+    await ws_server.close()
 
 
 # ── 1. Managed process lifecycle ─────────────────────────────────────
@@ -96,7 +145,7 @@ async def box_server():
 @pytest.mark.asyncio
 async def test_managed_process_start_and_query(box_server):
     """Start a managed process and query its status."""
-    server, client = box_server
+    ws_relay_url, client = box_server
 
     # Create session
     spec = BoxSpec(
@@ -133,7 +182,7 @@ async def test_managed_process_start_and_query(box_server):
 @pytest.mark.asyncio
 async def test_ws_stdio_attach_echo(box_server):
     """Attach to a managed process via WebSocket and verify bidirectional IO."""
-    server, client = box_server
+    ws_relay_url, client = box_server
 
     spec = BoxSpec(
         cmd='',
@@ -151,8 +200,8 @@ async def test_ws_stdio_attach_echo(box_server):
     )
     await client.start_managed_process('mcp-int-ws', proc_spec)
 
-    # Connect via WebSocket
-    ws_url = client.get_managed_process_websocket_url('mcp-int-ws')
+    # Connect via WebSocket (ws relay)
+    ws_url = client.get_managed_process_websocket_url('mcp-int-ws', ws_relay_url)
     session = aiohttp.ClientSession()
     try:
         async with session.ws_connect(ws_url) as ws:
@@ -177,7 +226,7 @@ async def test_ws_stdio_attach_echo(box_server):
 @pytest.mark.asyncio
 async def test_delete_session_cleans_up(box_server):
     """After deleting a session, it should no longer exist."""
-    server, client = box_server
+    ws_relay_url, client = box_server
 
     spec = BoxSpec(
         cmd='',
@@ -203,15 +252,15 @@ async def test_delete_session_cleans_up(box_server):
         await client.get_session('mcp-int-cleanup')
 
 
-# ── 4. GET /v1/sessions/{id} ────────────────────────────────────────
+# ── 4. GET session details ────────────────────────────────────────
 
 
 @requires_container
 @requires_socket
 @pytest.mark.asyncio
 async def test_get_session_returns_details(box_server):
-    """GET single session returns session details and managed process info."""
-    server, client = box_server
+    """Get single session returns session details and managed process info."""
+    ws_relay_url, client = box_server
 
     spec = BoxSpec(
         cmd='',
@@ -251,7 +300,7 @@ async def test_get_session_returns_details(box_server):
 @pytest.mark.asyncio
 async def test_process_exit_detected(box_server):
     """When a managed process exits, its status should reflect EXITED."""
-    server, client = box_server
+    ws_relay_url, client = box_server
 
     spec = BoxSpec(
         cmd='',
@@ -287,7 +336,7 @@ async def test_process_exit_detected(box_server):
 @pytest.mark.asyncio
 async def test_orphan_cleanup_preserves_own_containers(box_server):
     """Orphan cleanup should not remove containers belonging to the current instance."""
-    server, client = box_server
+    ws_relay_url, client = box_server
 
     # Create a session (container gets current instance ID label)
     spec = BoxSpec(