初始化提交

2026-08-21 16:00:13 +08:00
commit 887c1f449e
@@ -0,0 +1,475 @@
+"""沙箱客户端：封装对 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