287 lines
11 KiB
TypeScript
287 lines
11 KiB
TypeScript
// 工具：bash（前台 / 后台）、edit_file、question（阻塞 / 异步）。全部真实执行。
|
||
import { spawn, type ChildProcess } from 'node:child_process';
|
||
import { readFileSync, writeFileSync, mkdirSync } from 'node:fs';
|
||
import { dirname, resolve, relative, isAbsolute } from 'node:path';
|
||
import type { InboxItem, QuestionMode, ToolSpec, TraceBody } from '../shared/types.ts';
|
||
import type { LevelConfig } from './levels.ts';
|
||
|
||
export type ToolContext = {
|
||
workspace: string;
|
||
level: LevelConfig;
|
||
emit: (body: TraceBody) => void;
|
||
pushInbox: (item: Omit<InboxItem, 'id'>) => void;
|
||
onFileChanged: () => void;
|
||
};
|
||
|
||
type PendingQuestion = { mode: QuestionMode; resolve?: (answer: string) => void };
|
||
|
||
const FOREGROUND_TIMEOUT_MS = 180_000;
|
||
const RESULT_LIMIT = 6000;
|
||
const SUMMARY_LINES = 25;
|
||
|
||
export class ToolBox {
|
||
private ctx: ToolContext;
|
||
private children = new Set<ChildProcess>();
|
||
private jobSeq = 0;
|
||
private questionSeq = 0;
|
||
readonly runningJobs = new Map<string, { command: string; startedAt: number }>();
|
||
readonly pendingQuestions = new Map<string, PendingQuestion>();
|
||
|
||
constructor(ctx: ToolContext) {
|
||
this.ctx = ctx;
|
||
}
|
||
|
||
specs(): ToolSpec[] {
|
||
const { level } = this.ctx;
|
||
const bashParams: Record<string, unknown> = {
|
||
command: { type: 'string', description: '要执行的 shell 命令，在仓库根目录运行' },
|
||
};
|
||
if (level.backgroundBash) {
|
||
bashParams.run_in_background = {
|
||
type: 'boolean',
|
||
description: '为 true 时在后台运行，立即返回 job_id；命令退出后结果会以消息形式送达',
|
||
};
|
||
}
|
||
const specs: ToolSpec[] = [
|
||
{
|
||
type: 'function',
|
||
function: {
|
||
name: 'bash',
|
||
description: level.backgroundBash
|
||
? '执行 shell 命令。默认等命令退出后返回退出码和输出；run_in_background 为 true 时立即返回 job_id。'
|
||
: '执行 shell 命令，等命令退出后返回退出码和输出。',
|
||
parameters: { type: 'object', properties: bashParams, required: ['command'] },
|
||
},
|
||
},
|
||
{
|
||
type: 'function',
|
||
function: {
|
||
name: 'edit_file',
|
||
description:
|
||
'修改文件：把 old_string 替换成 new_string，old_string 必须在文件里恰好出现一次。old_string 为空时用 new_string 新建或覆盖整个文件。',
|
||
parameters: {
|
||
type: 'object',
|
||
properties: {
|
||
path: { type: 'string', description: '相对仓库根目录的路径' },
|
||
old_string: { type: 'string' },
|
||
new_string: { type: 'string' },
|
||
},
|
||
required: ['path', 'old_string', 'new_string'],
|
||
},
|
||
},
|
||
},
|
||
];
|
||
if (level.question) {
|
||
specs.push({
|
||
type: 'function',
|
||
function: {
|
||
name: 'question',
|
||
description:
|
||
level.question === 'blocking'
|
||
? '向用户提一个选择题，等用户回答后返回答案。'
|
||
: '向用户提一个选择题。立即返回 {"accepted":true,"question_id":...}，答案稍后以 <question_reply> 消息送达。',
|
||
parameters: {
|
||
type: 'object',
|
||
properties: {
|
||
question: { type: 'string' },
|
||
options: { type: 'array', items: { type: 'string' }, description: '2 到 4 个选项' },
|
||
},
|
||
required: ['question', 'options'],
|
||
},
|
||
},
|
||
});
|
||
}
|
||
return specs;
|
||
}
|
||
|
||
async run(callId: string, name: string, args: any): Promise<string> {
|
||
switch (name) {
|
||
case 'bash':
|
||
if (typeof args.command !== 'string') return '错误：缺少 command';
|
||
if (args.run_in_background && this.ctx.level.backgroundBash) {
|
||
return this.startBackground(callId, args.command);
|
||
}
|
||
return this.runForeground(callId, args.command);
|
||
case 'edit_file':
|
||
return this.editFile(args);
|
||
case 'question':
|
||
if (!this.ctx.level.question) break;
|
||
return this.ctx.level.question === 'blocking'
|
||
? this.askBlocking(callId, args)
|
||
: this.askAsync(callId, args);
|
||
}
|
||
return `错误：没有工具 ${name}`;
|
||
}
|
||
|
||
// 前台：工具调用一直挂着，直到子进程退出。
|
||
// 这里用的是 await，Node 事件循环没有被卡住，网页照样能收消息；但 Loop 停在这一步，模型看不到任何新东西。
|
||
private runForeground(callId: string, command: string): Promise<string> {
|
||
return new Promise((done) => {
|
||
const startedAt = Date.now();
|
||
const child = this.spawnShell(command);
|
||
let output = '';
|
||
const onData = (chunk: Buffer) => {
|
||
const text = chunk.toString();
|
||
output += text;
|
||
this.ctx.emit({ type: 'tool.output', callId, chunk: text });
|
||
};
|
||
child.stdout!.on('data', onData);
|
||
child.stderr!.on('data', onData);
|
||
const timer = setTimeout(() => child.kill('SIGKILL'), FOREGROUND_TIMEOUT_MS);
|
||
child.on('close', (code, signal) => {
|
||
clearTimeout(timer);
|
||
this.children.delete(child);
|
||
const ms = Date.now() - startedAt;
|
||
const status = signal ? `被信号 ${signal} 终止` : `exit_code: ${code}`;
|
||
done(`${status}（耗时 ${(ms / 1000).toFixed(1)} 秒）\n${tail(output, RESULT_LIMIT)}`);
|
||
});
|
||
});
|
||
}
|
||
|
||
// 后台：启动进程就返回 job_id。进程退出时，把退出码和日志摘要放进收件箱。
|
||
private startBackground(callId: string, command: string): string {
|
||
const jobId = `job_${++this.jobSeq}`;
|
||
const startedAt = Date.now();
|
||
const child = this.spawnShell(command);
|
||
this.runningJobs.set(jobId, { command, startedAt });
|
||
this.ctx.emit({ type: 'job.start', jobId, callId, command });
|
||
|
||
let output = '';
|
||
const onData = (chunk: Buffer) => {
|
||
const text = chunk.toString();
|
||
output += text;
|
||
this.ctx.emit({ type: 'job.output', jobId, chunk: text });
|
||
};
|
||
child.stdout!.on('data', onData);
|
||
child.stderr!.on('data', onData);
|
||
child.on('close', (code) => {
|
||
this.children.delete(child);
|
||
if (!this.runningJobs.delete(jobId)) return; // 已被 stop() 收掉
|
||
const durationMs = Date.now() - startedAt;
|
||
this.ctx.emit({ type: 'job.end', jobId, exitCode: code, durationMs });
|
||
this.ctx.pushInbox({
|
||
kind: 'job_done',
|
||
ref: jobId,
|
||
content:
|
||
`<background_job_done job_id="${jobId}" exit_code="${code}" duration_ms="${durationMs}">\n` +
|
||
`命令：${command}\n输出最后 ${SUMMARY_LINES} 行：\n${lastLines(output, SUMMARY_LINES)}\n` +
|
||
`</background_job_done>`,
|
||
});
|
||
});
|
||
return JSON.stringify({
|
||
job_id: jobId,
|
||
status: 'running',
|
||
note: '命令已在后台启动。退出后你会收到 <background_job_done> 消息。',
|
||
});
|
||
}
|
||
|
||
// AiWork 式：工具调用挂起，直到用户回答。答案就是工具结果。
|
||
private askBlocking(callId: string, args: any): Promise<string> {
|
||
const questionId = `q_${++this.questionSeq}`;
|
||
this.ctx.emit({
|
||
type: 'question.ask',
|
||
questionId,
|
||
callId,
|
||
mode: 'blocking',
|
||
question: String(args.question ?? ''),
|
||
options: Array.isArray(args.options) ? args.options.map(String) : [],
|
||
});
|
||
return new Promise((resolveAnswer) => {
|
||
this.pendingQuestions.set(questionId, {
|
||
mode: 'blocking',
|
||
resolve: (answer) => resolveAnswer(`用户回答：${answer}`),
|
||
});
|
||
});
|
||
}
|
||
|
||
// Codex 式：先回执 accepted，答案稍后带着 question_id 作为用户消息进收件箱。
|
||
private askAsync(callId: string, args: any): string {
|
||
const questionId = `q_${++this.questionSeq}`;
|
||
this.pendingQuestions.set(questionId, { mode: 'async' });
|
||
this.ctx.emit({
|
||
type: 'question.ask',
|
||
questionId,
|
||
callId,
|
||
mode: 'async',
|
||
question: String(args.question ?? ''),
|
||
options: Array.isArray(args.options) ? args.options.map(String) : [],
|
||
});
|
||
return JSON.stringify({ accepted: true, question_id: questionId });
|
||
}
|
||
|
||
answer(questionId: string, answer: string): void {
|
||
const pending = this.pendingQuestions.get(questionId);
|
||
if (!pending) throw new Error(`问题 ${questionId} 不存在或已回答`);
|
||
this.pendingQuestions.delete(questionId);
|
||
this.ctx.emit({ type: 'question.answer', questionId, answer });
|
||
if (pending.mode === 'blocking') {
|
||
pending.resolve!(answer);
|
||
} else {
|
||
this.ctx.pushInbox({
|
||
kind: 'question_reply',
|
||
ref: questionId,
|
||
content: `<question_reply question_id="${questionId}">${answer}</question_reply>`,
|
||
});
|
||
}
|
||
}
|
||
|
||
private editFile(args: any): string {
|
||
const path = String(args.path ?? '');
|
||
const full = resolve(this.ctx.workspace, path);
|
||
const rel = relative(this.ctx.workspace, full);
|
||
if (!path || rel.startsWith('..') || isAbsolute(rel)) return `错误：路径 ${path} 不在仓库内`;
|
||
const oldString = String(args.old_string ?? '');
|
||
const newString = String(args.new_string ?? '');
|
||
if (oldString === '') {
|
||
mkdirSync(dirname(full), { recursive: true });
|
||
writeFileSync(full, newString);
|
||
this.ctx.onFileChanged();
|
||
return `已写入 ${rel}`;
|
||
}
|
||
let text: string;
|
||
try {
|
||
text = readFileSync(full, 'utf8');
|
||
} catch {
|
||
return `错误：文件 ${rel} 不存在`;
|
||
}
|
||
const count = text.split(oldString).length - 1;
|
||
if (count !== 1) return `错误：old_string 在 ${rel} 里出现了 ${count} 次，需要恰好 1 次`;
|
||
writeFileSync(full, text.replace(oldString, () => newString));
|
||
this.ctx.onFileChanged();
|
||
return `已修改 ${rel}`;
|
||
}
|
||
|
||
private spawnShell(command: string): ChildProcess {
|
||
const env: NodeJS.ProcessEnv = { ...process.env, CI: '1', NO_COLOR: '1', FORCE_COLOR: '0' };
|
||
delete env.EASTMONEY_API_KEY;
|
||
const child = spawn('bash', ['-c', command], {
|
||
cwd: this.ctx.workspace,
|
||
env,
|
||
stdio: ['ignore', 'pipe', 'pipe'],
|
||
detached: true, // 独立进程组，stop() 时连同 npm 拉起的子进程一起结束
|
||
});
|
||
this.children.add(child);
|
||
return child;
|
||
}
|
||
|
||
stop(): void {
|
||
this.runningJobs.clear();
|
||
for (const child of this.children) {
|
||
try {
|
||
process.kill(-child.pid!, 'SIGKILL');
|
||
} catch {}
|
||
}
|
||
this.children.clear();
|
||
for (const pending of this.pendingQuestions.values()) pending.resolve?.('（会话已停止）');
|
||
this.pendingQuestions.clear();
|
||
}
|
||
}
|
||
|
||
function tail(text: string, limit: number): string {
|
||
return text.length <= limit ? text : `…（前面省略 ${text.length - limit} 字）\n${text.slice(-limit)}`;
|
||
}
|
||
|
||
function lastLines(text: string, count: number): string {
|
||
return text.trimEnd().split('\n').slice(-count).join('\n');
|
||
}
|