170 lines
12 KiB
TypeScript
170 lines
12 KiB
TypeScript
import assert from 'node:assert/strict'
|
|
import { mkdtemp, readFile, writeFile } from 'node:fs/promises'
|
|
import { tmpdir } from 'node:os'
|
|
import { join } from 'node:path'
|
|
import { SubagentScheduler, type SubagentWorker } from '../src/main/subagents/scheduler'
|
|
import { executeTurnTool, type ConversationToolContext } from '../src/main/agent/conversation-tools'
|
|
import { toolSpecs } from '../src/main/agent/tools'
|
|
import type { CapabilityToolSnapshot } from '../src/main/capabilities/runtime'
|
|
import type { ChatEvent, ChatMessage, ToolCallRecord } from '../src/shared/types'
|
|
import { applyAssistantEvent } from '../src/shared/message-parts'
|
|
import { getRevertChanges } from '../src/shared/revert'
|
|
import { assistantWireHistory } from '../src/main/agent/history'
|
|
|
|
const root = await mkdtemp(join(tmpdir(), 'ycode-agents-'))
|
|
const checks: string[] = []
|
|
function check(label: string, condition: unknown): void { assert.ok(condition, label); checks.push(label); console.log('PASS ' + label) }
|
|
const pause = (ms = 10): Promise<void> => new Promise(done => setTimeout(done, ms))
|
|
const base: CapabilityToolSnapshot = { specs: [], execute: async () => undefined }
|
|
function setup(enabled = true, hold = true, timeoutMs = 2000, resultText?: string) {
|
|
const id = Math.random().toString(36), controller = new AbortController(), events: ChatEvent[] = []
|
|
const policies = new Map<string, CapabilityToolSnapshot>(), prompts: string[] = [], release = new Map<string, () => void>()
|
|
let current = 0, peak = 0, disposed = 0
|
|
const scheduler = new SubagentScheduler({ runId: id, workspaceRoot: root, modelName: 'selected-model', enabled,
|
|
signal: controller.signal, emit: event => events.push(event), capabilities: base, timeoutMs,
|
|
worker: async (task, policy): Promise<SubagentWorker> => {
|
|
policies.set(task.id, policy)
|
|
return { async run(prompt, signal, emit) {
|
|
current++; peak = Math.max(peak, current); prompts.push(prompt)
|
|
emit({ runId: 'child', type: 'usage', usage: { inputTokens: 20, outputTokens: 5, totalTokens: 25 } })
|
|
try {
|
|
if (hold) await new Promise<void>((done, reject) => {
|
|
const onAbort = (): void => { signal.removeEventListener('abort', onAbort); reject(Object.assign(new Error('cancelled'), { name: 'AbortError' })) }
|
|
signal.addEventListener('abort', onAbort, { once: true })
|
|
release.set(task.id, () => { signal.removeEventListener('abort', onAbort); done() })
|
|
if (signal.aborted) onAbort()
|
|
})
|
|
emit({ runId: 'child', type: 'text', delta: 'RESULT:' + task.name })
|
|
return { content: resultText ?? 'RESULT:' + task.name, tools: [] }
|
|
} finally { current-- }
|
|
}, async dispose() { disposed++ } }
|
|
}
|
|
})
|
|
const ctx: ConversationToolContext = { runId: id, toolCallId: 'test-call', sessionId: 'parent-' + id, workspaceRoot: root,
|
|
permissionMode: 'full-access', agentMode: 'craft', signal: controller.signal, emit: event => events.push(event) }
|
|
const call = (name: string, args: Record<string, unknown> = {}) => executeTurnTool(name, args, ctx, scheduler.tools)
|
|
const spawn = async (name: string, extra = {}) => {
|
|
const result = await call('spawn_agent', { name, task: '独立调查并给出验收结果', ...extra })
|
|
assert.equal(result.ok, true, result.content); return JSON.parse(result.content) as { id: string; status: string }
|
|
}
|
|
return { scheduler, controller, ctx, call, spawn, events, policies, prompts, release,
|
|
get peak() { return peak }, get current() { return current }, get disposed() { return disposed } }
|
|
}
|
|
|
|
{
|
|
const state = setup(false)
|
|
check('disabled switch hides spawn and rejects a direct stale-name invocation', !state.scheduler.tools.specs.some(tool => tool.function.name === 'spawn_agent') && !(await state.call('spawn_agent', { name: 'x', task: 'x' })).ok)
|
|
state.scheduler.setEnabled(true)
|
|
check('enabling immediately exposes real scheduling tools', state.scheduler.tools.specs.some(tool => tool.function.name === 'spawn_agent'))
|
|
check('plan mode cannot bypass the scheduling mode gate', !(await executeTurnTool('spawn_agent', { name: 'x', task: 'x' }, { ...state.ctx, agentMode: 'plan' }, state.scheduler.tools)).ok)
|
|
const spawned = await Promise.all([1, 2, 3, 4, 5, 6].map(i => state.spawn('任务' + i)))
|
|
await pause()
|
|
check('six tasks use five workers with the sixth queued', state.peak === 5 && JSON.parse((await state.call('list_agents')).content).filter((view: { status: string }) => view.status === 'queued').length === 1)
|
|
const stranger = setup()
|
|
check('cross-turn agent identifiers cannot be read or cancelled', !(await stranger.call('cancel_agent', { id: spawned[0].id })).ok)
|
|
await stranger.spawn('跨会话排队')
|
|
check('the concurrency budget also applies across root conversations', stranger.current === 0)
|
|
await stranger.scheduler.close()
|
|
state.scheduler.setEnabled(false)
|
|
await state.scheduler.beforeFinish()
|
|
check('disabling cancels running and queued children without aborting the parent', !state.controller.signal.aborted && JSON.parse((await state.call('list_agents')).content).every((view: { status: string }) => view.status === 'cancelled'))
|
|
check('disabled switch also denies follow-up messages', !(await state.call('send_agent_message', { id: spawned[0].id, message: '继续' })).ok)
|
|
await state.scheduler.close()
|
|
check('cleanup disposes the five created child runtimes', state.disposed === 5)
|
|
}
|
|
{
|
|
const state = setup(true, false, 2000, '\u0000'.repeat(60000))
|
|
await state.spawn('长结果')
|
|
await pause()
|
|
const result = await state.call('wait_agent', { timeoutMs: 1000 })
|
|
check('large escaped results remain bounded complete JSON', result.content.length <= 30000 && JSON.parse(result.content)[0].truncated)
|
|
await state.scheduler.close()
|
|
}
|
|
{
|
|
const state = setup()
|
|
const writer = await state.spawn('写文件', { access: 'write', files: ['owned.txt'] })
|
|
check('write declarations reject directories, escapes and glob patterns', !(await state.call('spawn_agent', { name: 'bad', task: 'x', access: 'write', files: ['.'] })).ok
|
|
&& !(await state.call('spawn_agent', { name: 'bad', task: 'x', access: 'write', files: ['../escape.txt'] })).ok
|
|
&& !(await state.call('spawn_agent', { name: 'bad', task: 'x', access: 'write', files: ['*.ts'] })).ok)
|
|
await pause()
|
|
const policy = state.policies.get(writer.id)!
|
|
const workerCtx = { ...state.ctx, runId: 'worker', sessionId: 'child', toolCallId: 'child-write' }
|
|
check('child schemas hide commands, browsers and nested scheduling', policy.filterSpecs!(toolSpecs()).every(tool => ['list_dir', 'read_file', 'write_file', 'edit_file'].includes(tool.function.name)))
|
|
check('unlisted shell and nested spawn calls are rejected by the host', !(await executeTurnTool('execute_command', { command: 'echo forbidden' }, workerCtx, policy)).ok
|
|
&& !(await executeTurnTool('spawn_agent', { name: 'nested', task: 'nested' }, workerCtx, policy)).ok)
|
|
check('writer cannot modify another file', !(await executeTurnTool('write_file', { path: 'outside.txt', content: 'no' }, workerCtx, policy)).ok)
|
|
check('same-file sibling assignment is denied', !(await state.call('spawn_agent', { name: '冲突', task: 'x', access: 'write', files: ['owned.txt'] })).ok)
|
|
check('parent write and shell cannot collide with a child lease', !(await state.call('write_file', { path: 'owned.txt', content: 'parent' })).ok
|
|
&& !(await state.call('execute_command', { command: 'echo no' })).ok)
|
|
const childWrite = await executeTurnTool('write_file', { path: 'owned.txt', content: 'child' }, workerCtx, policy)
|
|
check('declared full-access child file write really reaches disk', childWrite.ok && await readFile(join(root, 'owned.txt'), 'utf8') === 'child')
|
|
state.release.get(writer.id)!()
|
|
await state.call('wait_agent', { ids: [writer.id], timeoutMs: 1000 })
|
|
const parentWrite = await state.call('write_file', { path: 'owned.txt', content: 'parent-after' })
|
|
check('file lease is released after child completion', parentWrite.ok && await readFile(join(root, 'owned.txt'), 'utf8') === 'parent-after')
|
|
const childTool: ToolCallRecord = { id: 'child-write', name: 'write_file', status: 'ok', args: {}, argsRaw: '{}', result: childWrite.data, metadata: childWrite.metadata }
|
|
const parentTool: ToolCallRecord = { ...childTool, id: 'parent-write', result: parentWrite.data, metadata: parentWrite.metadata }
|
|
const view = state.events.findLast(event => event.type === 'part' && event.part.type === 'subagent')
|
|
assert.ok(view?.type === 'part' && view.part.type === 'subagent')
|
|
const messages: ChatMessage[] = [{ id: 'user', role: 'user', content: 'change', createdAt: 0 },
|
|
{ id: 'assistant', role: 'assistant', content: '', createdAt: 1, tools: [parentTool],
|
|
parts: [{ type: 'subagent', agent: { ...view.part.agent, tools: [childTool] } }], status: 'done' }]
|
|
const changes = getRevertChanges(messages, 'user')
|
|
check('rewind includes child writes before later parent writes', changes.length === 1 && changes[0].before === null && changes[0].after === 'parent-after')
|
|
await state.scheduler.close()
|
|
}
|
|
{
|
|
const state = setup()
|
|
const reader = await state.spawn('只读')
|
|
await pause()
|
|
const policy = state.policies.get(reader.id)!
|
|
check('read-only child cannot write even with inherited full-access permission', !(await executeTurnTool('write_file', { path: 'read-only.txt', content: 'no' }, { ...state.ctx, runId: 'reader' }, policy)).ok)
|
|
const after = state.scheduler.beforeFinish()
|
|
state.release.get(reader.id)!()
|
|
check('parent finalization waits and supplies unread child results for synthesis', (await after)?.includes('RESULT:只读'))
|
|
check('collected results do not create an infinite finalization loop', await state.scheduler.beforeFinish() === undefined)
|
|
const part = state.events.findLast(event => event.type === 'part' && event.part.type === 'subagent')
|
|
assert.ok(part?.type === 'part')
|
|
check('HTTP cold history preserves child results as explicitly labeled task data', JSON.stringify(assistantWireHistory({
|
|
id: 'history', role: 'assistant', content: '', createdAt: 0, parts: [part.part]
|
|
})).includes('RESULT:只读'))
|
|
const follow = await state.call('send_agent_message', { id: reader.id, message: '补充核对' })
|
|
await pause()
|
|
state.release.get(reader.id)!()
|
|
await state.call('wait_agent', { ids: [reader.id], timeoutMs: 1000 })
|
|
check('completed child keeps its runtime for follow-up instructions', follow.ok && state.prompts.length === 2 && state.prompts[1] === '补充核对')
|
|
await state.scheduler.close()
|
|
}
|
|
{
|
|
const state = setup()
|
|
await state.spawn('取消测试')
|
|
await pause()
|
|
state.controller.abort()
|
|
await state.scheduler.close()
|
|
check('parent cancellation ends all children and releases their runtimes', state.current === 0 && state.disposed === 1)
|
|
}
|
|
{
|
|
const state = setup(true, true, 50)
|
|
await state.spawn('超时')
|
|
await pause(100)
|
|
const result = JSON.parse((await state.call('list_agents')).content)
|
|
check('child execution timeout reaches a terminal state', result[0].status === 'cancelled')
|
|
await state.scheduler.close()
|
|
}
|
|
{
|
|
const state = setup()
|
|
for (let i = 0; i < 16; i++) await state.spawn('有界任务' + i)
|
|
check('task creation is bounded per root turn', !(await state.call('spawn_agent', { name: 'overflow', task: 'overflow' })).ok)
|
|
const first = state.events.find(event => event.type === 'part' && event.part.type === 'subagent')
|
|
assert.ok(first?.type === 'part' && first.part.type === 'subagent')
|
|
let message: ChatMessage = { id: 'recover', role: 'assistant', content: '', createdAt: 0 }
|
|
message = applyAssistantEvent(message, first)
|
|
message = applyAssistantEvent(message, { ...first, part: { type: 'subagent', agent: { ...first.part.agent, status: 'running' } } })
|
|
check('streaming agent snapshots upsert in place', message.parts?.length === 1)
|
|
message = applyAssistantEvent(message, { runId: state.ctx.runId, type: 'cancelled' })
|
|
check('cold interrupted history never leaves a child shown as running', message.parts?.[0].type === 'subagent' && message.parts[0].agent.status === 'cancelled')
|
|
await state.scheduler.close()
|
|
}
|
|
await writeFile(join(root, 'report.json'), JSON.stringify({ result: 'PASS', checks }, null, 2))
|
|
console.log('SUBAGENT_REGRESSION_PASS ' + JSON.stringify({ checks: checks.length, directory: root }))
|