|
|
@@ -0,0 +1,342 @@
|
|
|
+// 脚本调用工具（docs/design/17-run-script.md）：模型写一段 Python 脚本，由网关在沙箱中执行；脚本通过 agent_tools 模块调用
|
|
|
+// MCP 工具（每次调用都经过网关：工具策略、结果处理与普通调用相同，只是不经过模型），循环、汇总、写文件都在脚本中完成，
|
|
|
+// 只有脚本打印的内容交回模型。脚本与网关之间用两条管道（fd 3 发请求、fd 4 收回复）传 JSON 行
|
|
|
+import { spawn } from 'node:child_process'
|
|
|
+import crypto from 'node:crypto'
|
|
|
+import fs from 'node:fs'
|
|
|
+import path from 'node:path'
|
|
|
+import type { ScriptProgressPayload } from '@agent-server/shared'
|
|
|
+import { z } from 'zod'
|
|
|
+import { config } from '../config'
|
|
|
+import { sandboxArgs } from '../services/office'
|
|
|
+import type { ToolResult } from './builtin'
|
|
|
+import type { GatewayRunContext } from './tokens'
|
|
|
+
|
|
|
+export const RUN_SCRIPT = 'run_script'
|
|
|
+
|
|
|
+/** 交给模型的输出上限；完整输出保存为文件 */
|
|
|
+const MAX_OUTPUT = 20_000
|
|
|
+const MAX_STDERR = 4000
|
|
|
+/** 每次调用返回给脚本的文本上限 */
|
|
|
+const MAX_CALL_TEXT = 1_000_000
|
|
|
+const PROGRESS_INTERVAL_MS = 1000
|
|
|
+
|
|
|
+export const runScriptSchema = z.object({
|
|
|
+ code: z.string().min(1).max(200_000).describe('Python 脚本（完整代码）'),
|
|
|
+ tools: z.array(z.string()).max(100).default([]).describe('脚本中会调用的、会修改数据的工具（工具全名）。只读工具不用列出；没有列出的写工具在脚本中调用会失败'),
|
|
|
+ description: z.string().max(200).default('').describe('一句话说明脚本做什么（显示在进度中）'),
|
|
|
+})
|
|
|
+
|
|
|
+export const RUN_SCRIPT_DESCRIPTION = `执行一段 Python 脚本，在脚本中批量调用上面列出的 MCP 工具（工具全名，形如 服务名__工具名），汇总后只返回脚本打印的内容。
|
|
|
+适合需要调用很多次工具（几十到上千次）、或要对大量结果做筛选统计的任务：比逐次调用快得多，也不占对话的上下文。
|
|
|
+脚本中：
|
|
|
+ from agent_tools import call, call_many, call_raw, progress, ToolError
|
|
|
+ call(tool, args) 调用一次；返回结构化结果（dict）或文本（str）；工具报错时抛出 ToolError
|
|
|
+ call_many(tool, args_list) 批量调用同一个工具（平台并发执行），按顺序返回结果列表；return_errors=True 时出错的位置返回 ToolError 而不抛出
|
|
|
+ call_raw(tool, args) 返回 {"text", "structured", "is_error"}，不抛出工具错误
|
|
|
+ progress(done, total, message) 报告进度，会显示在对话中
|
|
|
+运行环境：与 python 命令相同的 Python 运行环境（pandas、numpy 等可用），当前目录为会话工作区，可以读写工作区中的文件（把大结果写成文件，只打印摘要）；不能联网，不能访问工作区以外的文件。
|
|
|
+规则：只读工具可以直接调用；会修改数据的工具必须在 tools 参数中列出，执行前统一审批一次；破坏性工具（每次都要确认的）不能在脚本中调用，请直接调用。每个脚本调用工具的次数、并发数、运行时间有上限，运行较久时会转为后台执行，完成后结果会发给你。`
|
|
|
+
|
|
|
+/** 脚本中可以调用的工具 */
|
|
|
+export interface ScriptTool {
|
|
|
+ readOnly: boolean
|
|
|
+ /** 调用工具：返回给脚本的文本、结构化结果、是否出错 */
|
|
|
+ invoke: (args: Record<string, unknown>, signal: AbortSignal) => Promise<{ text: string; structured: unknown; isError: boolean }>
|
|
|
+}
|
|
|
+
|
|
|
+/** 脚本中 import 的辅助模块 */
|
|
|
+const HELPER = `"""AgentServer 提供的脚本辅助模块：调用 MCP 工具、报告进度（docs/design/17-run-script.md）"""
|
|
|
+import itertools as _it
|
|
|
+import json as _json
|
|
|
+import os as _os
|
|
|
+import threading as _th
|
|
|
+
|
|
|
+__all__ = ["call", "call_many", "call_raw", "progress", "ToolError"]
|
|
|
+
|
|
|
+_out = _os.fdopen(3, "w", encoding="utf-8")
|
|
|
+_in = _os.fdopen(4, "r", encoding="utf-8")
|
|
|
+_lock = _th.Lock()
|
|
|
+_ids = _it.count(1)
|
|
|
+_BATCH = 50
|
|
|
+
|
|
|
+
|
|
|
+class ToolError(Exception):
|
|
|
+ """工具返回错误，或平台拒绝了这次调用"""
|
|
|
+
|
|
|
+ def __init__(self, tool, message):
|
|
|
+ super().__init__(f"{tool}: {message}")
|
|
|
+ self.tool = tool
|
|
|
+ self.message = message
|
|
|
+
|
|
|
+
|
|
|
+def _send(msg):
|
|
|
+ _out.write(_json.dumps(msg, ensure_ascii=False) + "\\n")
|
|
|
+ _out.flush()
|
|
|
+
|
|
|
+
|
|
|
+def _request(msgs):
|
|
|
+ with _lock:
|
|
|
+ for m in msgs:
|
|
|
+ _send(m)
|
|
|
+ pending = {m["id"] for m in msgs}
|
|
|
+ got = {}
|
|
|
+ while pending:
|
|
|
+ line = _in.readline()
|
|
|
+ if not line:
|
|
|
+ raise RuntimeError("与平台的连接已断开")
|
|
|
+ r = _json.loads(line)
|
|
|
+ got[r["id"]] = r
|
|
|
+ pending.discard(r["id"])
|
|
|
+ return got
|
|
|
+
|
|
|
+
|
|
|
+def _raw(r):
|
|
|
+ return {"text": r.get("text", ""), "structured": r.get("structured"), "is_error": bool(r.get("error") or r.get("isError"))}
|
|
|
+
|
|
|
+
|
|
|
+def _value(tool, r):
|
|
|
+ if r.get("error"):
|
|
|
+ raise ToolError(tool, r["error"])
|
|
|
+ if r.get("isError"):
|
|
|
+ raise ToolError(tool, r.get("text", ""))
|
|
|
+ return r["structured"] if r.get("structured") is not None else r.get("text", "")
|
|
|
+
|
|
|
+
|
|
|
+def call_raw(tool, args=None):
|
|
|
+ i = next(_ids)
|
|
|
+ r = _request([{"id": i, "type": "call", "tool": tool, "args": args or {}}])[i]
|
|
|
+ if r.get("error"):
|
|
|
+ raise ToolError(tool, r["error"])
|
|
|
+ return _raw(r)
|
|
|
+
|
|
|
+
|
|
|
+def call(tool, args=None):
|
|
|
+ i = next(_ids)
|
|
|
+ return _value(tool, _request([{"id": i, "type": "call", "tool": tool, "args": args or {}}])[i])
|
|
|
+
|
|
|
+
|
|
|
+def call_many(tool, args_list, return_errors=False):
|
|
|
+ out = []
|
|
|
+ items = list(args_list)
|
|
|
+ for start in range(0, len(items), _BATCH):
|
|
|
+ msgs = [{"id": next(_ids), "type": "call", "tool": tool, "args": a or {}} for a in items[start:start + _BATCH]]
|
|
|
+ got = _request(msgs)
|
|
|
+ for m in msgs:
|
|
|
+ try:
|
|
|
+ out.append(_value(tool, got[m["id"]]))
|
|
|
+ except ToolError as e:
|
|
|
+ if not return_errors:
|
|
|
+ raise
|
|
|
+ out.append(e)
|
|
|
+ return out
|
|
|
+
|
|
|
+
|
|
|
+def progress(done, total=None, message=None):
|
|
|
+ with _lock:
|
|
|
+ _send({"type": "progress", "done": done, "total": total, "message": None if message is None else str(message)[:200]})
|
|
|
+`
|
|
|
+
|
|
|
+/** 辅助模块所在目录（启动时写入数据目录，沙箱中只读） */
|
|
|
+export function scriptHelperDir(): string {
|
|
|
+ const dir = path.join(config.dataDir, 'data', 'script-runtime')
|
|
|
+ const file = path.join(dir, 'agent_tools.py')
|
|
|
+ if (!fs.existsSync(file) || fs.readFileSync(file, 'utf8') !== HELPER) {
|
|
|
+ fs.mkdirSync(dir, { recursive: true })
|
|
|
+ fs.writeFileSync(file, HELPER)
|
|
|
+ }
|
|
|
+ return dir
|
|
|
+}
|
|
|
+
|
|
|
+export interface ScriptRunOptions {
|
|
|
+ ctx: GatewayRunContext
|
|
|
+ code: string
|
|
|
+ title: string
|
|
|
+ tools: Map<string, ScriptTool>
|
|
|
+ /** 只读模式：工作区只读 */
|
|
|
+ readOnly: boolean
|
|
|
+ signal: AbortSignal
|
|
|
+ limits: { maxCalls: number; concurrency: number; maxMinutes: number }
|
|
|
+ pythonRuntime: string
|
|
|
+}
|
|
|
+
|
|
|
+export interface ScriptSummary {
|
|
|
+ id: string
|
|
|
+ exitCode: number | null
|
|
|
+ calls: number
|
|
|
+ failures: number
|
|
|
+ byTool: Record<string, number>
|
|
|
+ ms: number
|
|
|
+ logDir: string
|
|
|
+}
|
|
|
+
|
|
|
+/** 在沙箱中执行脚本，返回交给模型的结果 */
|
|
|
+export async function runScript(o: ScriptRunOptions): Promise<{ result: ToolResult; summary: ScriptSummary }> {
|
|
|
+ const id = `scr_${Date.now().toString(36)}${crypto.randomBytes(3).toString('hex')}`
|
|
|
+ const ws = fs.realpathSync(o.ctx.workspaceDir)
|
|
|
+ const logDir = path.join(ws, '.agent', 'scripts', id)
|
|
|
+ fs.mkdirSync(logDir, { recursive: true })
|
|
|
+ const scriptFile = path.join(logDir, 'script.py')
|
|
|
+ fs.writeFileSync(scriptFile, o.code)
|
|
|
+ const callLog = fs.createWriteStream(path.join(logDir, 'calls.jsonl'))
|
|
|
+ const rel = (p: string) => path.relative(ws, p)
|
|
|
+
|
|
|
+ const started = Date.now()
|
|
|
+ const stats = { calls: 0, failures: 0, byTool: {} as Record<string, number> }
|
|
|
+ const reported = { done: null as number | null, total: null as number | null, message: null as string | null }
|
|
|
+ let lastProgress = 0
|
|
|
+ const emit = (status: ScriptProgressPayload['status'], force = false) => {
|
|
|
+ if (!o.ctx.progress) return
|
|
|
+ const now = Date.now()
|
|
|
+ if (!force && now - lastProgress < PROGRESS_INTERVAL_MS) return
|
|
|
+ lastProgress = now
|
|
|
+ try {
|
|
|
+ o.ctx.progress({ scriptId: id, title: o.title, status, calls: stats.calls, failures: stats.failures, ...reported, elapsedMs: now - started })
|
|
|
+ } catch {
|
|
|
+ // 会话已删除等：进度只是展示
|
|
|
+ }
|
|
|
+ }
|
|
|
+ const progressTimer = setInterval(() => emit('running'), PROGRESS_INTERVAL_MS)
|
|
|
+
|
|
|
+ const runtime = fs.realpathSync(o.pythonRuntime)
|
|
|
+ const helper = scriptHelperDir()
|
|
|
+ const sandbox = sandboxArgs(ws, [runtime, helper])
|
|
|
+ // 只读模式：工作区只读（脚本仍可以打印结果）
|
|
|
+ const args = o.readOnly ? sandbox.map((a, i) => (a === '--bind' && sandbox[i + 1] === ws ? '--ro-bind' : a)) : sandbox
|
|
|
+ const child = spawn('bwrap', [
|
|
|
+ ...args,
|
|
|
+ '--setenv', 'PYTHONPATH', helper,
|
|
|
+ '--setenv', 'PYTHONUNBUFFERED', '1',
|
|
|
+ '--setenv', 'PYTHONDONTWRITEBYTECODE', '1',
|
|
|
+ '--setenv', 'MPLBACKEND', 'Agg',
|
|
|
+ '--setenv', 'MPLCONFIGDIR', '/tmp',
|
|
|
+ '--chdir', ws,
|
|
|
+ path.join(runtime, 'bin/python'), scriptFile,
|
|
|
+ ], { stdio: ['ignore', 'pipe', 'pipe', 'pipe', 'pipe'] })
|
|
|
+ const toScript = child.stdio[4] as NodeJS.WritableStream
|
|
|
+ const fromScript = child.stdio[3] as NodeJS.ReadableStream
|
|
|
+
|
|
|
+ let stdout = ''
|
|
|
+ let stdoutBytes = 0
|
|
|
+ const stdoutFile = fs.createWriteStream(path.join(logDir, 'output.txt'))
|
|
|
+ child.stdout!.on('data', (d: Buffer) => {
|
|
|
+ stdoutFile.write(d)
|
|
|
+ stdoutBytes += d.length
|
|
|
+ if (stdout.length < MAX_OUTPUT * 2) stdout += d.toString('utf8')
|
|
|
+ })
|
|
|
+ let stderr = ''
|
|
|
+ child.stderr!.on('data', (d: Buffer) => {
|
|
|
+ stderr = (stderr + d.toString('utf8')).slice(-MAX_STDERR * 4)
|
|
|
+ })
|
|
|
+
|
|
|
+ // ---- 脚本的工具调用：限制总次数与并发 ----
|
|
|
+ let active = 0
|
|
|
+ const queue: (() => void)[] = []
|
|
|
+ const acquire = () => new Promise<void>(resolve => {
|
|
|
+ if (active < o.limits.concurrency) {
|
|
|
+ active++
|
|
|
+ resolve()
|
|
|
+ } else queue.push(resolve)
|
|
|
+ })
|
|
|
+ const release = () => {
|
|
|
+ const next = queue.shift()
|
|
|
+ if (next) next()
|
|
|
+ else active--
|
|
|
+ }
|
|
|
+ const controller = new AbortController()
|
|
|
+ const reply = (msg: Record<string, unknown>) => {
|
|
|
+ if (!toScript.writable) return
|
|
|
+ try {
|
|
|
+ toScript.write(`${JSON.stringify(msg)}\n`)
|
|
|
+ } catch {
|
|
|
+ // 脚本已退出
|
|
|
+ }
|
|
|
+ }
|
|
|
+ const handleCall = async (m: { id: number; tool?: unknown; args?: unknown }) => {
|
|
|
+ const tool = typeof m.tool === 'string' ? m.tool : ''
|
|
|
+ const t = o.tools.get(tool)
|
|
|
+ if (!t) return reply({ id: m.id, error: `工具不可用：${tool}（只读工具可以直接调用；会修改数据的工具需要在 tools 参数中列出；破坏性工具不能在脚本中调用）` })
|
|
|
+ if (stats.calls >= o.limits.maxCalls) return reply({ id: m.id, error: `已达到每个脚本调用工具的上限（${o.limits.maxCalls} 次）` })
|
|
|
+ stats.calls++
|
|
|
+ stats.byTool[tool] = (stats.byTool[tool] ?? 0) + 1
|
|
|
+ await acquire()
|
|
|
+ const t0 = Date.now()
|
|
|
+ try {
|
|
|
+ const args = m.args && typeof m.args === 'object' && !Array.isArray(m.args) ? m.args as Record<string, unknown> : {}
|
|
|
+ const r = await t.invoke(args, controller.signal)
|
|
|
+ if (r.isError) stats.failures++
|
|
|
+ callLog.write(`${JSON.stringify({ at: new Date(t0).toISOString(), tool, args, ms: Date.now() - t0, isError: r.isError, text: r.text.slice(0, 500) })}\n`)
|
|
|
+ reply({ id: m.id, text: r.text.slice(0, MAX_CALL_TEXT), structured: r.structured ?? null, isError: r.isError })
|
|
|
+ } catch (err) {
|
|
|
+ stats.failures++
|
|
|
+ callLog.write(`${JSON.stringify({ at: new Date(t0).toISOString(), tool, ms: Date.now() - t0, error: (err as Error).message })}\n`)
|
|
|
+ reply({ id: m.id, error: (err as Error).message })
|
|
|
+ } finally {
|
|
|
+ release()
|
|
|
+ emit('running')
|
|
|
+ }
|
|
|
+ }
|
|
|
+ let buf = ''
|
|
|
+ fromScript.on('data', (d: Buffer) => {
|
|
|
+ buf += d.toString('utf8')
|
|
|
+ let nl: number
|
|
|
+ while ((nl = buf.indexOf('\n')) >= 0) {
|
|
|
+ const line = buf.slice(0, nl)
|
|
|
+ buf = buf.slice(nl + 1)
|
|
|
+ let m: { type?: string; id?: number; tool?: unknown; args?: unknown; done?: unknown; total?: unknown; message?: unknown }
|
|
|
+ try {
|
|
|
+ m = JSON.parse(line)
|
|
|
+ } catch {
|
|
|
+ continue
|
|
|
+ }
|
|
|
+ if (m.type === 'call' && typeof m.id === 'number') void handleCall(m as { id: number })
|
|
|
+ else if (m.type === 'progress') {
|
|
|
+ reported.done = typeof m.done === 'number' ? m.done : null
|
|
|
+ reported.total = typeof m.total === 'number' ? m.total : null
|
|
|
+ reported.message = typeof m.message === 'string' ? m.message : null
|
|
|
+ emit('running')
|
|
|
+ }
|
|
|
+ }
|
|
|
+ })
|
|
|
+
|
|
|
+ // ---- 结束：正常退出、超时、中断 ----
|
|
|
+ let stopReason: string | null = null
|
|
|
+ const kill = (reason: string) => {
|
|
|
+ stopReason ??= reason
|
|
|
+ controller.abort()
|
|
|
+ child.kill('SIGKILL')
|
|
|
+ }
|
|
|
+ const timer = setTimeout(() => kill(`超过 ${o.limits.maxMinutes} 分钟，已停止`), o.limits.maxMinutes * 60_000)
|
|
|
+ const onAbort = () => kill('已取消')
|
|
|
+ o.signal.addEventListener('abort', onAbort, { once: true })
|
|
|
+ emit('running', true)
|
|
|
+
|
|
|
+ const exitCode = await new Promise<number | null>(resolve => {
|
|
|
+ child.on('error', err => {
|
|
|
+ stderr += `\n${err.message}`
|
|
|
+ resolve(null)
|
|
|
+ })
|
|
|
+ child.on('close', code => resolve(code))
|
|
|
+ })
|
|
|
+ clearTimeout(timer)
|
|
|
+ clearInterval(progressTimer)
|
|
|
+ o.signal.removeEventListener('abort', onAbort)
|
|
|
+ controller.abort()
|
|
|
+ await Promise.all([new Promise(r => callLog.end(r)), new Promise(r => stdoutFile.end(r))])
|
|
|
+
|
|
|
+ const ms = Date.now() - started
|
|
|
+ const ok = exitCode === 0 && !stopReason
|
|
|
+ emit(stopReason === '已取消' ? 'cancelled' : ok ? 'completed' : 'failed', true)
|
|
|
+
|
|
|
+ const outFile = path.join(logDir, 'output.txt')
|
|
|
+ const shown = stdout.length > MAX_OUTPUT
|
|
|
+ ? `${stdout.slice(0, MAX_OUTPUT)}\n…（输出共 ${stdoutBytes} 字节，以上为前一部分；完整输出：${rel(outFile)}）`
|
|
|
+ : stdout.trim() || '（脚本没有输出）'
|
|
|
+ const head = `脚本${ok ? '执行完成' : stopReason ? `被中止：${stopReason}` : `出错退出（退出码 ${exitCode}）`}。调用工具 ${stats.calls} 次${stats.failures ? `（失败 ${stats.failures} 次）` : ''}，用时 ${(ms / 1000).toFixed(1)} 秒。调用明细：${rel(path.join(logDir, 'calls.jsonl'))}`
|
|
|
+ const parts = [head, '', '输出：', shown]
|
|
|
+ if (!ok && stderr.trim()) parts.push('', '错误信息：', stderr.trim().slice(-MAX_STDERR))
|
|
|
+ return {
|
|
|
+ result: { content: [{ type: 'text', text: parts.join('\n') }], ...(ok ? {} : { isError: true }) },
|
|
|
+ summary: { id, exitCode, calls: stats.calls, failures: stats.failures, byTool: stats.byTool, ms, logDir: rel(logDir) },
|
|
|
+ }
|
|
|
+}
|
Jerry Liang