530 lines
21 KiB
Python
530 lines
21 KiB
Python
"""百工 Tool Runtime：司命评估 → 直接执行/沙箱执行 → 一炁事件回写 的编排核心。
|
||
|
||
调用链：
|
||
ToolRuntime.execute / execute_stream
|
||
→ SafetyGovernor.evaluate（deny→记 tool.failed 并抛 PermissionDenied；
|
||
ask→抛 RequiresConfirmation 交上层确认后重试）
|
||
→ 执行路径二选一：
|
||
- ToolSpec.executor 非 None：构造执行上下文（含 http_client 等）
|
||
直接调用 executor，跳过沙箱（HTTP 类工具如 fetch_url）。
|
||
- executor 为 None：ToolSpec.build_request 构造沙箱请求（isolated_workspace 工具按
|
||
run 隔离 workdir）→ SandboxClient.exec / exec_stream（携带司命授权上下文；
|
||
事件经 EventStoreSink 落 EventStore）。
|
||
→ ToolResult / 流式 ToolStreamEvent
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import re
|
||
import uuid
|
||
from collections.abc import AsyncIterator
|
||
from contextlib import suppress
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
import httpx
|
||
|
||
from app.db.session import AsyncSessionLocal
|
||
from app.integrations.sandbox.audit_sink import AuditEventSink
|
||
from app.integrations.sandbox.auth import AuthorizationContext
|
||
from app.integrations.sandbox.bridge import EventStoreSink
|
||
from app.integrations.sandbox.client import SandboxClient
|
||
from app.integrations.sandbox.events import (
|
||
CompositeEventSink,
|
||
EventSink,
|
||
completed_event,
|
||
failed_event,
|
||
started_event,
|
||
)
|
||
from app.integrations.sandbox.exceptions import SandboxError
|
||
from app.integrations.sandbox.models import (
|
||
DirEntry,
|
||
ExecResponse,
|
||
FileReadRequest,
|
||
FileReadResponse,
|
||
)
|
||
from app.modules.runtime.events import Event, EventStore
|
||
from app.modules.safety.exceptions import PermissionDenied, RequiresConfirmation
|
||
from app.modules.safety.governor import SafetyDecision, SafetyGovernor
|
||
from app.modules.tool_runtime.definition import ToolInput, ToolResult, ToolSpec
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ToolStreamEvent:
|
||
"""工具流式事件：stdout/stderr 增量（base64 str，透传给前端解码）或 done（结束汇总 dict）。"""
|
||
|
||
kind: str # "stdout" | "stderr" | "done"
|
||
payload: str | dict[str, Any]
|
||
|
||
|
||
class ToolRuntime:
|
||
"""百工工具运行时。"""
|
||
|
||
_ISOLATE_SAFE = re.compile(r"[^A-Za-z0-9_.-]")
|
||
|
||
def __init__(
|
||
self,
|
||
sandbox: SandboxClient,
|
||
governor: SafetyGovernor,
|
||
event_store: EventStore,
|
||
specs: dict[str, ToolSpec] | None = None,
|
||
*,
|
||
http_client: httpx.AsyncClient | None = None,
|
||
) -> None:
|
||
self._sandbox = sandbox
|
||
self._governor = governor
|
||
self._store = event_store
|
||
self._http_client = http_client
|
||
# 不在启动期快照 builtin_specs()——那样删除/修改内置工具后需重启后端才生效。
|
||
# builtin 工具每次经 _resolve_specs() 实时取（builtin_specs 是纯函数，返回新 dict）。
|
||
# 调用方显式传入的 specs 视为"覆盖层"，与 register 的工具一起叠加在 builtin 之上。
|
||
self._custom_specs: dict[str, ToolSpec] = dict(specs) if specs else {}
|
||
self._registered_specs: dict[str, ToolSpec] = {}
|
||
|
||
@property
|
||
def http_client(self) -> httpx.AsyncClient | None:
|
||
"""共享 HTTP 客户端（供 MCP executor 等复用，生命周期同 runtime）。"""
|
||
return self._http_client
|
||
|
||
def _resolve_specs(self) -> dict[str, ToolSpec]:
|
||
"""合并 builtin（实时）+ custom（调用方传入）+ registered（动态注册）。
|
||
|
||
后者覆盖前者同名，保证覆盖语义；builtin 每次现取，builtin_specs() 的任何变更
|
||
（如新增/删除/修改内置工具）立即对当前 ToolRuntime 实例生效，无需重启。
|
||
"""
|
||
from app.application.tool_runtime.builtins import builtin_specs
|
||
|
||
merged: dict[str, ToolSpec] = {}
|
||
merged.update(builtin_specs())
|
||
merged.update(self._custom_specs)
|
||
merged.update(self._registered_specs)
|
||
return merged
|
||
|
||
def register(self, spec: ToolSpec) -> None:
|
||
"""注册或覆盖工具规格（叠加在 builtin 与调用方传入的 specs 之上）。"""
|
||
self._registered_specs[spec.name] = spec
|
||
|
||
def unregister(self, name: str) -> None:
|
||
"""移除动态注册的工具（builtin / custom 不受影响）。"""
|
||
self._registered_specs.pop(name, None)
|
||
|
||
async def reload_governor(self) -> None:
|
||
"""从 tool_policies 表重新加载司命策略（写操作后调用，使变更即时生效）。
|
||
|
||
读取所有 enabled 策略构造 SafetyGovernor 覆盖当前实例。governor 的 policy
|
||
字典作为覆盖层（见 SafetyGovernor.evaluate）：命中的工具用表里的 risk_level /
|
||
permission_policy，未命中的用 ToolSpec 自带 risk_level 推导默认策略。
|
||
"""
|
||
from app.db.session import AsyncSessionLocal
|
||
from app.repositories import governance as governance_repo
|
||
|
||
async with AsyncSessionLocal() as session:
|
||
policies = await governance_repo.list_tool_policies(session)
|
||
data = {
|
||
"tools": {
|
||
p.name: {
|
||
"risk_level": p.risk_level,
|
||
"permission_policy": p.permission_policy,
|
||
}
|
||
for p in policies
|
||
if p.enabled
|
||
}
|
||
}
|
||
self._governor = SafetyGovernor.from_policy_dict(data)
|
||
|
||
def list_specs(self) -> list[ToolSpec]:
|
||
"""返回当前可用工具规格（用于 API 暴露工具元数据）。
|
||
|
||
builtin 实时合并：删除/修改 builtin_specs() 中的工具定义后无需重启即可生效。
|
||
"""
|
||
return list(self._resolve_specs().values())
|
||
|
||
async def execute(
|
||
self,
|
||
tool: str,
|
||
input: ToolInput,
|
||
*,
|
||
run_id: str | None = None,
|
||
character_id: str | None = None,
|
||
trace_id: str | None = None,
|
||
confirmed: bool = False,
|
||
context: dict[str, Any] | None = None,
|
||
) -> ToolResult:
|
||
cancellation = (context or {}).get("cancellation")
|
||
if cancellation is not None:
|
||
cancellation.raise_if_cancelled()
|
||
spec = self._resolve_specs().get(tool)
|
||
if spec is None:
|
||
await self._record_failed(
|
||
tool, run_id, character_id, trace_id, "permission_denied", f"未知工具: {tool}"
|
||
)
|
||
raise PermissionDenied(f"未知工具: {tool}")
|
||
|
||
decision = self._governor.evaluate(
|
||
tool=tool, risk_level=spec.risk_level, confirmed=confirmed
|
||
)
|
||
if decision.decision is SafetyDecision.deny:
|
||
await self._record_failed(
|
||
tool, run_id, character_id, trace_id, "permission_denied", decision.reason
|
||
)
|
||
raise PermissionDenied(decision.reason)
|
||
if decision.decision is SafetyDecision.ask:
|
||
# 不执行；上层确认后以 confirmed=True 重试。携带 risk_level 供前端展示。
|
||
raise RequiresConfirmation(decision.reason, risk_level=decision.risk_level)
|
||
|
||
# ---- 直接执行路径（executor 工具，如 fetch_url）----
|
||
if spec.executor is not None:
|
||
exec_context = self._build_exec_context(context, trace_id=trace_id)
|
||
sink = self._make_sink(run_id, character_id)
|
||
await sink.emit(
|
||
started_event(
|
||
tool=spec.name,
|
||
argv=[],
|
||
risk_level=spec.risk_level,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
)
|
||
)
|
||
try:
|
||
result = await spec.executor(input, exec_context)
|
||
except Exception as exc: # noqa: BLE001 - executor 任意错误统一兜底
|
||
await sink.emit(
|
||
failed_event(
|
||
tool=spec.name,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
error_code="executor_failed",
|
||
error_message=str(exc),
|
||
)
|
||
)
|
||
raise SandboxError(f"工具 {spec.name} 直接执行失败: {exc}") from exc
|
||
|
||
await sink.emit(
|
||
completed_event(
|
||
tool=spec.name,
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
exit_code=result.exit_code if result.exit_code is not None else 0,
|
||
duration_ms=result.duration_ms,
|
||
timeout=result.timeout,
|
||
truncated=result.truncated,
|
||
stdout_len=len(result.stdout),
|
||
stderr_len=len(result.stderr),
|
||
)
|
||
)
|
||
# 补齐 trace_id（executor 可能未填）
|
||
if not result.trace_id:
|
||
result = ToolResult(
|
||
ok=result.ok,
|
||
exit_code=result.exit_code,
|
||
stdout=result.stdout,
|
||
stderr=result.stderr,
|
||
duration_ms=result.duration_ms,
|
||
timeout=result.timeout,
|
||
truncated=result.truncated,
|
||
trace_id=trace_id or "",
|
||
)
|
||
return result
|
||
|
||
# ---- 沙箱执行路径 ----
|
||
request = spec.build_request(input)
|
||
if trace_id and not request.trace_id:
|
||
request.trace_id = trace_id
|
||
# 文件类工具按角色隔离工作目录（同一角色跨 run 文件互通，不同角色互不可见）
|
||
if spec.isolated_workspace:
|
||
request.workdir = self._isolate_workdir(
|
||
request.workdir, subject_id=character_id, run_id=run_id, trace_id=trace_id
|
||
)
|
||
|
||
auth_ctx = AuthorizationContext(
|
||
trace_id=request.trace_id or trace_id or "",
|
||
risk_level=spec.risk_level,
|
||
tool=spec.name,
|
||
character_id=character_id,
|
||
run_id=run_id,
|
||
)
|
||
# 本次调用绑定带 run/character 上下文的 sink，使沙箱事件落入一炁事件存储与审计表
|
||
sink = self._make_sink(run_id, character_id)
|
||
|
||
try:
|
||
resp: ExecResponse = await self._sandbox.exec(request, auth_ctx, event_sink=sink)
|
||
except asyncio.CancelledError:
|
||
# 取消 HTTP await 不代表远端子进程已停止；尽力调用沙箱 cancel，随后保持取消语义。
|
||
if request.trace_id:
|
||
with suppress(Exception):
|
||
await asyncio.shield(self._sandbox.cancel(request.trace_id, auth_ctx))
|
||
raise
|
||
except SandboxError:
|
||
# 沙箱错误事件已由 sink 记录（tool.failed），此处不重复写入
|
||
raise
|
||
|
||
return ToolResult(
|
||
ok=resp.ok,
|
||
exit_code=resp.exit_code,
|
||
stdout=resp.stdout,
|
||
stderr=resp.stderr,
|
||
duration_ms=resp.duration_ms,
|
||
timeout=resp.timeout,
|
||
truncated=resp.truncated,
|
||
trace_id=resp.trace_id,
|
||
)
|
||
|
||
async def execute_stream(
|
||
self,
|
||
tool: str,
|
||
input: ToolInput,
|
||
*,
|
||
run_id: str | None = None,
|
||
character_id: str | None = None,
|
||
trace_id: str | None = None,
|
||
confirmed: bool = False,
|
||
) -> AsyncIterator[ToolStreamEvent]:
|
||
"""流式执行：司命评估后转发沙箱 stdout/stderr 增量，末尾一个 done。
|
||
|
||
started/completed/failed 落一炁事件存储；stdout/stderr 增量只透传不落库。
|
||
司命 deny 记 tool.failed 并抛 PermissionDenied；ask 抛 RequiresConfirmation
|
||
（流式端点要求调用方先经非流式 invoke 完成 confirmation，再以 confirmed=True 调用）。
|
||
"""
|
||
spec = self._resolve_specs().get(tool)
|
||
if spec is None:
|
||
await self._record_failed(
|
||
tool, run_id, character_id, trace_id, "permission_denied", f"未知工具: {tool}"
|
||
)
|
||
raise PermissionDenied(f"未知工具: {tool}")
|
||
|
||
decision = self._governor.evaluate(
|
||
tool=tool, risk_level=spec.risk_level, confirmed=confirmed
|
||
)
|
||
if decision.decision is SafetyDecision.deny:
|
||
await self._record_failed(
|
||
tool, run_id, character_id, trace_id, "permission_denied", decision.reason
|
||
)
|
||
raise PermissionDenied(decision.reason)
|
||
if decision.decision is SafetyDecision.ask:
|
||
raise RequiresConfirmation(decision.reason, risk_level=decision.risk_level)
|
||
|
||
request = spec.build_request(input)
|
||
if trace_id and not request.trace_id:
|
||
request.trace_id = trace_id
|
||
if spec.isolated_workspace:
|
||
request.workdir = self._isolate_workdir(
|
||
request.workdir, subject_id=character_id, run_id=run_id, trace_id=trace_id
|
||
)
|
||
|
||
effective_trace = request.trace_id or trace_id or ""
|
||
auth_ctx = AuthorizationContext(
|
||
trace_id=effective_trace,
|
||
risk_level=spec.risk_level,
|
||
tool=spec.name,
|
||
character_id=character_id,
|
||
run_id=run_id,
|
||
)
|
||
sink = self._make_sink(run_id, character_id)
|
||
|
||
await sink.emit(
|
||
started_event(
|
||
tool=spec.name,
|
||
argv=request.argv,
|
||
risk_level=spec.risk_level,
|
||
trace_id=effective_trace or None,
|
||
run_id=run_id,
|
||
)
|
||
)
|
||
|
||
final: dict[str, Any] = {}
|
||
try:
|
||
async for ev in self._sandbox.exec_stream(request, auth_ctx):
|
||
if ev.kind in ("stdout", "stderr"):
|
||
yield ToolStreamEvent(ev.kind, ev.payload)
|
||
elif ev.kind == "done" and isinstance(ev.payload, dict):
|
||
final = ev.payload
|
||
except SandboxError:
|
||
await sink.emit(
|
||
failed_event(
|
||
tool=spec.name,
|
||
trace_id=effective_trace or None,
|
||
run_id=run_id,
|
||
error_code="sandbox_error",
|
||
error_message="沙箱流式执行失败",
|
||
)
|
||
)
|
||
raise
|
||
|
||
await sink.emit(
|
||
completed_event(
|
||
tool=spec.name,
|
||
trace_id=effective_trace or None,
|
||
run_id=run_id,
|
||
exit_code=final.get("exit_code"),
|
||
duration_ms=int(final.get("duration_ms", 0)),
|
||
timeout=bool(final.get("timeout", False)),
|
||
truncated=bool(final.get("truncated", False)),
|
||
stdout_len=0, # 流式不收集完整文本
|
||
stderr_len=0,
|
||
)
|
||
)
|
||
yield ToolStreamEvent("done", final)
|
||
|
||
def _make_sink(self, run_id: str | None, character_id: str | None) -> EventSink:
|
||
"""构造本次执行的事件 sink 组合：一炁流程事件 + 治理审计，互不感知。
|
||
|
||
EventStoreSink 落 events 表（流程，run 内 sequence），AuditEventSink 落 audit_events
|
||
表（治理审计，含 decision / risk_level / argv 脱敏摘要）。两者经 CompositeEventSink 扇出。
|
||
"""
|
||
return CompositeEventSink(
|
||
[
|
||
EventStoreSink(self._store, run_id=run_id, character_id=character_id),
|
||
AuditEventSink(AsyncSessionLocal, run_id=run_id, character_id=character_id),
|
||
]
|
||
)
|
||
|
||
def _isolate_workdir(
|
||
self,
|
||
base: str | None,
|
||
*,
|
||
subject_id: str | None,
|
||
run_id: str | None,
|
||
trace_id: str | None,
|
||
) -> str:
|
||
"""把工具 workspace 根目录隔离到「主体」级子目录，防跨主体信息泄漏。
|
||
|
||
隔离主体 subject_id：当前为 character_id（一个角色一个持久 workspace，跨会话/
|
||
跨 run 文件互通，符合"一个角色拥有其全部文件"的语义）；多用户上线后应改为真实
|
||
用户标识（owner_scope / user_id），届时仅需调用方改传参，本函数逻辑不变。
|
||
|
||
run_id / trace_id 仅在 subject_id 缺失时回退（如脱离聊天流程的直接调用），
|
||
保证既有隔离又不至于因缺主体而落到共享根目录。
|
||
|
||
subject_id 经 sanitize（仅保留 ``[A-Za-z0-9_.-]``）从源头杜绝路径注入；
|
||
沙箱端另有 canonicalize + ``starts_with(base)`` 兜底。
|
||
"""
|
||
root = base or "workspace"
|
||
key = subject_id or run_id or trace_id or f"adhoc_{uuid.uuid4().hex[:8]}"
|
||
safe = self._ISOLATE_SAFE.sub("_", key).strip("._-") or "adhoc"
|
||
return f"{root}/{safe}"
|
||
|
||
def _build_exec_context(
|
||
self, context: dict[str, Any] | None, *, trace_id: str | None
|
||
) -> dict[str, Any]:
|
||
"""构造 executor 执行上下文：合并 runtime 持有的 http_client、settings 中的密钥等。
|
||
|
||
调用方传入的 context（per-request 覆盖）优先于 runtime 默认值。
|
||
"""
|
||
exec_context: dict[str, Any] = {"trace_id": trace_id}
|
||
if self._http_client is not None:
|
||
exec_context["http_client"] = self._http_client
|
||
if context:
|
||
exec_context.update(context)
|
||
return exec_context
|
||
|
||
async def read_sandbox_file(
|
||
self,
|
||
path: str,
|
||
*,
|
||
run_id: str | None = None,
|
||
character_id: str | None = None,
|
||
trace_id: str | None = None,
|
||
isolated: bool = True,
|
||
) -> FileReadResponse:
|
||
"""读取沙箱内文件（仅供文件预览端点调用）。
|
||
|
||
Args:
|
||
path: 文件路径。
|
||
isolated=True 时视为工作区相对路径（write_file 传入的 path），
|
||
自动补全 ``workspace/<character_id>/`` 隔离前缀；
|
||
isolated=False 时作为沙箱 base 下的完整相对路径。
|
||
run_id: 关联的 run ID（可选，用于授权上下文）。
|
||
character_id: 人物 ID（可选）。提供时自动构造隔离工作区路径。
|
||
trace_id: 追踪 ID（可选，自动生成）。
|
||
isolated: 是否按 isolated_workspace 规则补全路径前缀。
|
||
|
||
Returns:
|
||
FileReadResponse，含 base64url-no-pad 编码的文件内容。
|
||
"""
|
||
effective_path = path
|
||
if isolated:
|
||
workdir = self._isolate_workdir(
|
||
"workspace", subject_id=character_id, run_id=run_id, trace_id=trace_id
|
||
)
|
||
effective_path = f"{workdir}/{path.lstrip('/')}"
|
||
|
||
effective_trace = trace_id or str(uuid.uuid4())
|
||
auth_ctx = AuthorizationContext(
|
||
trace_id=effective_trace,
|
||
risk_level="L0",
|
||
tool="file_read",
|
||
character_id=character_id,
|
||
run_id=run_id,
|
||
)
|
||
return await self._sandbox.file_read(
|
||
FileReadRequest(path=effective_path),
|
||
auth_ctx,
|
||
)
|
||
|
||
async def list_sandbox_dir(
|
||
self,
|
||
path: str = ".",
|
||
*,
|
||
run_id: str | None = None,
|
||
character_id: str | None = None,
|
||
trace_id: str | None = None,
|
||
isolated: bool = True,
|
||
) -> list[DirEntry]:
|
||
"""列出沙箱工作区目录（仅供文件浏览端点调用）。
|
||
|
||
Args:
|
||
path: 目录路径，默认工作区根。
|
||
isolated: True 时自动应用 isolated_workspace 隔离前缀。
|
||
|
||
返回排序后的目录条目列表。
|
||
"""
|
||
from app.integrations.sandbox.models import FileListRequest
|
||
|
||
effective_path = path
|
||
if isolated:
|
||
workdir = self._isolate_workdir(
|
||
"workspace", subject_id=character_id, run_id=run_id, trace_id=trace_id
|
||
)
|
||
effective_path = f"{workdir}/{path.lstrip('/')}"
|
||
|
||
effective_trace = trace_id or str(uuid.uuid4())
|
||
auth_ctx = AuthorizationContext(
|
||
trace_id=effective_trace,
|
||
risk_level="L0",
|
||
tool="file_list",
|
||
character_id=character_id,
|
||
run_id=run_id,
|
||
)
|
||
result = await self._sandbox.file_list(
|
||
FileListRequest(path=effective_path),
|
||
auth_ctx,
|
||
)
|
||
return result.entries
|
||
|
||
async def _record_failed(
|
||
self,
|
||
tool: str,
|
||
run_id: str | None,
|
||
character_id: str | None,
|
||
trace_id: str | None,
|
||
error_code: str,
|
||
error_message: str,
|
||
) -> None:
|
||
await self._store.append(
|
||
Event(
|
||
type="tool.failed",
|
||
source="tool",
|
||
visibility="public",
|
||
trace_id=trace_id,
|
||
run_id=run_id,
|
||
character_id=character_id,
|
||
payload={
|
||
"tool": tool,
|
||
"phase": "safety",
|
||
"error_code": error_code,
|
||
"error_message": error_message,
|
||
},
|
||
)
|
||
)
|