forked from wrenn/python-sdk
502 lines
16 KiB
Python
502 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import builtins
|
|
import json
|
|
from collections.abc import AsyncIterator, Iterator
|
|
from dataclasses import dataclass
|
|
from typing import Literal, overload
|
|
|
|
import httpx
|
|
import httpx_ws
|
|
|
|
from wrenn.exceptions import handle_response
|
|
|
|
|
|
@dataclass
|
|
class CommandResult:
|
|
"""Result from a foreground command execution."""
|
|
|
|
stdout: str
|
|
stderr: str
|
|
exit_code: int
|
|
duration_ms: int | None = None
|
|
|
|
|
|
@dataclass
|
|
class CommandHandle:
|
|
"""Handle for a background process."""
|
|
|
|
pid: int
|
|
tag: str
|
|
capsule_id: str
|
|
|
|
|
|
@dataclass
|
|
class ProcessInfo:
|
|
"""Information about a running process."""
|
|
|
|
pid: int
|
|
tag: str | None = None
|
|
cmd: str | None = None
|
|
args: list[str] | None = None
|
|
|
|
|
|
class StreamEvent:
|
|
"""Base class for streaming exec events."""
|
|
|
|
__slots__ = ("type",)
|
|
|
|
def __init__(self, type: str) -> None:
|
|
self.type = type
|
|
|
|
|
|
class StreamStartEvent(StreamEvent):
|
|
__slots__ = ("pid",)
|
|
|
|
def __init__(self, pid: int) -> None:
|
|
super().__init__("start")
|
|
self.pid = pid
|
|
|
|
|
|
class StreamStdoutEvent(StreamEvent):
|
|
__slots__ = ("data",)
|
|
|
|
def __init__(self, data: str) -> None:
|
|
super().__init__("stdout")
|
|
self.data = data
|
|
|
|
|
|
class StreamStderrEvent(StreamEvent):
|
|
__slots__ = ("data",)
|
|
|
|
def __init__(self, data: str) -> None:
|
|
super().__init__("stderr")
|
|
self.data = data
|
|
|
|
|
|
class StreamExitEvent(StreamEvent):
|
|
__slots__ = ("exit_code",)
|
|
|
|
def __init__(self, exit_code: int) -> None:
|
|
super().__init__("exit")
|
|
self.exit_code = exit_code
|
|
|
|
|
|
class StreamErrorEvent(StreamEvent):
|
|
__slots__ = ("data",)
|
|
|
|
def __init__(self, data: str) -> None:
|
|
super().__init__("error")
|
|
self.data = data
|
|
|
|
|
|
def _parse_stream_event(raw: dict) -> StreamEvent:
|
|
t = raw.get("type")
|
|
if t == "start":
|
|
return StreamStartEvent(pid=raw.get("pid", 0))
|
|
if t == "stdout":
|
|
return StreamStdoutEvent(data=raw.get("data", ""))
|
|
if t == "stderr":
|
|
return StreamStderrEvent(data=raw.get("data", ""))
|
|
if t == "exit":
|
|
return StreamExitEvent(exit_code=raw.get("exit_code", -1))
|
|
if t == "error":
|
|
return StreamErrorEvent(data=raw.get("data", ""))
|
|
return StreamEvent(type=t or "unknown")
|
|
|
|
|
|
def _decode_exec_response(data: dict) -> CommandResult:
|
|
stdout = data.get("stdout") or ""
|
|
stderr = data.get("stderr") or ""
|
|
if data.get("encoding") == "base64":
|
|
stdout = base64.b64decode(stdout).decode("utf-8", errors="replace")
|
|
if stderr:
|
|
stderr = base64.b64decode(stderr).decode("utf-8", errors="replace")
|
|
return CommandResult(
|
|
stdout=stdout,
|
|
stderr=stderr,
|
|
exit_code=data.get("exit_code", -1),
|
|
duration_ms=data.get("duration_ms"),
|
|
)
|
|
|
|
|
|
class Commands:
|
|
"""Sync command execution interface. Accessed via ``capsule.commands``."""
|
|
|
|
def __init__(self, capsule_id: str, http: httpx.Client) -> None:
|
|
self._capsule_id = capsule_id
|
|
self._http = http
|
|
|
|
@overload
|
|
def run(
|
|
self,
|
|
cmd: str,
|
|
*,
|
|
background: Literal[False] = ...,
|
|
timeout: int | None = 30,
|
|
envs: dict[str, str] | None = None,
|
|
cwd: str | None = None,
|
|
tag: str | None = None,
|
|
) -> CommandResult: ...
|
|
|
|
@overload
|
|
def run(
|
|
self,
|
|
cmd: str,
|
|
*,
|
|
background: Literal[True],
|
|
timeout: int | None = 30,
|
|
envs: dict[str, str] | None = None,
|
|
cwd: str | None = None,
|
|
tag: str | None = None,
|
|
) -> CommandHandle: ...
|
|
|
|
def run(
|
|
self,
|
|
cmd: str,
|
|
*,
|
|
background: bool = False,
|
|
timeout: int | None = 30,
|
|
envs: dict[str, str] | None = None,
|
|
cwd: str | None = None,
|
|
tag: str | None = None,
|
|
) -> CommandResult | CommandHandle:
|
|
"""Execute a shell command inside the capsule.
|
|
|
|
Args:
|
|
cmd (str): Shell command string to execute.
|
|
background (bool): If ``True``, launch the process in the
|
|
background and return a :class:`CommandHandle` immediately.
|
|
Defaults to ``False``.
|
|
timeout (int | None): Seconds before the foreground command times
|
|
out. Ignored for background commands. Defaults to ``30``.
|
|
envs (dict[str, str] | None): Additional environment variables
|
|
to set for the process.
|
|
cwd (str | None): Working directory for the process.
|
|
tag (str | None): Optional label attached to background processes
|
|
for later retrieval via :meth:`connect`.
|
|
|
|
Returns:
|
|
CommandResult: stdout, stderr, exit code, and duration for
|
|
foreground commands (``background=False``).
|
|
|
|
CommandHandle: PID and tag for background commands
|
|
(``background=True``).
|
|
"""
|
|
payload: dict = {
|
|
"cmd": "/bin/sh",
|
|
"args": ["-c", cmd],
|
|
"background": background,
|
|
}
|
|
if timeout is not None and not background:
|
|
payload["timeout_sec"] = timeout
|
|
if envs is not None:
|
|
payload["envs"] = envs
|
|
if cwd is not None:
|
|
payload["cwd"] = cwd
|
|
if tag is not None:
|
|
payload["tag"] = tag
|
|
|
|
http_timeout: httpx.Timeout | None = None
|
|
if not background and timeout is not None:
|
|
http_timeout = httpx.Timeout(timeout + 10, connect=5.0)
|
|
|
|
resp = self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/exec",
|
|
json=payload,
|
|
timeout=http_timeout,
|
|
)
|
|
data = handle_response(resp)
|
|
assert isinstance(data, dict)
|
|
|
|
if background:
|
|
return CommandHandle(
|
|
pid=data.get("pid", 0),
|
|
tag=data.get("tag", ""),
|
|
capsule_id=self._capsule_id,
|
|
)
|
|
return _decode_exec_response(data)
|
|
|
|
def list(self) -> list[ProcessInfo]:
|
|
"""List all running background processes in the capsule.
|
|
|
|
Returns:
|
|
list[ProcessInfo]: Running processes with their PID, tag, and
|
|
command information.
|
|
"""
|
|
resp = self._http.get(f"/v1/capsules/{self._capsule_id}/processes")
|
|
data = handle_response(resp)
|
|
assert isinstance(data, dict)
|
|
return [
|
|
ProcessInfo(
|
|
pid=p.get("pid", 0),
|
|
tag=p.get("tag"),
|
|
cmd=p.get("cmd"),
|
|
args=p.get("args"),
|
|
)
|
|
for p in data.get("processes", [])
|
|
]
|
|
|
|
def kill(self, pid: int) -> None:
|
|
"""Send SIGKILL to a background process.
|
|
|
|
Args:
|
|
pid (int): PID of the process to kill.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If no process with the given PID exists.
|
|
"""
|
|
resp = self._http.delete(f"/v1/capsules/{self._capsule_id}/processes/{pid}")
|
|
handle_response(resp)
|
|
|
|
def connect(self, pid: int) -> Iterator[StreamEvent]:
|
|
"""Connect to a running background process and stream its output.
|
|
|
|
Args:
|
|
pid (int): PID of the background process to attach to.
|
|
|
|
Yields:
|
|
StreamEvent: Successive output events. Stops on
|
|
:class:`StreamExitEvent` or :class:`StreamErrorEvent`.
|
|
"""
|
|
with httpx_ws.connect_ws(
|
|
f"/v1/capsules/{self._capsule_id}/processes/{pid}/stream",
|
|
self._http,
|
|
) as ws: # type: httpx_ws.WebSocketSession
|
|
while True:
|
|
try:
|
|
raw = ws.receive_json()
|
|
event = _parse_stream_event(raw)
|
|
yield event
|
|
if event.type in ("exit", "error"):
|
|
break
|
|
except httpx_ws.WebSocketDisconnect:
|
|
break
|
|
|
|
def stream(
|
|
self, cmd: str, args: builtins.list[str] | None = None
|
|
) -> Iterator[StreamEvent]:
|
|
"""Execute a command via WebSocket, streaming output as events.
|
|
|
|
Args:
|
|
cmd (str): Command to execute.
|
|
args (list[str] | None): Additional arguments for the command.
|
|
When omitted, *cmd* is interpreted as a shell command
|
|
string and executed via ``/bin/sh -c``.
|
|
|
|
Yields:
|
|
StreamEvent: Successive events including :class:`StreamStartEvent`,
|
|
:class:`StreamStdoutEvent`, :class:`StreamStderrEvent`,
|
|
:class:`StreamExitEvent`, and :class:`StreamErrorEvent`.
|
|
"""
|
|
with httpx_ws.connect_ws(
|
|
f"/v1/capsules/{self._capsule_id}/exec/stream",
|
|
self._http,
|
|
) as ws: # type: httpx_ws.WebSocketSession
|
|
if args:
|
|
start_msg: dict = {"type": "start", "cmd": cmd, "args": args}
|
|
else:
|
|
start_msg = {"type": "start", "cmd": "/bin/sh", "args": ["-c", cmd]}
|
|
ws.send_text(json.dumps(start_msg))
|
|
while True:
|
|
try:
|
|
raw = ws.receive_json()
|
|
event = _parse_stream_event(raw)
|
|
yield event
|
|
if event.type in ("exit", "error"):
|
|
break
|
|
except httpx_ws.WebSocketDisconnect:
|
|
break
|
|
|
|
|
|
class AsyncCommands:
|
|
"""Async command execution interface. Accessed via ``capsule.commands``."""
|
|
|
|
def __init__(self, capsule_id: str, http: httpx.AsyncClient) -> None:
|
|
self._capsule_id = capsule_id
|
|
self._http = http
|
|
|
|
@overload
|
|
async def run(
|
|
self,
|
|
cmd: str,
|
|
*,
|
|
background: Literal[False] = ...,
|
|
timeout: int | None = 30,
|
|
envs: dict[str, str] | None = None,
|
|
cwd: str | None = None,
|
|
tag: str | None = None,
|
|
) -> CommandResult: ...
|
|
|
|
@overload
|
|
async def run(
|
|
self,
|
|
cmd: str,
|
|
*,
|
|
background: Literal[True],
|
|
timeout: int | None = 30,
|
|
envs: dict[str, str] | None = None,
|
|
cwd: str | None = None,
|
|
tag: str | None = None,
|
|
) -> CommandHandle: ...
|
|
|
|
async def run(
|
|
self,
|
|
cmd: str,
|
|
*,
|
|
background: bool = False,
|
|
timeout: int | None = 30,
|
|
envs: dict[str, str] | None = None,
|
|
cwd: str | None = None,
|
|
tag: str | None = None,
|
|
) -> CommandResult | CommandHandle:
|
|
"""Execute a shell command inside the capsule.
|
|
|
|
Args:
|
|
cmd (str): Shell command string to execute.
|
|
background (bool): If ``True``, launch the process in the
|
|
background and return a :class:`CommandHandle` immediately.
|
|
Defaults to ``False``.
|
|
timeout (int | None): Seconds before the foreground command times
|
|
out. Ignored for background commands. Defaults to ``30``.
|
|
envs (dict[str, str] | None): Additional environment variables
|
|
to set for the process.
|
|
cwd (str | None): Working directory for the process.
|
|
tag (str | None): Optional label attached to background processes
|
|
for later retrieval via :meth:`connect`.
|
|
|
|
Returns:
|
|
CommandResult: stdout, stderr, exit code, and duration for
|
|
foreground commands (``background=False``).
|
|
|
|
CommandHandle: PID and tag for background commands
|
|
(``background=True``).
|
|
"""
|
|
payload: dict = {
|
|
"cmd": "/bin/sh",
|
|
"args": ["-c", cmd],
|
|
"background": background,
|
|
}
|
|
if timeout is not None and not background:
|
|
payload["timeout_sec"] = timeout
|
|
if envs is not None:
|
|
payload["envs"] = envs
|
|
if cwd is not None:
|
|
payload["cwd"] = cwd
|
|
if tag is not None:
|
|
payload["tag"] = tag
|
|
|
|
http_timeout: httpx.Timeout | None = None
|
|
if not background and timeout is not None:
|
|
http_timeout = httpx.Timeout(timeout + 10, connect=5.0)
|
|
|
|
resp = await self._http.post(
|
|
f"/v1/capsules/{self._capsule_id}/exec",
|
|
json=payload,
|
|
timeout=http_timeout,
|
|
)
|
|
data = handle_response(resp)
|
|
assert isinstance(data, dict)
|
|
|
|
if background:
|
|
return CommandHandle(
|
|
pid=data.get("pid", 0),
|
|
tag=data.get("tag", ""),
|
|
capsule_id=self._capsule_id,
|
|
)
|
|
return _decode_exec_response(data)
|
|
|
|
async def list(self) -> list[ProcessInfo]:
|
|
"""List all running background processes in the capsule.
|
|
|
|
Returns:
|
|
list[ProcessInfo]: Running processes with their PID, tag, and
|
|
command information.
|
|
"""
|
|
resp = await self._http.get(f"/v1/capsules/{self._capsule_id}/processes")
|
|
data = handle_response(resp)
|
|
assert isinstance(data, dict)
|
|
return [
|
|
ProcessInfo(
|
|
pid=p.get("pid", 0),
|
|
tag=p.get("tag"),
|
|
cmd=p.get("cmd"),
|
|
args=p.get("args"),
|
|
)
|
|
for p in data.get("processes", [])
|
|
]
|
|
|
|
async def kill(self, pid: int) -> None:
|
|
"""Send SIGKILL to a background process.
|
|
|
|
Args:
|
|
pid (int): PID of the process to kill.
|
|
|
|
Raises:
|
|
WrennNotFoundError: If no process with the given PID exists.
|
|
"""
|
|
resp = await self._http.delete(
|
|
f"/v1/capsules/{self._capsule_id}/processes/{pid}"
|
|
)
|
|
handle_response(resp)
|
|
|
|
async def connect(self, pid: int) -> AsyncIterator[StreamEvent]:
|
|
"""Connect to a running background process and stream its output.
|
|
|
|
Args:
|
|
pid (int): PID of the background process to attach to.
|
|
|
|
Yields:
|
|
StreamEvent: Successive output events. Stops on
|
|
:class:`StreamExitEvent` or :class:`StreamErrorEvent`.
|
|
"""
|
|
async with httpx_ws.aconnect_ws(
|
|
f"/v1/capsules/{self._capsule_id}/processes/{pid}/stream",
|
|
self._http,
|
|
) as ws: # type: httpx_ws.AsyncWebSocketSession
|
|
try:
|
|
while True:
|
|
raw = await ws.receive_json()
|
|
event = _parse_stream_event(raw)
|
|
yield event
|
|
if event.type in ("exit", "error"):
|
|
break
|
|
except httpx_ws.WebSocketDisconnect:
|
|
pass
|
|
|
|
async def stream(
|
|
self, cmd: str, args: builtins.list[str] | None = None
|
|
) -> AsyncIterator[StreamEvent]:
|
|
"""Execute a command via WebSocket, streaming output as events.
|
|
|
|
Args:
|
|
cmd (str): Command to execute.
|
|
args (list[str] | None): Additional arguments for the command.
|
|
When omitted, *cmd* is interpreted as a shell command
|
|
string and executed via ``/bin/sh -c``.
|
|
|
|
Yields:
|
|
StreamEvent: Successive events including :class:`StreamStartEvent`,
|
|
:class:`StreamStdoutEvent`, :class:`StreamStderrEvent`,
|
|
:class:`StreamExitEvent`, and :class:`StreamErrorEvent`.
|
|
"""
|
|
async with httpx_ws.aconnect_ws(
|
|
f"/v1/capsules/{self._capsule_id}/exec/stream",
|
|
self._http,
|
|
) as ws: # type: httpx_ws.AsyncWebSocketSession
|
|
if args:
|
|
start_msg: dict = {"type": "start", "cmd": cmd, "args": args}
|
|
else:
|
|
start_msg = {"type": "start", "cmd": "/bin/sh", "args": ["-c", cmd]}
|
|
await ws.send_text(json.dumps(start_msg))
|
|
try:
|
|
while True:
|
|
raw = await ws.receive_json()
|
|
event = _parse_stream_event(raw)
|
|
yield event
|
|
if event.type in ("exit", "error"):
|
|
break
|
|
except httpx_ws.WebSocketDisconnect:
|
|
pass
|