from __future__ import annotations import asyncio from collections import deque import pytest from investibook.clients.run9 import Run9OutputChunk from investibook.core.run9_process import ( Run9ProcessTransport, Run9StreamWriter, build_run9_box_id, ) class FakeRun9Client: def __init__(self, chunks: list[Run9OutputChunk] | None = None) -> None: self.chunks = deque(chunks or []) self.stdin_writes: list[tuple[bytes, bool]] = [] self.started: dict | None = None self.killed: list[str] = [] self.closed = False async def start_background_exec(self, **kwargs): self.started = kwargs return {"exec_id": "exec-1", "state": "running"} async def pull_output(self, **kwargs): return self.chunks.popleft() async def write_stdin(self, *, exec_id, data, close_stdin=False): assert exec_id == "exec-1" self.stdin_writes.append((data, close_stdin)) async def kill_exec(self, exec_id): self.killed.append(exec_id) async def aclose(self): self.closed = True def test_box_id_is_stable_per_workspace_and_independent_of_runtime_owner(): first = build_run9_box_id( workspace_id="ws-123", environment="prod" ) second = build_run9_box_id( workspace_id="ws-123", environment="prod" ) assert first == second assert first == "ib-prod-ws-123" def test_process_pull_wait_can_switch_after_handshake(): proc = object.__new__(Run9ProcessTransport) proc._pull_wait_ms = 50 proc.set_pull_wait_ms(500) assert proc._pull_wait_ms == 500 @pytest.mark.asyncio async def test_stream_writer_flushes_exact_bytes_on_drain(): client = FakeRun9Client() writer = Run9StreamWriter(client=client, exec_id="exec-1") # type: ignore[arg-type] assert isinstance(writer, asyncio.StreamWriter) writer.write(b"first") writer.write(memoryview(b"-second")) await writer.drain() writer.close() await writer.wait_closed() assert client.stdin_writes == [(b"first-second", False), (b"", True)] @pytest.mark.asyncio async def test_stream_writer_fails_closed_when_run9_write_fails(): class FailingClient(FakeRun9Client): async def write_stdin(self, *, exec_id, data, close_stdin=False): if not self.stdin_writes: self.stdin_writes.append((data, close_stdin)) raise RuntimeError("temporary failure") await super().write_stdin( exec_id=exec_id, data=data, close_stdin=close_stdin, ) client = FailingClient() writer = Run9StreamWriter(client=client, exec_id="exec-1") # type: ignore[arg-type] writer.write(b"message") with pytest.raises(RuntimeError, match="temporary failure"): await writer.drain() with pytest.raises(ConnectionError, match="previously failed"): await writer.drain() with pytest.raises(ConnectionError, match="previously failed"): writer.write(b"another-message") assert client.stdin_writes == [(b"message", False)] @pytest.mark.asyncio async def test_process_preserves_output_chunk_boundaries_as_one_byte_stream(): client = FakeRun9Client( [ Run9OutputChunk( data=b'{"jsonrpc":"2.0","result":"\xe4\xbd', cursor="1", state="running", ), Run9OutputChunk( data=b'\xa0\xe5\xa5\xbd"}\n', cursor="2", state="succeeded", exit_code=0, ), ] ) proc = await Run9ProcessTransport.start( client=client, # type: ignore[arg-type] box_id="box-1", command=["opencode", "acp"], workdir="/workspace", pull_wait_ms=1, ) line = await proc.stdout.readline() assert line.decode("utf-8") == '{"jsonrpc":"2.0","result":"你好"}\n' assert await proc.wait() == 0 assert proc.returncode == 0 assert client.started is not None command = client.started["command"] assert command[-2:] == ["opencode", "acp"] assert "2>>" in command[2] await proc.aclose() assert client.closed is True @pytest.mark.asyncio async def test_process_rejects_exec_that_finishes_before_running(): client = FakeRun9Client( [Run9OutputChunk(data=b"", cursor="1", state="succeeded", exit_code=0)] ) with pytest.raises(RuntimeError, match="exited before entering running state"): await Run9ProcessTransport.start( client=client, # type: ignore[arg-type] box_id="box-1", command=["opencode", "acp"], workdir="/workspace", pull_wait_ms=1, ) assert client.closed is True @pytest.mark.asyncio async def test_process_can_wait_for_remote_startup_output(): class ReadyClient(FakeRun9Client): def __init__(self): super().__init__( [ Run9OutputChunk(data=b"", cursor="1", state="pending"), Run9OutputChunk( data=b'{"type":"run9_transport_ready"}\n', cursor="2", state="running", ), ] ) self.finish = asyncio.Event() async def pull_output(self, **kwargs): if self.chunks: return self.chunks.popleft() await self.finish.wait() return Run9OutputChunk( data=b"", cursor="3", state="cancelled", exit_code=1 ) async def kill_exec(self, exec_id): await super().kill_exec(exec_id) self.finish.set() client = ReadyClient() proc = await Run9ProcessTransport.start( client=client, # type: ignore[arg-type] box_id="box-1", command=["node", "pi-core.mjs"], workdir="/workspace", pull_wait_ms=1, wait_for_first_output=True, ) assert await proc.stdout.readline() == b'{"type":"run9_transport_ready"}\n' assert proc.startup_timings["start_exec_ms"] >= 0 assert proc.startup_timings["first_output_ms"] >= 0 await proc.aclose() @pytest.mark.asyncio async def test_process_prioritizes_pending_stdin_before_next_output_pull(): class ContendedClient(FakeRun9Client): def __init__(self): super().__init__() self.pull_count = 0 self.second_pull_started = asyncio.Event() self.release_second_pull = asyncio.Event() self.finish = asyncio.Event() self.order: list[str] = [] async def pull_output(self, **kwargs): self.pull_count += 1 if self.pull_count == 1: return Run9OutputChunk( data=b'{"type":"ready"}\n', cursor="1", state="running" ) if self.pull_count == 2: self.second_pull_started.set() await self.release_second_pull.wait() return Run9OutputChunk(data=b"", cursor="2", state="running") self.order.append("pull") await self.finish.wait() return Run9OutputChunk( data=b"", cursor="3", state="cancelled", exit_code=1 ) async def write_stdin(self, *, exec_id, data, close_stdin=False): self.order.append("write") await super().write_stdin( exec_id=exec_id, data=data, close_stdin=close_stdin, ) async def kill_exec(self, exec_id): await super().kill_exec(exec_id) self.finish.set() client = ContendedClient() proc = await Run9ProcessTransport.start( client=client, # type: ignore[arg-type] box_id="box-1", command=["node", "pi-core.mjs"], workdir="/workspace", pull_wait_ms=10_000, ) assert await proc.stdout.readline() == b'{"type":"ready"}\n' await client.second_pull_started.wait() proc.stdin.write(b'{"type":"get_state"}\n') drain_task = asyncio.create_task(proc.stdin.drain()) client.release_second_pull.set() await drain_task assert client.order[0] == "write" await proc.aclose()