476 lines
17 KiB
Python
476 lines
17 KiB
Python
"""沙箱客户端：封装对 Rust 沙箱服务的调用，含司命授权签发与一炁事件回写。
|
||
|
||
使用约束（见 backend/CLAUDE.md）：禁止在调用处裸 ``async with httpx.AsyncClient()``，
|
||
应复用注入的 ``AsyncClient``（长生命周期，由上层依赖注入管理）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import time
|
||
from collections.abc import AsyncIterator
|
||
from dataclasses import dataclass
|
||
from typing import Any, NoReturn, TypeVar
|
||
|
||
import httpx
|
||
from pydantic import BaseModel
|
||
|
||
from app.core.config import SandboxSettings
|
||
from app.integrations.sandbox.auth import AuthorizationContext, sign_authorization
|
||
from app.integrations.sandbox.events import (
|
||
EventSink,
|
||
LoggingEventSink,
|
||
completed_event,
|
||
failed_event,
|
||
started_event,
|
||
)
|
||
from app.integrations.sandbox.exceptions import (
|
||
SandboxBadRequest,
|
||
SandboxError,
|
||
SandboxExecFailed,
|
||
SandboxNotFound,
|
||
SandboxUnauthorized,
|
||
SandboxUnavailable,
|
||
)
|
||
from app.integrations.sandbox.models import (
|
||
CancelResponse,
|
||
ExecRequest,
|
||
ExecResponse,
|
||
FileListRequest,
|
||
FileListResponse,
|
||
FileReadRequest,
|
||
FileReadResponse,
|
||
FileStatRequest,
|
||
FileStatResponse,
|
||
FileWriteRequest,
|
||
FileWriteResponse,
|
||
)
|
||
|
||
_T = TypeVar("_T", bound=BaseModel)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class SandboxStreamEvent:
|
||
"""沙箱流式事件：stdout/stderr 增量（base64 str，保真透传）或 done（结束汇总 dict）。"""
|
||
|
||
kind: str # "stdout" | "stderr" | "done"
|
||
payload: str | dict[str, Any]
|
||
|
||
|
||
class SandboxClient:
|
||
"""沙箱服务异步客户端。
|
||
|
||
授权语义：调用方需先经司命审核（在本类之外完成），通过 ``AuthorizationContext`` 携带
|
||
授权上下文；本类据此签发短期 HMAC 凭证（exec 含本次请求的 body_hash 绑定），放入
|
||
``Authorization`` 头，由沙箱端校验。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
http_client: httpx.AsyncClient,
|
||
settings: SandboxSettings,
|
||
event_sink: EventSink | None = None,
|
||
) -> None:
|
||
if settings.sandbox_auth_enabled and not settings.sandbox_auth_secret:
|
||
raise ValueError("auth_enabled=True 时 auth_secret 不能为空（须与沙箱端共享）")
|
||
self._http = http_client
|
||
self._settings = settings
|
||
self._sink: EventSink = event_sink or LoggingEventSink()
|
||
|
||
async def health(self) -> bool:
|
||
"""健康探测，失败返回 False 而不抛异常。"""
|
||
try:
|
||
resp = await self._http.get(f"{self._settings.rust_sandbox_url}/healthz", timeout=5.0)
|
||
except httpx.HTTPError:
|
||
return False
|
||
return resp.status_code == 200
|
||
|
||
async def exec(
|
||
self,
|
||
request: ExecRequest,
|
||
auth_ctx: AuthorizationContext,
|
||
*,
|
||
event_sink: EventSink | None = None,
|
||
) -> ExecResponse:
|
||
"""同步执行：收集全部输出后返回。授权失败、网络错误或沙箱错误分别抛出对应异常，
|
||
并回写 tool.failed 事件。``event_sink`` 可按调用覆盖，便于编排层按 run 绑定上下文。"""
|
||
sink = event_sink or self._sink
|
||
trace_id = request.trace_id or auth_ctx.trace_id
|
||
run_id = auth_ctx.run_id
|
||
tool = auth_ctx.tool
|
||
|
||
# 请求体带上解析后的 trace_id，便于沙箱日志关联
|
||
if request.trace_id is None and trace_id is not None:
|
||
request.trace_id = trace_id
|
||
|
||
await sink.emit(
|
||
started_event(
|
||
tool=tool,
|
||
argv=request.argv or [f"code:{request.language}", request.code],
|
||
risk_level=auth_ctx.risk_level,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
)
|
||
)
|
||
|
||
try:
|
||
resp = await self._http.post(
|
||
f"{self._settings.rust_sandbox_url}/v1/exec",
|
||
json=request.model_dump(exclude_none=True),
|
||
headers=self._exec_auth_headers(auth_ctx, request),
|
||
timeout=self._settings.http_timeout(request.timeout_ms),
|
||
)
|
||
except httpx.HTTPError as exc:
|
||
await sink.emit(
|
||
failed_event(
|
||
tool=tool,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
error_code="sandbox_unavailable",
|
||
error_message=str(exc),
|
||
)
|
||
)
|
||
raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc
|
||
|
||
if resp.status_code >= 400:
|
||
exc_type: type[SandboxError]
|
||
if resp.status_code == 401:
|
||
exc_type = SandboxUnauthorized
|
||
elif resp.status_code == 400:
|
||
exc_type = SandboxBadRequest
|
||
else:
|
||
exc_type = SandboxExecFailed
|
||
await self._fail(tool, trace_id, run_id, resp, exc_type, sink)
|
||
|
||
result = ExecResponse.model_validate(resp.json())
|
||
await sink.emit(
|
||
completed_event(
|
||
tool=tool,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
exit_code=result.exit_code,
|
||
duration_ms=result.duration_ms,
|
||
timeout=result.timeout,
|
||
truncated=result.truncated,
|
||
stdout_len=len(result.stdout),
|
||
stderr_len=len(result.stderr),
|
||
)
|
||
)
|
||
return result
|
||
|
||
async def exec_stream(
|
||
self,
|
||
request: ExecRequest,
|
||
auth_ctx: AuthorizationContext,
|
||
) -> AsyncIterator[SandboxStreamEvent]:
|
||
"""流式执行：实时产出 stdout/stderr 增量（base64 str）与 done 汇总（dict）。
|
||
|
||
消费沙箱 ``POST /v1/exec/stream`` 的 SSE。不回写一炁事件——started/completed/failed
|
||
由编排层 ``ToolRuntime.execute_stream`` 负责，本方法只透传执行级输出。
|
||
"""
|
||
if request.trace_id is None and auth_ctx.trace_id is not None:
|
||
request.trace_id = auth_ctx.trace_id
|
||
|
||
try:
|
||
async with self._http.stream(
|
||
"POST",
|
||
f"{self._settings.rust_sandbox_url}/v1/exec/stream",
|
||
json=request.model_dump(exclude_none=True),
|
||
headers=self._exec_auth_headers(auth_ctx, request),
|
||
timeout=self._settings.http_timeout(request.timeout_ms),
|
||
) as resp:
|
||
if resp.status_code == 401:
|
||
raise SandboxUnauthorized(await self._read_error(resp))
|
||
if resp.status_code == 400:
|
||
raise SandboxBadRequest(await self._read_error(resp))
|
||
if resp.status_code >= 400:
|
||
raise SandboxExecFailed(await self._read_error(resp))
|
||
async for event_type, data in self._iter_sse(resp):
|
||
if event_type in ("stdout", "stderr"):
|
||
yield SandboxStreamEvent(event_type, data)
|
||
elif event_type == "done":
|
||
try:
|
||
yield SandboxStreamEvent("done", json.loads(data))
|
||
except ValueError:
|
||
continue
|
||
except httpx.HTTPError as exc:
|
||
raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc
|
||
|
||
async def cancel(
|
||
self,
|
||
trace_id: str,
|
||
auth_ctx: AuthorizationContext,
|
||
) -> bool:
|
||
"""取消运行中的任务（按 trace_id）。命中返回 True；任务不存在或已完成抛 SandboxNotFound。
|
||
|
||
cancel 不回写一炁事件（由编排层按需记录），凭证不绑定 body_hash（危害有限）。
|
||
"""
|
||
try:
|
||
resp = await self._http.post(
|
||
f"{self._settings.rust_sandbox_url}/v1/cancel",
|
||
json={"trace_id": trace_id},
|
||
headers=self._auth_headers(auth_ctx),
|
||
timeout=10.0,
|
||
)
|
||
except httpx.HTTPError as exc:
|
||
raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc
|
||
|
||
if resp.status_code == 404:
|
||
raise SandboxNotFound(f"任务不存在或已完成: {trace_id}")
|
||
if resp.status_code >= 400:
|
||
code, message = self._extract_error(resp)
|
||
if resp.status_code == 401:
|
||
raise SandboxUnauthorized(message)
|
||
raise SandboxExecFailed(message, code=code)
|
||
return CancelResponse.model_validate(resp.json()).cancelled
|
||
|
||
# ---- 文件 IO（不经 shell，二进制安全；与 Rust fileio.rs 契约对齐） ----
|
||
|
||
async def file_read(
|
||
self,
|
||
request: FileReadRequest,
|
||
auth_ctx: AuthorizationContext,
|
||
*,
|
||
event_sink: EventSink | None = None,
|
||
) -> FileReadResponse:
|
||
return await self._file_op(
|
||
"read", "/v1/file/read", request, FileReadResponse, auth_ctx, event_sink=event_sink
|
||
)
|
||
|
||
async def file_write(
|
||
self,
|
||
request: FileWriteRequest,
|
||
auth_ctx: AuthorizationContext,
|
||
*,
|
||
event_sink: EventSink | None = None,
|
||
) -> FileWriteResponse:
|
||
return await self._file_op(
|
||
"write", "/v1/file/write", request, FileWriteResponse, auth_ctx, event_sink=event_sink
|
||
)
|
||
|
||
async def file_list(
|
||
self,
|
||
request: FileListRequest,
|
||
auth_ctx: AuthorizationContext,
|
||
*,
|
||
event_sink: EventSink | None = None,
|
||
) -> FileListResponse:
|
||
return await self._file_op(
|
||
"list", "/v1/file/list", request, FileListResponse, auth_ctx, event_sink=event_sink
|
||
)
|
||
|
||
async def file_stat(
|
||
self,
|
||
request: FileStatRequest,
|
||
auth_ctx: AuthorizationContext,
|
||
*,
|
||
event_sink: EventSink | None = None,
|
||
) -> FileStatResponse:
|
||
return await self._file_op(
|
||
"stat", "/v1/file/stat", request, FileStatResponse, auth_ctx, event_sink=event_sink
|
||
)
|
||
|
||
async def _file_op(
|
||
self,
|
||
op: str,
|
||
endpoint: str,
|
||
request: FileReadRequest | FileWriteRequest | FileListRequest | FileStatRequest,
|
||
response_cls: type[_T],
|
||
auth_ctx: AuthorizationContext,
|
||
*,
|
||
event_sink: EventSink | None = None,
|
||
) -> _T:
|
||
"""文件操作通用流程：签发 file_op 凭证 → POST → 解析响应，并回写一炁事件。
|
||
|
||
与 exec 共用错误映射（``_fail``）与事件构造；argv 记 ``[file:<op>, <path>]``
|
||
供审计 input_summary。文件操作在主进程运行（非 exec 子进程），无超时参数，
|
||
用默认 HTTP 超时（``http_timeout(None)``）。
|
||
"""
|
||
sink = event_sink or self._sink
|
||
trace_id = auth_ctx.trace_id
|
||
run_id = auth_ctx.run_id
|
||
tool = auth_ctx.tool
|
||
path = request.path
|
||
|
||
await sink.emit(
|
||
started_event(
|
||
tool=tool,
|
||
argv=[f"file:{op}", path],
|
||
risk_level=auth_ctx.risk_level,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
)
|
||
)
|
||
|
||
t0 = time.monotonic()
|
||
try:
|
||
resp = await self._http.post(
|
||
f"{self._settings.rust_sandbox_url}{endpoint}",
|
||
json=request.model_dump(exclude_none=True),
|
||
headers=self._file_auth_headers(auth_ctx, op, path),
|
||
timeout=self._settings.http_timeout(None),
|
||
)
|
||
except httpx.HTTPError as exc:
|
||
await sink.emit(
|
||
failed_event(
|
||
tool=tool,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
error_code="sandbox_unavailable",
|
||
error_message=str(exc),
|
||
)
|
||
)
|
||
raise SandboxUnavailable(f"沙箱不可达: {exc}") from exc
|
||
|
||
if resp.status_code >= 400:
|
||
exc_type: type[SandboxError]
|
||
if resp.status_code == 401:
|
||
exc_type = SandboxUnauthorized
|
||
elif resp.status_code == 400:
|
||
exc_type = SandboxBadRequest
|
||
elif resp.status_code == 404:
|
||
exc_type = SandboxNotFound
|
||
else:
|
||
exc_type = SandboxExecFailed
|
||
await self._fail(tool, trace_id, run_id, resp, exc_type, sink)
|
||
|
||
result = response_cls.model_validate(resp.json())
|
||
# completed 事件：read 带 returned 字节数 + 截断标记，write 带 bytes_written，
|
||
# list/stat 无字节数语义，以 0 占位（exit_code=0 表示操作成功）。
|
||
stdout_len = getattr(result, "returned", None) or getattr(result, "bytes_written", 0)
|
||
truncated = bool(getattr(result, "truncated", False))
|
||
await sink.emit(
|
||
completed_event(
|
||
tool=tool,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
exit_code=0,
|
||
duration_ms=int((time.monotonic() - t0) * 1000),
|
||
timeout=False,
|
||
truncated=truncated,
|
||
stdout_len=stdout_len,
|
||
stderr_len=0,
|
||
)
|
||
)
|
||
return result
|
||
|
||
def _file_auth_headers(
|
||
self,
|
||
auth_ctx: AuthorizationContext,
|
||
op: str,
|
||
path: str,
|
||
) -> dict[str, str]:
|
||
"""签发文件操作凭证：file_op 绑定 {op, path}（exec 路径 body_hash 为 None，不绑定）。"""
|
||
if not self._settings.sandbox_auth_enabled:
|
||
return {}
|
||
token = sign_authorization(
|
||
self._settings.sandbox_auth_secret,
|
||
auth_ctx,
|
||
argv=None,
|
||
stdin=None,
|
||
ttl_seconds=self._settings.sandbox_auth_ttl_seconds,
|
||
file_op=(op, path),
|
||
)
|
||
return {"Authorization": f"Bearer {token}"}
|
||
|
||
def _auth_headers(
|
||
self,
|
||
auth_ctx: AuthorizationContext,
|
||
argv: list[str] | None = None,
|
||
stdin: str | None = None,
|
||
*,
|
||
code: str | None = None,
|
||
language: str | None = None,
|
||
) -> dict[str, str]:
|
||
"""签发授权凭证。
|
||
|
||
- exec 命令模式：传 argv/stdin，按 compute_body_hash 绑定。
|
||
- exec 代码模式：传 code/language/stdin，按 compute_body_hash_code 绑定。
|
||
- cancel：都不传，不绑定 body_hash。
|
||
"""
|
||
if not self._settings.sandbox_auth_enabled:
|
||
return {}
|
||
token = sign_authorization(
|
||
self._settings.sandbox_auth_secret,
|
||
auth_ctx,
|
||
argv=argv,
|
||
stdin=stdin,
|
||
code=code,
|
||
language=language,
|
||
ttl_seconds=self._settings.sandbox_auth_ttl_seconds,
|
||
)
|
||
return {"Authorization": f"Bearer {token}"}
|
||
|
||
def _exec_auth_headers(
|
||
self,
|
||
auth_ctx: AuthorizationContext,
|
||
request: ExecRequest,
|
||
) -> dict[str, str]:
|
||
"""按 code/argv 模式分流签发 exec 授权头。"""
|
||
if request.code is not None:
|
||
return self._auth_headers(
|
||
auth_ctx,
|
||
stdin=request.stdin,
|
||
code=request.code,
|
||
language=request.language,
|
||
)
|
||
return self._auth_headers(auth_ctx, request.argv, request.stdin)
|
||
|
||
@staticmethod
|
||
async def _read_error(resp: httpx.Response) -> str:
|
||
"""流式响应的错误体：先 aread 再复用 _extract_error。"""
|
||
await resp.aread()
|
||
_, message = SandboxClient._extract_error(resp)
|
||
return message
|
||
|
||
@staticmethod
|
||
async def _iter_sse(resp: httpx.Response) -> AsyncIterator[tuple[str, str]]:
|
||
"""极简 SSE 解析：按空行分界，聚合 ``event:`` 与 ``data:`` 行（忽略 id/retry/注释/心跳）。"""
|
||
event_type = "message"
|
||
data_lines: list[str] = []
|
||
async for line in resp.aiter_lines():
|
||
if line == "":
|
||
if data_lines:
|
||
yield event_type, "\n".join(data_lines)
|
||
event_type = "message"
|
||
data_lines = []
|
||
elif line.startswith("event:"):
|
||
event_type = line[len("event:") :].strip()
|
||
elif line.startswith("data:"):
|
||
data_lines.append(line[len("data:") :].lstrip())
|
||
# 其余行（id:/retry:/注释/心跳）忽略
|
||
if data_lines:
|
||
yield event_type, "\n".join(data_lines)
|
||
|
||
async def _fail(
|
||
self,
|
||
tool: str,
|
||
trace_id: str | None,
|
||
run_id: str | None,
|
||
resp: httpx.Response,
|
||
exc_type: type[SandboxError],
|
||
sink: EventSink,
|
||
) -> NoReturn:
|
||
code, message = self._extract_error(resp)
|
||
await sink.emit(
|
||
failed_event(
|
||
tool=tool,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
error_code=code,
|
||
error_message=message,
|
||
)
|
||
)
|
||
raise exc_type(message)
|
||
|
||
@staticmethod
|
||
def _extract_error(resp: httpx.Response) -> tuple[str, str]:
|
||
try:
|
||
body = resp.json()
|
||
except ValueError:
|
||
return "exec_failed", resp.text
|
||
err = body.get("error") if isinstance(body, dict) else None
|
||
if isinstance(err, dict):
|
||
return str(err.get("code", "exec_failed")), str(err.get("message", resp.text))
|
||
return "exec_failed", resp.text
|