Add persistent workspace shell sessions
Let agents inhabit a workspace across separate calls instead of only submitting one-shot execs. Add workspace shell open/read/write/signal/close across the CLI, Python SDK, and MCP server, with persisted shell records, a local PTY-backed mock implementation, and guest-agent support for real Firecracker workspaces. Mark the 2.5.0 roadmap milestone done, refresh docs/examples and the release metadata, and verify with uv lock, UV_CACHE_DIR=.uv-cache make check, and UV_CACHE_DIR=.uv-cache make dist-check.
This commit is contained in:
parent
2de31306b6
commit
3f8293ad24
28 changed files with 3265 additions and 81 deletions
|
|
@ -1,6 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
|
|
@ -50,6 +51,11 @@ def test_pyro_create_server_registers_vm_run(tmp_path: Path) -> None:
|
|||
assert "vm_create" in tool_names
|
||||
assert "workspace_create" in tool_names
|
||||
assert "workspace_sync_push" in tool_names
|
||||
assert "shell_open" in tool_names
|
||||
assert "shell_read" in tool_names
|
||||
assert "shell_write" in tool_names
|
||||
assert "shell_signal" in tool_names
|
||||
assert "shell_close" in tool_names
|
||||
|
||||
|
||||
def test_pyro_vm_run_tool_executes(tmp_path: Path) -> None:
|
||||
|
|
@ -130,6 +136,16 @@ def test_pyro_workspace_methods_delegate_to_manager(tmp_path: Path) -> None:
|
|||
(updated_dir / "more.txt").write_text("more\n", encoding="utf-8")
|
||||
synced = pyro.push_workspace_sync(workspace_id, updated_dir, dest="subdir")
|
||||
executed = pyro.exec_workspace(workspace_id, command="cat note.txt")
|
||||
opened = pyro.open_shell(workspace_id)
|
||||
shell_id = str(opened["shell_id"])
|
||||
written = pyro.write_shell(workspace_id, shell_id, input="pwd")
|
||||
read = pyro.read_shell(workspace_id, shell_id)
|
||||
deadline = time.time() + 5
|
||||
while "/workspace" not in str(read["output"]) and time.time() < deadline:
|
||||
read = pyro.read_shell(workspace_id, shell_id, cursor=0)
|
||||
time.sleep(0.05)
|
||||
signaled = pyro.signal_shell(workspace_id, shell_id)
|
||||
closed = pyro.close_shell(workspace_id, shell_id)
|
||||
status = pyro.status_workspace(workspace_id)
|
||||
logs = pyro.logs_workspace(workspace_id)
|
||||
deleted = pyro.delete_workspace(workspace_id)
|
||||
|
|
@ -137,6 +153,10 @@ def test_pyro_workspace_methods_delegate_to_manager(tmp_path: Path) -> None:
|
|||
assert executed["stdout"] == "ok\n"
|
||||
assert created["workspace_seed"]["mode"] == "directory"
|
||||
assert synced["workspace_sync"]["destination"] == "/workspace/subdir"
|
||||
assert written["input_length"] == 3
|
||||
assert "/workspace" in read["output"]
|
||||
assert signaled["signal"] == "INT"
|
||||
assert closed["closed"] is True
|
||||
assert status["command_count"] == 1
|
||||
assert logs["count"] == 1
|
||||
assert deleted["deleted"] is True
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ def test_cli_subcommand_help_includes_examples_and_guidance() -> None:
|
|||
assert "pyro workspace create debian:12 --seed-path ./repo" in workspace_help
|
||||
assert "pyro workspace sync push WORKSPACE_ID ./repo --dest src" in workspace_help
|
||||
assert "pyro workspace exec WORKSPACE_ID" in workspace_help
|
||||
assert "pyro workspace shell open WORKSPACE_ID" in workspace_help
|
||||
|
||||
workspace_create_help = _subparser_choice(
|
||||
_subparser_choice(parser, "workspace"),
|
||||
|
|
@ -92,6 +93,41 @@ def test_cli_subcommand_help_includes_examples_and_guidance() -> None:
|
|||
assert "--dest" in workspace_sync_push_help
|
||||
assert "Import host content into `/workspace`" in workspace_sync_push_help
|
||||
|
||||
workspace_shell_help = _subparser_choice(
|
||||
_subparser_choice(parser, "workspace"),
|
||||
"shell",
|
||||
).format_help()
|
||||
assert "pyro workspace shell open WORKSPACE_ID" in workspace_shell_help
|
||||
assert "Use `workspace exec` for one-shot commands." in workspace_shell_help
|
||||
|
||||
workspace_shell_open_help = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "open"
|
||||
).format_help()
|
||||
assert "--cwd" in workspace_shell_open_help
|
||||
assert "--cols" in workspace_shell_open_help
|
||||
assert "--rows" in workspace_shell_open_help
|
||||
|
||||
workspace_shell_read_help = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "read"
|
||||
).format_help()
|
||||
assert "Shell output is written to stdout." in workspace_shell_read_help
|
||||
|
||||
workspace_shell_write_help = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "write"
|
||||
).format_help()
|
||||
assert "--input" in workspace_shell_write_help
|
||||
assert "--no-newline" in workspace_shell_write_help
|
||||
|
||||
workspace_shell_signal_help = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "signal"
|
||||
).format_help()
|
||||
assert "--signal" in workspace_shell_signal_help
|
||||
|
||||
workspace_shell_close_help = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "close"
|
||||
).format_help()
|
||||
assert "Close a persistent workspace shell" in workspace_shell_close_help
|
||||
|
||||
|
||||
def test_cli_run_prints_json(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
@ -681,6 +717,226 @@ def test_cli_workspace_status_and_delete_print_json(
|
|||
assert deleted["deleted"] is True
|
||||
|
||||
|
||||
def test_cli_workspace_shell_open_and_read_human(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
class StubPyro:
|
||||
def open_shell(
|
||||
self,
|
||||
workspace_id: str,
|
||||
*,
|
||||
cwd: str,
|
||||
cols: int,
|
||||
rows: int,
|
||||
) -> dict[str, Any]:
|
||||
assert workspace_id == "workspace-123"
|
||||
assert cwd == "/workspace"
|
||||
assert cols == 120
|
||||
assert rows == 30
|
||||
return {
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": "shell-123",
|
||||
"state": "running",
|
||||
"cwd": cwd,
|
||||
"cols": cols,
|
||||
"rows": rows,
|
||||
"started_at": 1.0,
|
||||
"ended_at": None,
|
||||
"exit_code": None,
|
||||
"execution_mode": "guest_vsock",
|
||||
}
|
||||
|
||||
def read_shell(
|
||||
self,
|
||||
workspace_id: str,
|
||||
shell_id: str,
|
||||
*,
|
||||
cursor: int,
|
||||
max_chars: int,
|
||||
) -> dict[str, Any]:
|
||||
assert workspace_id == "workspace-123"
|
||||
assert shell_id == "shell-123"
|
||||
assert cursor == 0
|
||||
assert max_chars == 1024
|
||||
return {
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
"state": "running",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"started_at": 1.0,
|
||||
"ended_at": None,
|
||||
"exit_code": None,
|
||||
"execution_mode": "guest_vsock",
|
||||
"cursor": 0,
|
||||
"next_cursor": 14,
|
||||
"output": "pyro$ pwd\n",
|
||||
"truncated": False,
|
||||
}
|
||||
|
||||
class OpenParser:
|
||||
def parse_args(self) -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
command="workspace",
|
||||
workspace_command="shell",
|
||||
workspace_shell_command="open",
|
||||
workspace_id="workspace-123",
|
||||
cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
json=False,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cli, "_build_parser", lambda: OpenParser())
|
||||
monkeypatch.setattr(cli, "Pyro", StubPyro)
|
||||
cli.main()
|
||||
|
||||
class ReadParser:
|
||||
def parse_args(self) -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
command="workspace",
|
||||
workspace_command="shell",
|
||||
workspace_shell_command="read",
|
||||
workspace_id="workspace-123",
|
||||
shell_id="shell-123",
|
||||
cursor=0,
|
||||
max_chars=1024,
|
||||
json=False,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cli, "_build_parser", lambda: ReadParser())
|
||||
cli.main()
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert "pyro$ pwd\n" in captured.out
|
||||
assert "[workspace-shell-open] workspace_id=workspace-123 shell_id=shell-123" in captured.err
|
||||
assert "[workspace-shell-read] workspace_id=workspace-123 shell_id=shell-123" in captured.err
|
||||
|
||||
|
||||
def test_cli_workspace_shell_write_signal_close_json(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
class StubPyro:
|
||||
def write_shell(
|
||||
self,
|
||||
workspace_id: str,
|
||||
shell_id: str,
|
||||
*,
|
||||
input: str,
|
||||
append_newline: bool,
|
||||
) -> dict[str, Any]:
|
||||
assert workspace_id == "workspace-123"
|
||||
assert shell_id == "shell-123"
|
||||
assert input == "pwd"
|
||||
assert append_newline is False
|
||||
return {
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
"state": "running",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"started_at": 1.0,
|
||||
"ended_at": None,
|
||||
"exit_code": None,
|
||||
"execution_mode": "guest_vsock",
|
||||
"input_length": 3,
|
||||
"append_newline": False,
|
||||
}
|
||||
|
||||
def signal_shell(
|
||||
self,
|
||||
workspace_id: str,
|
||||
shell_id: str,
|
||||
*,
|
||||
signal_name: str,
|
||||
) -> dict[str, Any]:
|
||||
assert signal_name == "INT"
|
||||
return {
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
"state": "running",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"started_at": 1.0,
|
||||
"ended_at": None,
|
||||
"exit_code": None,
|
||||
"execution_mode": "guest_vsock",
|
||||
"signal": signal_name,
|
||||
}
|
||||
|
||||
def close_shell(self, workspace_id: str, shell_id: str) -> dict[str, Any]:
|
||||
return {
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
"state": "stopped",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"started_at": 1.0,
|
||||
"ended_at": 2.0,
|
||||
"exit_code": 0,
|
||||
"execution_mode": "guest_vsock",
|
||||
"closed": True,
|
||||
}
|
||||
|
||||
class WriteParser:
|
||||
def parse_args(self) -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
command="workspace",
|
||||
workspace_command="shell",
|
||||
workspace_shell_command="write",
|
||||
workspace_id="workspace-123",
|
||||
shell_id="shell-123",
|
||||
input="pwd",
|
||||
no_newline=True,
|
||||
json=True,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cli, "_build_parser", lambda: WriteParser())
|
||||
monkeypatch.setattr(cli, "Pyro", StubPyro)
|
||||
cli.main()
|
||||
written = json.loads(capsys.readouterr().out)
|
||||
assert written["append_newline"] is False
|
||||
|
||||
class SignalParser:
|
||||
def parse_args(self) -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
command="workspace",
|
||||
workspace_command="shell",
|
||||
workspace_shell_command="signal",
|
||||
workspace_id="workspace-123",
|
||||
shell_id="shell-123",
|
||||
signal="INT",
|
||||
json=True,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cli, "_build_parser", lambda: SignalParser())
|
||||
cli.main()
|
||||
signaled = json.loads(capsys.readouterr().out)
|
||||
assert signaled["signal"] == "INT"
|
||||
|
||||
class CloseParser:
|
||||
def parse_args(self) -> argparse.Namespace:
|
||||
return argparse.Namespace(
|
||||
command="workspace",
|
||||
workspace_command="shell",
|
||||
workspace_shell_command="close",
|
||||
workspace_id="workspace-123",
|
||||
shell_id="shell-123",
|
||||
json=True,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(cli, "_build_parser", lambda: CloseParser())
|
||||
cli.main()
|
||||
closed = json.loads(capsys.readouterr().out)
|
||||
assert closed["closed"] is True
|
||||
|
||||
|
||||
def test_cli_workspace_exec_json_error_exits_nonzero(
|
||||
monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -18,6 +18,12 @@ from pyro_mcp.contract import (
|
|||
PUBLIC_CLI_ENV_SUBCOMMANDS,
|
||||
PUBLIC_CLI_RUN_FLAGS,
|
||||
PUBLIC_CLI_WORKSPACE_CREATE_FLAGS,
|
||||
PUBLIC_CLI_WORKSPACE_SHELL_CLOSE_FLAGS,
|
||||
PUBLIC_CLI_WORKSPACE_SHELL_OPEN_FLAGS,
|
||||
PUBLIC_CLI_WORKSPACE_SHELL_READ_FLAGS,
|
||||
PUBLIC_CLI_WORKSPACE_SHELL_SIGNAL_FLAGS,
|
||||
PUBLIC_CLI_WORKSPACE_SHELL_SUBCOMMANDS,
|
||||
PUBLIC_CLI_WORKSPACE_SHELL_WRITE_FLAGS,
|
||||
PUBLIC_CLI_WORKSPACE_SUBCOMMANDS,
|
||||
PUBLIC_CLI_WORKSPACE_SYNC_PUSH_FLAGS,
|
||||
PUBLIC_CLI_WORKSPACE_SYNC_SUBCOMMANDS,
|
||||
|
|
@ -86,6 +92,37 @@ def test_public_cli_help_lists_commands_and_run_flags() -> None:
|
|||
).format_help()
|
||||
for flag in PUBLIC_CLI_WORKSPACE_SYNC_PUSH_FLAGS:
|
||||
assert flag in workspace_sync_push_help_text
|
||||
workspace_shell_help_text = _subparser_choice(
|
||||
_subparser_choice(parser, "workspace"),
|
||||
"shell",
|
||||
).format_help()
|
||||
for subcommand_name in PUBLIC_CLI_WORKSPACE_SHELL_SUBCOMMANDS:
|
||||
assert subcommand_name in workspace_shell_help_text
|
||||
workspace_shell_open_help_text = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "open"
|
||||
).format_help()
|
||||
for flag in PUBLIC_CLI_WORKSPACE_SHELL_OPEN_FLAGS:
|
||||
assert flag in workspace_shell_open_help_text
|
||||
workspace_shell_read_help_text = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "read"
|
||||
).format_help()
|
||||
for flag in PUBLIC_CLI_WORKSPACE_SHELL_READ_FLAGS:
|
||||
assert flag in workspace_shell_read_help_text
|
||||
workspace_shell_write_help_text = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "write"
|
||||
).format_help()
|
||||
for flag in PUBLIC_CLI_WORKSPACE_SHELL_WRITE_FLAGS:
|
||||
assert flag in workspace_shell_write_help_text
|
||||
workspace_shell_signal_help_text = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "signal"
|
||||
).format_help()
|
||||
for flag in PUBLIC_CLI_WORKSPACE_SHELL_SIGNAL_FLAGS:
|
||||
assert flag in workspace_shell_signal_help_text
|
||||
workspace_shell_close_help_text = _subparser_choice(
|
||||
_subparser_choice(_subparser_choice(parser, "workspace"), "shell"), "close"
|
||||
).format_help()
|
||||
for flag in PUBLIC_CLI_WORKSPACE_SHELL_CLOSE_FLAGS:
|
||||
assert flag in workspace_shell_close_help_text
|
||||
|
||||
demo_help_text = _subparser_choice(parser, "demo").format_help()
|
||||
for subcommand_name in PUBLIC_CLI_DEMO_SUBCOMMANDS:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
|
|
@ -34,6 +35,11 @@ def test_create_server_registers_vm_tools(tmp_path: Path) -> None:
|
|||
assert "workspace_create" in tool_names
|
||||
assert "workspace_logs" in tool_names
|
||||
assert "workspace_sync_push" in tool_names
|
||||
assert "shell_open" in tool_names
|
||||
assert "shell_read" in tool_names
|
||||
assert "shell_write" in tool_names
|
||||
assert "shell_signal" in tool_names
|
||||
assert "shell_close" in tool_names
|
||||
|
||||
|
||||
def test_vm_run_round_trip(tmp_path: Path) -> None:
|
||||
|
|
@ -190,6 +196,11 @@ def test_workspace_tools_round_trip(tmp_path: Path) -> None:
|
|||
dict[str, Any],
|
||||
dict[str, Any],
|
||||
dict[str, Any],
|
||||
dict[str, Any],
|
||||
dict[str, Any],
|
||||
dict[str, Any],
|
||||
dict[str, Any],
|
||||
dict[str, Any],
|
||||
]:
|
||||
server = create_server(manager=manager)
|
||||
created = _extract_structured(
|
||||
|
|
@ -225,18 +236,88 @@ def test_workspace_tools_round_trip(tmp_path: Path) -> None:
|
|||
},
|
||||
)
|
||||
)
|
||||
opened = _extract_structured(
|
||||
await server.call_tool("shell_open", {"workspace_id": workspace_id})
|
||||
)
|
||||
shell_id = str(opened["shell_id"])
|
||||
written = _extract_structured(
|
||||
await server.call_tool(
|
||||
"shell_write",
|
||||
{
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
"input": "pwd",
|
||||
},
|
||||
)
|
||||
)
|
||||
read = _extract_structured(
|
||||
await server.call_tool(
|
||||
"shell_read",
|
||||
{
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
},
|
||||
)
|
||||
)
|
||||
deadline = time.time() + 5
|
||||
while "/workspace" not in str(read["output"]) and time.time() < deadline:
|
||||
read = _extract_structured(
|
||||
await server.call_tool(
|
||||
"shell_read",
|
||||
{
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
"cursor": 0,
|
||||
},
|
||||
)
|
||||
)
|
||||
await asyncio.sleep(0.05)
|
||||
signaled = _extract_structured(
|
||||
await server.call_tool(
|
||||
"shell_signal",
|
||||
{
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
},
|
||||
)
|
||||
)
|
||||
closed = _extract_structured(
|
||||
await server.call_tool(
|
||||
"shell_close",
|
||||
{
|
||||
"workspace_id": workspace_id,
|
||||
"shell_id": shell_id,
|
||||
},
|
||||
)
|
||||
)
|
||||
logs = _extract_structured(
|
||||
await server.call_tool("workspace_logs", {"workspace_id": workspace_id})
|
||||
)
|
||||
deleted = _extract_structured(
|
||||
await server.call_tool("workspace_delete", {"workspace_id": workspace_id})
|
||||
)
|
||||
return created, synced, executed, logs, deleted
|
||||
return created, synced, executed, opened, written, read, signaled, closed, logs, deleted
|
||||
|
||||
created, synced, executed, logs, deleted = asyncio.run(_run())
|
||||
(
|
||||
created,
|
||||
synced,
|
||||
executed,
|
||||
opened,
|
||||
written,
|
||||
read,
|
||||
signaled,
|
||||
closed,
|
||||
logs,
|
||||
deleted,
|
||||
) = asyncio.run(_run())
|
||||
assert created["state"] == "started"
|
||||
assert created["workspace_seed"]["mode"] == "directory"
|
||||
assert synced["workspace_sync"]["destination"] == "/workspace/subdir"
|
||||
assert executed["stdout"] == "more\n"
|
||||
assert opened["state"] == "running"
|
||||
assert written["input_length"] == 3
|
||||
assert "/workspace" in read["output"]
|
||||
assert signaled["signal"] == "INT"
|
||||
assert closed["closed"] is True
|
||||
assert logs["count"] == 1
|
||||
assert deleted["deleted"] is True
|
||||
|
|
|
|||
|
|
@ -105,6 +105,130 @@ def test_vsock_exec_client_upload_archive_round_trip(
|
|||
assert stub.closed is True
|
||||
|
||||
|
||||
def test_vsock_exec_client_shell_round_trip(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(socket, "AF_VSOCK", 40, raising=False)
|
||||
responses = [
|
||||
json.dumps(
|
||||
{
|
||||
"shell_id": "shell-1",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"state": "running",
|
||||
"started_at": 1.0,
|
||||
"ended_at": None,
|
||||
"exit_code": None,
|
||||
}
|
||||
).encode("utf-8"),
|
||||
json.dumps(
|
||||
{
|
||||
"shell_id": "shell-1",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"state": "running",
|
||||
"started_at": 1.0,
|
||||
"ended_at": None,
|
||||
"exit_code": None,
|
||||
"cursor": 0,
|
||||
"next_cursor": 12,
|
||||
"output": "pyro$ pwd\n",
|
||||
"truncated": False,
|
||||
}
|
||||
).encode("utf-8"),
|
||||
json.dumps(
|
||||
{
|
||||
"shell_id": "shell-1",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"state": "running",
|
||||
"started_at": 1.0,
|
||||
"ended_at": None,
|
||||
"exit_code": None,
|
||||
"input_length": 3,
|
||||
"append_newline": True,
|
||||
}
|
||||
).encode("utf-8"),
|
||||
json.dumps(
|
||||
{
|
||||
"shell_id": "shell-1",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"state": "running",
|
||||
"started_at": 1.0,
|
||||
"ended_at": None,
|
||||
"exit_code": None,
|
||||
"signal": "INT",
|
||||
}
|
||||
).encode("utf-8"),
|
||||
json.dumps(
|
||||
{
|
||||
"shell_id": "shell-1",
|
||||
"cwd": "/workspace",
|
||||
"cols": 120,
|
||||
"rows": 30,
|
||||
"state": "stopped",
|
||||
"started_at": 1.0,
|
||||
"ended_at": 2.0,
|
||||
"exit_code": 0,
|
||||
"closed": True,
|
||||
}
|
||||
).encode("utf-8"),
|
||||
]
|
||||
stubs = [StubSocket(response) for response in responses]
|
||||
remaining = list(stubs)
|
||||
|
||||
def socket_factory(family: int, sock_type: int) -> StubSocket:
|
||||
assert family == socket.AF_VSOCK
|
||||
assert sock_type == socket.SOCK_STREAM
|
||||
return remaining.pop(0)
|
||||
|
||||
client = VsockExecClient(socket_factory=socket_factory)
|
||||
opened = client.open_shell(
|
||||
1234,
|
||||
5005,
|
||||
shell_id="shell-1",
|
||||
cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
)
|
||||
assert opened.shell_id == "shell-1"
|
||||
read = client.read_shell(1234, 5005, shell_id="shell-1", cursor=0, max_chars=1024)
|
||||
assert read.output == "pyro$ pwd\n"
|
||||
write = client.write_shell(
|
||||
1234,
|
||||
5005,
|
||||
shell_id="shell-1",
|
||||
input_text="pwd",
|
||||
append_newline=True,
|
||||
)
|
||||
assert write["input_length"] == 3
|
||||
signaled = client.signal_shell(1234, 5005, shell_id="shell-1", signal_name="INT")
|
||||
assert signaled["signal"] == "INT"
|
||||
closed = client.close_shell(1234, 5005, shell_id="shell-1")
|
||||
assert closed["closed"] is True
|
||||
open_request = json.loads(stubs[0].sent.decode("utf-8").strip())
|
||||
assert open_request["action"] == "open_shell"
|
||||
assert open_request["shell_id"] == "shell-1"
|
||||
|
||||
|
||||
def test_vsock_exec_client_raises_agent_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(socket, "AF_VSOCK", 40, raising=False)
|
||||
stub = StubSocket(b'{"error":"shell is unavailable"}')
|
||||
client = VsockExecClient(socket_factory=lambda family, sock_type: stub)
|
||||
with pytest.raises(RuntimeError, match="shell is unavailable"):
|
||||
client.open_shell(
|
||||
1234,
|
||||
5005,
|
||||
shell_id="shell-1",
|
||||
cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
)
|
||||
|
||||
|
||||
def test_vsock_exec_client_rejects_bad_json(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(socket, "AF_VSOCK", 40, raising=False)
|
||||
stub = StubSocket(b"[]")
|
||||
|
|
|
|||
|
|
@ -454,6 +454,73 @@ def test_workspace_sync_push_rejects_destination_outside_workspace(tmp_path: Pat
|
|||
manager.push_workspace_sync(workspace_id, source_path=source_dir, dest="../escape")
|
||||
|
||||
|
||||
def test_workspace_shell_lifecycle_and_rehydration(tmp_path: Path) -> None:
|
||||
manager = VmManager(
|
||||
backend_name="mock",
|
||||
base_dir=tmp_path / "vms",
|
||||
network_manager=TapNetworkManager(enabled=False),
|
||||
)
|
||||
created = manager.create_workspace(
|
||||
environment="debian:12-base",
|
||||
allow_host_compat=True,
|
||||
)
|
||||
workspace_id = str(created["workspace_id"])
|
||||
|
||||
opened = manager.open_shell(workspace_id)
|
||||
shell_id = str(opened["shell_id"])
|
||||
assert opened["state"] == "running"
|
||||
|
||||
manager.write_shell(workspace_id, shell_id, input_text="pwd")
|
||||
|
||||
output = ""
|
||||
deadline = time.time() + 5
|
||||
while time.time() < deadline:
|
||||
read = manager.read_shell(workspace_id, shell_id, cursor=0, max_chars=65536)
|
||||
output = str(read["output"])
|
||||
if "/workspace" in output:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
assert "/workspace" in output
|
||||
|
||||
manager_rehydrated = VmManager(
|
||||
backend_name="mock",
|
||||
base_dir=tmp_path / "vms",
|
||||
network_manager=TapNetworkManager(enabled=False),
|
||||
)
|
||||
second_opened = manager_rehydrated.open_shell(workspace_id)
|
||||
second_shell_id = str(second_opened["shell_id"])
|
||||
assert second_shell_id != shell_id
|
||||
|
||||
manager_rehydrated.write_shell(workspace_id, second_shell_id, input_text="printf 'ok\\n'")
|
||||
second_output = ""
|
||||
deadline = time.time() + 5
|
||||
while time.time() < deadline:
|
||||
read = manager_rehydrated.read_shell(
|
||||
workspace_id,
|
||||
second_shell_id,
|
||||
cursor=0,
|
||||
max_chars=65536,
|
||||
)
|
||||
second_output = str(read["output"])
|
||||
if "ok" in second_output:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
assert "ok" in second_output
|
||||
|
||||
logs = manager.logs_workspace(workspace_id)
|
||||
assert logs["count"] == 0
|
||||
|
||||
closed = manager.close_shell(workspace_id, shell_id)
|
||||
assert closed["closed"] is True
|
||||
with pytest.raises(ValueError, match="does not exist"):
|
||||
manager.read_shell(workspace_id, shell_id)
|
||||
|
||||
deleted = manager.delete_workspace(workspace_id)
|
||||
assert deleted["deleted"] is True
|
||||
with pytest.raises(ValueError, match="does not exist"):
|
||||
manager_rehydrated.read_shell(workspace_id, second_shell_id)
|
||||
|
||||
|
||||
def test_workspace_create_rejects_unsafe_seed_archive(tmp_path: Path) -> None:
|
||||
archive_path = tmp_path / "bad.tgz"
|
||||
with tarfile.open(archive_path, "w:gz") as archive:
|
||||
|
|
|
|||
220
tests/test_workspace_shells.py
Normal file
220
tests/test_workspace_shells.py
Normal file
|
|
@ -0,0 +1,220 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import cast
|
||||
|
||||
import pytest
|
||||
|
||||
from pyro_mcp.workspace_shells import (
|
||||
create_local_shell,
|
||||
get_local_shell,
|
||||
remove_local_shell,
|
||||
shell_signal_arg_help,
|
||||
shell_signal_names,
|
||||
)
|
||||
|
||||
|
||||
def _read_until(
|
||||
workspace_id: str,
|
||||
shell_id: str,
|
||||
text: str,
|
||||
*,
|
||||
timeout_seconds: float = 5.0,
|
||||
) -> dict[str, object]:
|
||||
deadline = time.time() + timeout_seconds
|
||||
payload = get_local_shell(workspace_id=workspace_id, shell_id=shell_id).read(
|
||||
cursor=0,
|
||||
max_chars=65536,
|
||||
)
|
||||
while text not in str(payload["output"]) and time.time() < deadline:
|
||||
time.sleep(0.05)
|
||||
payload = get_local_shell(workspace_id=workspace_id, shell_id=shell_id).read(
|
||||
cursor=0,
|
||||
max_chars=65536,
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
def test_workspace_shells_round_trip(tmp_path: Path) -> None:
|
||||
session = create_local_shell(
|
||||
workspace_id="workspace-1",
|
||||
shell_id="shell-1",
|
||||
cwd=tmp_path,
|
||||
display_cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
)
|
||||
try:
|
||||
assert session.summary()["state"] == "running"
|
||||
write = session.write("printf 'hello\\n'", append_newline=True)
|
||||
assert write["input_length"] == 16
|
||||
payload = _read_until("workspace-1", "shell-1", "hello")
|
||||
assert "hello" in str(payload["output"])
|
||||
assert cast(int, payload["next_cursor"]) >= cast(int, payload["cursor"])
|
||||
assert isinstance(payload["truncated"], bool)
|
||||
session.write("sleep 60", append_newline=True)
|
||||
signaled = session.send_signal("INT")
|
||||
assert signaled["signal"] == "INT"
|
||||
finally:
|
||||
closed = session.close()
|
||||
assert closed["closed"] is True
|
||||
|
||||
|
||||
def test_workspace_shell_registry_helpers(tmp_path: Path) -> None:
|
||||
session = create_local_shell(
|
||||
workspace_id="workspace-2",
|
||||
shell_id="shell-2",
|
||||
cwd=tmp_path,
|
||||
display_cwd="/workspace/subdir",
|
||||
cols=80,
|
||||
rows=24,
|
||||
)
|
||||
assert get_local_shell(workspace_id="workspace-2", shell_id="shell-2") is session
|
||||
assert shell_signal_names() == ("HUP", "INT", "TERM", "KILL")
|
||||
assert "HUP" in shell_signal_arg_help()
|
||||
with pytest.raises(RuntimeError, match="already exists"):
|
||||
create_local_shell(
|
||||
workspace_id="workspace-2",
|
||||
shell_id="shell-2",
|
||||
cwd=tmp_path,
|
||||
display_cwd="/workspace/subdir",
|
||||
cols=80,
|
||||
rows=24,
|
||||
)
|
||||
removed = remove_local_shell(workspace_id="workspace-2", shell_id="shell-2")
|
||||
assert removed is session
|
||||
assert remove_local_shell(workspace_id="workspace-2", shell_id="shell-2") is None
|
||||
with pytest.raises(ValueError, match="does not exist"):
|
||||
get_local_shell(workspace_id="workspace-2", shell_id="shell-2")
|
||||
closed = session.close()
|
||||
assert closed["closed"] is True
|
||||
|
||||
|
||||
def test_workspace_shells_error_after_exit(tmp_path: Path) -> None:
|
||||
session = create_local_shell(
|
||||
workspace_id="workspace-3",
|
||||
shell_id="shell-3",
|
||||
cwd=tmp_path,
|
||||
display_cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
)
|
||||
session.write("exit", append_newline=True)
|
||||
deadline = time.time() + 5
|
||||
while session.summary()["state"] != "stopped" and time.time() < deadline:
|
||||
time.sleep(0.05)
|
||||
assert session.summary()["state"] == "stopped"
|
||||
with pytest.raises(RuntimeError, match="not running"):
|
||||
session.write("pwd", append_newline=True)
|
||||
with pytest.raises(RuntimeError, match="not running"):
|
||||
session.send_signal("INT")
|
||||
closed = session.close()
|
||||
assert closed["closed"] is True
|
||||
|
||||
|
||||
def test_workspace_shells_reject_invalid_signal(tmp_path: Path) -> None:
|
||||
session = create_local_shell(
|
||||
workspace_id="workspace-4",
|
||||
shell_id="shell-4",
|
||||
cwd=tmp_path,
|
||||
display_cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
)
|
||||
try:
|
||||
with pytest.raises(ValueError, match="unsupported shell signal"):
|
||||
session.send_signal("BOGUS")
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
|
||||
def test_workspace_shells_init_failure_closes_ptys(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
def _boom(*args: object, **kwargs: object) -> object:
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr(subprocess, "Popen", _boom)
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
create_local_shell(
|
||||
workspace_id="workspace-5",
|
||||
shell_id="shell-5",
|
||||
cwd=tmp_path,
|
||||
display_cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
)
|
||||
|
||||
|
||||
def test_workspace_shells_write_and_signal_runtime_errors(
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
session = create_local_shell(
|
||||
workspace_id="workspace-6",
|
||||
shell_id="shell-6",
|
||||
cwd=tmp_path,
|
||||
display_cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
)
|
||||
try:
|
||||
with session._lock: # noqa: SLF001
|
||||
session._master_fd = None # noqa: SLF001
|
||||
with pytest.raises(RuntimeError, match="transport is unavailable"):
|
||||
session.write("pwd", append_newline=True)
|
||||
|
||||
with session._lock: # noqa: SLF001
|
||||
master_fd, slave_fd = os.pipe()
|
||||
os.close(slave_fd)
|
||||
session._master_fd = master_fd # noqa: SLF001
|
||||
|
||||
def _raise_write(fd: int, data: bytes) -> int:
|
||||
del fd, data
|
||||
raise OSError("broken")
|
||||
|
||||
monkeypatch.setattr("pyro_mcp.workspace_shells.os.write", _raise_write)
|
||||
with pytest.raises(RuntimeError, match="failed to write"):
|
||||
session.write("pwd", append_newline=True)
|
||||
|
||||
def _raise_killpg(pid: int, signum: int) -> None:
|
||||
del pid, signum
|
||||
raise ProcessLookupError()
|
||||
|
||||
monkeypatch.setattr("pyro_mcp.workspace_shells.os.killpg", _raise_killpg)
|
||||
with pytest.raises(RuntimeError, match="not running"):
|
||||
session.send_signal("INT")
|
||||
finally:
|
||||
try:
|
||||
session.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def test_workspace_shells_refresh_process_state_updates_exit_code(tmp_path: Path) -> None:
|
||||
session = create_local_shell(
|
||||
workspace_id="workspace-7",
|
||||
shell_id="shell-7",
|
||||
cwd=tmp_path,
|
||||
display_cwd="/workspace",
|
||||
cols=120,
|
||||
rows=30,
|
||||
)
|
||||
try:
|
||||
class StubProcess:
|
||||
def poll(self) -> int:
|
||||
return 7
|
||||
|
||||
session._process = StubProcess() # type: ignore[assignment] # noqa: SLF001
|
||||
session._refresh_process_state() # noqa: SLF001
|
||||
assert session.summary()["state"] == "stopped"
|
||||
assert session.summary()["exit_code"] == 7
|
||||
finally:
|
||||
try:
|
||||
session.close()
|
||||
except Exception:
|
||||
pass
|
||||
Loading…
Add table
Add a link
Reference in a new issue