130 lines
5.0 KiB
Python
130 lines
5.0 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import socket
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from uuid import UUID, uuid4
|
|
|
|
from .redaction import redact_text
|
|
|
|
|
|
class AgentError(RuntimeError):
|
|
pass
|
|
|
|
|
|
class AgentClient:
|
|
PROTOCOL_VERSION = 1
|
|
MAX_MESSAGE_BYTES = 64 * 1024
|
|
|
|
def __init__(self, socket_path: Path, *, timeout: int = 30):
|
|
self.socket_path = Path(socket_path)
|
|
self.timeout = timeout
|
|
|
|
def request(self, method: str, path: str, **values: Any) -> dict[str, Any]:
|
|
request = {
|
|
"protocol_version": self.PROTOCOL_VERSION,
|
|
"method": method,
|
|
"path": path,
|
|
**values,
|
|
}
|
|
encoded = json.dumps(request, separators=(",", ":"), ensure_ascii=True).encode("utf-8") + b"\n"
|
|
if len(encoded) > self.MAX_MESSAGE_BYTES:
|
|
raise AgentError("Agent request exceeds 64 KiB.")
|
|
if not hasattr(socket, "AF_UNIX"):
|
|
raise AgentError("Unix sockets are unavailable on this platform.")
|
|
try:
|
|
with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client:
|
|
client.settimeout(self.timeout)
|
|
client.connect(str(self.socket_path))
|
|
client.sendall(encoded)
|
|
client.shutdown(socket.SHUT_WR)
|
|
response = bytearray()
|
|
while b"\n" not in response:
|
|
chunk = client.recv(4096)
|
|
if not chunk:
|
|
break
|
|
response.extend(chunk)
|
|
if len(response) > self.MAX_MESSAGE_BYTES:
|
|
raise AgentError("Agent response exceeds 64 KiB.")
|
|
except (OSError, TimeoutError) as exc:
|
|
raise AgentError(f"Agent connection failed: {type(exc).__name__}") from exc
|
|
line, separator, remainder = bytes(response).partition(b"\n")
|
|
if not separator or remainder:
|
|
raise AgentError("Agent must return exactly one JSON line.")
|
|
try:
|
|
payload = json.loads(line.decode("utf-8"))
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
raise AgentError("Agent returned invalid JSON.") from exc
|
|
if not isinstance(payload, dict) or payload.get("protocol_version") != self.PROTOCOL_VERSION:
|
|
raise AgentError("Agent returned an incompatible protocol response.")
|
|
status = payload.get("status")
|
|
body = payload.get("body")
|
|
if not isinstance(status, int) or not isinstance(body, dict):
|
|
raise AgentError("Agent returned a malformed response.")
|
|
if status < 200 or status >= 300:
|
|
message = body.get("error") or body.get("detail") or f"Agent returned status {status}."
|
|
raise AgentError(redact_text(message, limit=2_000))
|
|
return body
|
|
|
|
def capabilities(self) -> dict[str, Any]:
|
|
return self.request("GET", "/v1/capabilities")
|
|
|
|
def submit_operation(
|
|
self,
|
|
*,
|
|
request_id: str,
|
|
action: str,
|
|
plugin_slug: str,
|
|
version: str | None,
|
|
approved_payload_sha256: str | None,
|
|
requested_by: str,
|
|
) -> UUID:
|
|
idempotency_key = uuid4()
|
|
body = self.request(
|
|
"POST",
|
|
"/v1/operations",
|
|
idempotency_key=str(idempotency_key),
|
|
body={
|
|
"request_id": request_id,
|
|
"action": action,
|
|
"plugin_slug": plugin_slug,
|
|
"version": version,
|
|
"approved_payload_sha256": approved_payload_sha256,
|
|
"requested_by": requested_by,
|
|
},
|
|
)
|
|
raw_id = body.get("operation_id") or body.get("id")
|
|
try:
|
|
return UUID(str(raw_id))
|
|
except (ValueError, TypeError) as exc:
|
|
raise AgentError("Agent did not return a valid operation UUID.") from exc
|
|
|
|
def wait_for_operation(
|
|
self,
|
|
operation_id: UUID,
|
|
*,
|
|
overall_timeout: int,
|
|
poll_interval: float,
|
|
) -> dict[str, Any]:
|
|
deadline = time.monotonic() + overall_timeout
|
|
while True:
|
|
body = self.get_operation(operation_id)
|
|
state = body.get("state") or body.get("status")
|
|
if state in {"succeeded", "completed"}:
|
|
result = body.get("result", body)
|
|
if not isinstance(result, dict):
|
|
raise AgentError("Agent operation result is malformed.")
|
|
return result
|
|
if state in {"failed", "errored", "cancelled"}:
|
|
raise AgentError(redact_text(body.get("error") or f"Agent operation {state}."))
|
|
if state not in {"queued", "pending", "running", "accepted"}:
|
|
raise AgentError("Agent returned an unknown operation state.")
|
|
if time.monotonic() >= deadline:
|
|
raise AgentError("Timed out waiting for the host agent operation.")
|
|
time.sleep(poll_interval)
|
|
|
|
def get_operation(self, operation_id: UUID) -> dict[str, Any]:
|
|
return self.request("GET", f"/v1/operations/{operation_id}")
|