mirror of
https://github.com/usestrix/strix.git
synced 2026-08-22 02:58:39 +02:00
feat(tui): replace Textual with a Go/Bubble Tea interface (#941)
This commit is contained in:
oyasumi
@@ -0,0 +1,560 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import socket
|
||||
import struct
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
from agents.tool import ToolOutputImage
|
||||
|
||||
from strix.config.settings import DEFAULT_MAX_TURNS
|
||||
from strix.interface.tui.backend.controller import TuiController
|
||||
from strix.interface.tui.backend.projection import terminal_projection
|
||||
from strix.interface.tui.backend.protocol import (
|
||||
MAX_COMMAND_BYTES,
|
||||
PROTOCOL_CAPABILITIES,
|
||||
PROTOCOL_VERSION,
|
||||
ProtocolHandshakeError,
|
||||
envelope,
|
||||
)
|
||||
from strix.interface.tui.backend.server import TuiBackendServer
|
||||
from strix.interface.tui.live_view import TuiLiveView
|
||||
|
||||
|
||||
def args() -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
needs_setup=True,
|
||||
targets_info=[],
|
||||
instruction=None,
|
||||
scan_mode="deep",
|
||||
max_budget_usd=None,
|
||||
max_turns=DEFAULT_MAX_TURNS,
|
||||
scope_mode="auto",
|
||||
diff_base=None,
|
||||
local_sources=[],
|
||||
diff_scope={"active": False},
|
||||
user_explicit_instruction=None,
|
||||
run_name=None,
|
||||
)
|
||||
|
||||
|
||||
async def send_message(connection: socket.socket, message: dict[str, object]) -> None:
|
||||
raw = json.dumps(message).encode()
|
||||
await asyncio.get_running_loop().sock_sendall(connection, struct.pack(">I", len(raw)) + raw)
|
||||
|
||||
|
||||
async def receive_exactly(connection: socket.socket, size: int) -> bytes:
|
||||
chunks: list[bytes] = []
|
||||
while size:
|
||||
chunk = await asyncio.get_running_loop().sock_recv(connection, size)
|
||||
if not chunk:
|
||||
raise EOFError
|
||||
chunks.append(chunk)
|
||||
size -= len(chunk)
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
async def receive_frame(connection: socket.socket) -> tuple[int, dict[str, Any]]:
|
||||
size = struct.unpack(">I", await receive_exactly(connection, 4))[0]
|
||||
value = json.loads(await receive_exactly(connection, size))
|
||||
assert isinstance(value, dict)
|
||||
return size, value
|
||||
|
||||
|
||||
async def receive_message(connection: socket.socket) -> dict[str, Any]:
|
||||
return (await receive_frame(connection))[1]
|
||||
|
||||
|
||||
async def start_server(
|
||||
server: TuiBackendServer, backend: socket.socket, child: socket.socket
|
||||
) -> dict[str, Any]:
|
||||
start_task = asyncio.create_task(server.start(backend))
|
||||
hello = await receive_message(child)
|
||||
await send_message(
|
||||
child,
|
||||
{
|
||||
"version": PROTOCOL_VERSION,
|
||||
"type": "ready",
|
||||
"payload": {"capabilities": list(PROTOCOL_CAPABILITIES)},
|
||||
},
|
||||
)
|
||||
await asyncio.wait_for(start_task, timeout=1)
|
||||
return hello
|
||||
|
||||
|
||||
async def receive_until(
|
||||
connection: socket.socket,
|
||||
message_type: str,
|
||||
*,
|
||||
request_id: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
for _ in range(100):
|
||||
message = await asyncio.wait_for(receive_message(connection), timeout=2)
|
||||
if message.get("type") != message_type:
|
||||
continue
|
||||
if request_id is not None and message.get("request_id") != request_id:
|
||||
continue
|
||||
return message
|
||||
raise AssertionError(f"did not receive {message_type}")
|
||||
|
||||
|
||||
async def receive_initial_state(connection: socket.socket) -> None:
|
||||
state_received = False
|
||||
complete: set[str] = set()
|
||||
while not state_received or complete != {"agents", "events", "vulnerabilities"}:
|
||||
message = await asyncio.wait_for(receive_message(connection), timeout=2)
|
||||
if message["type"] == "state":
|
||||
state_received = True
|
||||
elif message["type"] == "collection_bootstrap":
|
||||
payload = message["payload"]
|
||||
if payload["done"]:
|
||||
complete.add(payload["collection"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_server_requires_ready_before_state_or_commands() -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
server = TuiBackendServer(TuiController(args()))
|
||||
start_task = asyncio.create_task(server.start(backend))
|
||||
try:
|
||||
hello = await receive_message(child)
|
||||
assert hello == {
|
||||
"version": 3,
|
||||
"type": "hello",
|
||||
"payload": {"capabilities": list(PROTOCOL_CAPABILITIES)},
|
||||
}
|
||||
with pytest.raises(TimeoutError):
|
||||
await asyncio.wait_for(receive_message(child), timeout=0.1)
|
||||
assert not start_task.done()
|
||||
|
||||
await send_message(
|
||||
child,
|
||||
{
|
||||
"version": 3,
|
||||
"type": "ready",
|
||||
"payload": {"capabilities": list(PROTOCOL_CAPABILITIES)},
|
||||
},
|
||||
)
|
||||
await asyncio.wait_for(start_task, timeout=1)
|
||||
assert server.activated is True
|
||||
assert (await receive_until(child, "state"))["payload"]["revision"] == 1
|
||||
finally:
|
||||
child.close()
|
||||
start_task.cancel()
|
||||
await server.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("version", "capabilities"),
|
||||
[
|
||||
(2, list(PROTOCOL_CAPABILITIES)),
|
||||
(3, ["state-revisions"]),
|
||||
],
|
||||
)
|
||||
async def test_server_rejects_handshake_mismatch(version: int, capabilities: list[str]) -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
server = TuiBackendServer(TuiController(args()))
|
||||
start_task = asyncio.create_task(server.start(backend))
|
||||
try:
|
||||
await receive_message(child)
|
||||
await send_message(
|
||||
child,
|
||||
{"version": version, "type": "ready", "payload": {"capabilities": capabilities}},
|
||||
)
|
||||
with pytest.raises(ProtocolHandshakeError, match="mismatch"):
|
||||
await asyncio.wait_for(start_task, timeout=1)
|
||||
assert server.activated is False
|
||||
finally:
|
||||
child.close()
|
||||
await server.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_server_command_round_trip_over_inherited_socket() -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
server = TuiBackendServer(TuiController(args()))
|
||||
await start_server(server, backend, child)
|
||||
try:
|
||||
await send_message(
|
||||
child,
|
||||
{
|
||||
"version": 3,
|
||||
"type": "setup.add_target",
|
||||
"request_id": "test-1",
|
||||
"payload": {"target": "example.com"},
|
||||
},
|
||||
)
|
||||
result = await receive_until(child, "command_result", request_id="test-1")
|
||||
assert result["payload"]["ok"] is True
|
||||
assert result["payload"]["command"] == "setup.add_target"
|
||||
state = await receive_until(child, "state")
|
||||
assert state["payload"]["revision"] >= 1
|
||||
assert state["payload"]["state"]["targets"] == ["example.com"]
|
||||
finally:
|
||||
child.close()
|
||||
await server.close()
|
||||
|
||||
|
||||
def test_unicode_heavy_setup_state_stays_within_control_frame_limit() -> None:
|
||||
controller = TuiController(args())
|
||||
controller.instruction = "🔒" * 10_000
|
||||
controller.targets = [f"https://例え.{index}/" + "界" * 500 for index in range(20)]
|
||||
controller.error = "失" * 10_000
|
||||
controller.messages = [
|
||||
{"id": str(index), "text": "警" * 10_000, "level": "warning"} for index in range(10)
|
||||
]
|
||||
controller.report_state = cast(
|
||||
"Any",
|
||||
SimpleNamespace(
|
||||
caido_url="https://例え.example/" + "道" * 10_000,
|
||||
get_total_llm_usage=lambda: {f"model-{index}": "費" * 10_000 for index in range(20)},
|
||||
),
|
||||
)
|
||||
server = TuiBackendServer(controller)
|
||||
|
||||
snapshot = controller.snapshot()
|
||||
encoded = server._encode(envelope("state", {"revision": 1, "state": snapshot}))
|
||||
|
||||
assert len(encoded) <= MAX_COMMAND_BYTES
|
||||
assert "🔒".encode() in encoded
|
||||
assert snapshot["projection_truncated"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_persistence_error_does_not_kill_command_reader(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
controller = TuiController(args())
|
||||
calls = 0
|
||||
|
||||
async def handle(command: str, payload: dict[str, Any]) -> dict[str, Any]:
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 1:
|
||||
raise OSError("disk is read-only")
|
||||
return {"command": command, "payload": payload}
|
||||
|
||||
monkeypatch.setattr(controller, "handle", handle)
|
||||
server = TuiBackendServer(controller)
|
||||
await start_server(server, backend, child)
|
||||
try:
|
||||
for request_id in ("persist-1", "persist-2"):
|
||||
await send_message(
|
||||
child,
|
||||
{
|
||||
"version": 3,
|
||||
"type": "setup.select_model",
|
||||
"request_id": request_id,
|
||||
"payload": {"provider": "openai", "model": "openai/gpt-5"},
|
||||
},
|
||||
)
|
||||
result = await receive_until(child, "command_result", request_id=request_id)
|
||||
if request_id == "persist-1":
|
||||
assert result["payload"]["error"] == {
|
||||
"code": "persistence_error",
|
||||
"message": "disk is read-only",
|
||||
"retryable": True,
|
||||
}
|
||||
else:
|
||||
assert result["payload"]["ok"] is True
|
||||
assert server._reader_task is not None and not server._reader_task.done()
|
||||
finally:
|
||||
child.close()
|
||||
await server.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_version_error_is_correlated_and_next_command_succeeds() -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
server = TuiBackendServer(TuiController(args()))
|
||||
await start_server(server, backend, child)
|
||||
try:
|
||||
await send_message(
|
||||
child,
|
||||
{
|
||||
"version": 2,
|
||||
"type": "setup.add_target",
|
||||
"request_id": "bad-version",
|
||||
"payload": {"target": "ignored.example"},
|
||||
},
|
||||
)
|
||||
rejected = await receive_until(child, "command_result", request_id="bad-version")
|
||||
assert rejected["payload"]["error"]["code"] == "invalid_request"
|
||||
|
||||
await send_message(
|
||||
child,
|
||||
{
|
||||
"version": 3,
|
||||
"type": "setup.add_target",
|
||||
"request_id": "after-error",
|
||||
"payload": {"target": "example.com"},
|
||||
},
|
||||
)
|
||||
accepted = await receive_until(child, "command_result", request_id="after-error")
|
||||
assert accepted["payload"]["ok"] is True
|
||||
finally:
|
||||
child.close()
|
||||
await server.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_collection_bootstrap_is_chunked_deltas_are_incremental_and_idle_is_silent() -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
controller = TuiController(args())
|
||||
report_state = SimpleNamespace(
|
||||
vulnerability_reports=[],
|
||||
caido_url=None,
|
||||
get_total_llm_usage=dict,
|
||||
)
|
||||
controller.report_state = cast("Any", report_state)
|
||||
content = "x" * (64 * 1024)
|
||||
for index in range(80):
|
||||
controller.live_view.record_user_message(f"agent-{index}", content)
|
||||
server = TuiBackendServer(controller)
|
||||
await start_server(server, backend, child)
|
||||
try:
|
||||
event_frames = 0
|
||||
event_count = 0
|
||||
complete: set[str] = set()
|
||||
state_received = False
|
||||
while not (complete == {"agents", "events", "vulnerabilities"} and state_received):
|
||||
size, message = await asyncio.wait_for(receive_frame(child), timeout=5)
|
||||
if message["type"] == "state":
|
||||
state_received = True
|
||||
if message["type"] != "collection_bootstrap":
|
||||
continue
|
||||
payload = message["payload"]
|
||||
if payload["collection"] == "events":
|
||||
event_frames += 1
|
||||
event_count += len(payload["items"])
|
||||
assert size <= 4 * 1024 * 1024
|
||||
if payload["done"]:
|
||||
complete.add(payload["collection"])
|
||||
assert event_frames >= 2
|
||||
assert event_count == 80
|
||||
|
||||
server.notify_changed()
|
||||
with pytest.raises(TimeoutError):
|
||||
await asyncio.wait_for(receive_message(child), timeout=0.2)
|
||||
|
||||
controller.live_view.record_user_message("agent-new", "delta")
|
||||
controller.notify_changed()
|
||||
delta = await receive_until(child, "collection_delta")
|
||||
assert delta["payload"]["collection"] == "events"
|
||||
assert delta["payload"]["base_revision"] == 1
|
||||
assert len(delta["payload"]["operations"]) == 1
|
||||
|
||||
report_state.vulnerability_reports.append(
|
||||
{"id": "vuln-0001", "title": "Incremental finding", "severity": "high"}
|
||||
)
|
||||
controller.notify_changed()
|
||||
finding_delta = await receive_until(child, "collection_delta")
|
||||
assert finding_delta["payload"]["collection"] == "vulnerabilities"
|
||||
assert len(finding_delta["payload"]["operations"]) == 1
|
||||
finally:
|
||||
child.close()
|
||||
await server.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agents_collection_has_no_state_cap_and_sends_delete_and_resync() -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
controller = TuiController(args())
|
||||
for index in range(40):
|
||||
controller.live_view.upsert_agent(
|
||||
f"agent-{index}",
|
||||
name=f"Agent {index}",
|
||||
status="running",
|
||||
)
|
||||
server = TuiBackendServer(controller)
|
||||
await start_server(server, backend, child)
|
||||
try:
|
||||
agents: list[dict[str, Any]] = []
|
||||
complete: set[str] = set()
|
||||
state: dict[str, Any] | None = None
|
||||
while state is None or complete != {"agents", "events", "vulnerabilities"}:
|
||||
message = await asyncio.wait_for(receive_message(child), timeout=2)
|
||||
if message["type"] == "state":
|
||||
state = message["payload"]["state"]
|
||||
elif message["type"] == "collection_bootstrap":
|
||||
payload = message["payload"]
|
||||
if payload["collection"] == "agents":
|
||||
agents.extend(payload["items"])
|
||||
if payload["done"]:
|
||||
complete.add(payload["collection"])
|
||||
|
||||
assert "agents" not in state
|
||||
assert len(agents) == 40
|
||||
|
||||
controller.live_view.agents.pop("agent-7")
|
||||
controller.notify_changed()
|
||||
delta = await receive_until(child, "collection_delta")
|
||||
assert delta["payload"]["collection"] == "agents"
|
||||
assert delta["payload"]["operations"] == [{"op": "delete", "id": "agent-7"}]
|
||||
|
||||
await send_message(
|
||||
child,
|
||||
{
|
||||
"version": 3,
|
||||
"type": "collection.resync",
|
||||
"request_id": "resync-agents",
|
||||
"payload": {"collection": "agents"},
|
||||
},
|
||||
)
|
||||
result = await receive_until(child, "command_result", request_id="resync-agents")
|
||||
assert result["payload"]["ok"] is True
|
||||
bootstrap = await receive_until(child, "collection_bootstrap")
|
||||
assert bootstrap["payload"]["collection"] == "agents"
|
||||
assert bootstrap["payload"]["revision"] == 3
|
||||
assert len(bootstrap["payload"]["items"]) == 39
|
||||
finally:
|
||||
child.close()
|
||||
await server.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bootstrap_larger_than_64_mib_has_no_total_message_ceiling(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
server = TuiBackendServer(TuiController(args()))
|
||||
shared_projection = "x" * (1024 * 1024)
|
||||
items = [{"id": f"event-{index}", "content": shared_projection} for index in range(65)]
|
||||
frames: list[dict[str, Any]] = []
|
||||
encoded_sizes: list[int] = []
|
||||
|
||||
async def capture(message: dict[str, Any]) -> None:
|
||||
encoded_sizes.append(len(server._encode(message)))
|
||||
frames.append(message)
|
||||
|
||||
monkeypatch.setattr(server, "_send", capture)
|
||||
|
||||
await server._send_collection_frames(
|
||||
"collection_bootstrap",
|
||||
{"collection": "events", "revision": 1},
|
||||
"items",
|
||||
items,
|
||||
)
|
||||
|
||||
assert sum(len(item["content"]) for item in items) > 64 * 1024 * 1024
|
||||
assert len(frames) > 16
|
||||
assert max(encoded_sizes) <= 4 * 1024 * 1024
|
||||
assert frames[0]["payload"]["cursor"] == 0
|
||||
assert frames[-1]["payload"]["next_cursor"] == len(items)
|
||||
assert frames[-1]["payload"]["done"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_terminal_projection_is_truncated_without_mutating_history() -> None:
|
||||
controller = TuiController(args())
|
||||
durable = "x" * (2 * 1024 * 1024)
|
||||
controller.live_view.record_user_message("agent", durable)
|
||||
|
||||
projected = controller.collection("events")
|
||||
|
||||
assert len(projected[0]["data"]["content"]) < len(durable)
|
||||
assert controller.live_view.events[0]["data"]["content"] == durable
|
||||
|
||||
|
||||
def test_terminal_projection_strips_ansi_osc_and_c1_controls() -> None:
|
||||
controller = TuiController(args())
|
||||
hostile = "safe\x1b[31mred\x1b[0m\x1b]52;c;Y2xpcGJvYXJk\x07\x85tail"
|
||||
controller.live_view.record_user_message("agent", hostile)
|
||||
|
||||
projected = controller.collection_snapshot("events")[1][0]["data"]["content"]
|
||||
|
||||
assert projected == "saferedtail"
|
||||
assert "\x1b" not in projected
|
||||
|
||||
hostile_mapping = {"header\x1b]52;c;Y2xpcA==\x07": "value"}
|
||||
assert list(terminal_projection(hostile_mapping)) == ["header"]
|
||||
assert list(TuiBackendServer._sanitize_wire_value(hostile_mapping)) == ["header"]
|
||||
|
||||
|
||||
def test_terminal_event_history_is_bounded_without_changing_durable_sessions() -> None:
|
||||
controller = TuiController(args())
|
||||
for index in range(10_050):
|
||||
controller.live_view.record_user_message("agent", f"message-{index}")
|
||||
|
||||
_cursor, projected = controller.collection_snapshot("events")
|
||||
|
||||
assert len(controller.live_view.events) == 10_000
|
||||
assert len(projected) == 5_000
|
||||
assert projected[0]["data"]["content"] == "message-5050"
|
||||
assert projected[-1]["data"]["content"] == "message-10049"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_oversized_command_frame_is_rejected_before_payload_read() -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
server = TuiBackendServer(TuiController(args()))
|
||||
await start_server(server, backend, child)
|
||||
try:
|
||||
await asyncio.get_running_loop().sock_sendall(
|
||||
child, struct.pack(">I", MAX_COMMAND_BYTES + 1)
|
||||
)
|
||||
assert server._reader_task is not None
|
||||
await asyncio.wait_for(server._reader_task, timeout=1)
|
||||
assert server._socket is None
|
||||
finally:
|
||||
child.close()
|
||||
await server.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_server_stops_when_peer_closes() -> None:
|
||||
backend, child = socket.socketpair()
|
||||
child.setblocking(False) # noqa: FBT003
|
||||
server = TuiBackendServer(TuiController(args()))
|
||||
await start_server(server, backend, child)
|
||||
child.close()
|
||||
try:
|
||||
assert server._reader_task is not None
|
||||
await asyncio.wait_for(server._reader_task, timeout=1)
|
||||
finally:
|
||||
await server.close()
|
||||
|
||||
|
||||
def test_image_data_uri_survives_terminal_projection() -> None:
|
||||
uri = "data:image/png;base64," + "A" * 100_000
|
||||
assert terminal_projection(uri) == uri
|
||||
assert terminal_projection({"type": "image", "image_url": uri})["image_url"] == uri
|
||||
|
||||
oversized = "data:image/png;base64," + "A" * (3 * 1024 * 1024)
|
||||
assert terminal_projection(oversized) == "[image omitted from terminal projection]"
|
||||
|
||||
|
||||
def test_view_image_tool_output_is_normalized_to_image_dict() -> None:
|
||||
uri = "data:image/png;base64," + "B" * 4000
|
||||
view = TuiLiveView()
|
||||
view._record_tool_output_data(
|
||||
"agent",
|
||||
{
|
||||
"call_id": "c1",
|
||||
"tool_name": "view_image",
|
||||
"output": ToolOutputImage(type="image", image_url=uri),
|
||||
},
|
||||
)
|
||||
view._record_tool_output_data(
|
||||
"agent",
|
||||
{
|
||||
"call_id": "c2",
|
||||
"tool_name": "view_image",
|
||||
"output": [{"type": "input_image", "image_url": uri}],
|
||||
},
|
||||
)
|
||||
for event in view.events:
|
||||
assert event["data"]["result"] == {"type": "image", "image_url": uri}
|
||||
Reference in New Issue
Block a user