feat: add Pi subagent scheduler and child runtime

This commit is contained in:
2026-08-23 12:40:22 +08:00
parent b806c78139
commit f1c7cd8ad4
21 changed files with 1986 additions and 53 deletions

View File

@@ -58,6 +58,52 @@ function apply(
}
describe('Pi event projector', () => {
it('projects only subagent.v1 details and never exposes unknown raw details', async () => {
const projector = new PiEventProjector({ createId: () => 'unused' });
let snapshot = emptySnapshot();
snapshot.nodes.push({
kind: 'tool', id: 'tool-subagent', toolCallId: 'call-subagent', toolName: 'subagent',
title: 'subagent', inputText: '{}', status: 'running', output: [],
});
snapshot = apply(snapshot, await projector.project(snapshot, {
type: 'tool_execution_update',
toolCallId: 'call-subagent',
partialResult: {
content: [],
details: {
schema: 'subagent.v1', dispatchId: 'dispatch-a', mode: 'parallel',
tasks: [{
taskId: 'task-a', agentId: 'agent-a', toolProfile: 'read-only',
status: 'complete', summary: 'done', usage: { inputTokens: 3, outputTokens: 5 },
}],
},
},
}));
expect(snapshot.nodes).toContainEqual(expect.objectContaining({
kind: 'subagent', id: 'subagent:dispatch-a', runId: 'run-a',
details: expect.objectContaining({ schema: 'subagent.v1' }),
}));
expect(snapshot.nodes).toContainEqual(expect.objectContaining({
kind: 'tool', id: 'tool-subagent',
details: expect.objectContaining({ dispatchId: 'dispatch-a' }),
}));
snapshot.nodes.push({
kind: 'tool', id: 'tool-unknown', toolCallId: 'call-unknown', toolName: 'subagent',
title: 'subagent', inputText: '{}', status: 'running', output: [],
});
snapshot = apply(snapshot, await projector.project(snapshot, {
type: 'tool_execution_end',
toolCallId: 'call-unknown', isError: false,
result: { details: { schema: 'subagent.v2', raw: 'RAW_SECRET_DETAILS' } },
}));
expect(JSON.stringify(snapshot)).not.toContain('RAW_SECRET_DETAILS');
expect(snapshot.nodes).toContainEqual(expect.objectContaining({
kind: 'tool', id: 'tool-unknown',
output: [expect.objectContaining({ text: 'Subagent details are unavailable for this version.' })],
}));
});
it('keeps one assistant UI identity while content-index deltas become an authoritative message', async () => {
const projector = new PiEventProjector({
createId: () => 'assistant-ui-a',

View File

@@ -6,8 +6,14 @@ import path from 'node:path';
import { pathToFileURL } from 'node:url';
import { afterEach, describe, expect, it } from 'vitest';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
import { PiSubagentScheduler } from '../../electron/coding-runtime/pi/subagent';
import { PiProcessBudget } from '../../electron/coding-runtime/pi/worker-pool';
type ExtensionHandler = (...arguments_: unknown[]) => Promise<unknown> | unknown;
type ExtensionTool = {
name: string;
execute?: (...arguments_: unknown[]) => Promise<unknown>;
};
const roots: string[] = [];
const hosts: PiManagedExtensionHost[] = [];
@@ -22,6 +28,15 @@ describe('Makelore Pi extension bundle', () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-bundle-'));
roots.push(root);
const host = new PiManagedExtensionHost();
const scheduler = new PiSubagentScheduler({
processBudget: new PiProcessBudget(8),
openChild: async (input) => ({
id: input.taskId,
async run() { return { summary: `done ${input.agentId}` }; },
async stop() {},
}),
});
host.configureSubagents({ scheduler });
hosts.push(host);
const extensionWorker = await host.registerWorker({
conversationId: 'conversation-a1', generation: 1, projectId: 'project-a', extensionsDir: root,
@@ -38,22 +53,44 @@ describe('Makelore Pi extension bundle', () => {
bridge: process.env.MAKELORE_PI_BRIDGE_URL,
token: process.env.MAKELORE_PI_WORKER_TOKEN,
context: process.env.MAKELORE_PI_CONTEXT_FILE,
role: process.env.MAKELORE_PI_WORKER_ROLE,
};
Object.assign(process.env, extensionWorker.env);
try {
const module = await import(/* @vite-ignore */ pathToFileURL(extensionWorker.extensionPath).href) as {
default(factory: {
registerTool(tool: { name: string }): void;
registerTool(tool: ExtensionTool): void;
on(event: string, handler: ExtensionHandler): void;
}): void;
};
const handlers = new Map<string, ExtensionHandler>();
const tools: string[] = [];
const tools = new Map<string, ExtensionTool>();
module.default({
registerTool: (tool) => tools.push(tool.name),
registerTool: (tool) => tools.set(tool.name, tool),
on: (event, handler) => handlers.set(event, handler),
});
expect(tools).toEqual(['ask_user']);
expect([...tools.keys()]).toEqual(['ask_user', 'subagent']);
const updates: unknown[] = [];
const subagentResult = await tools.get('subagent')?.execute?.(
'subagent-1',
{
mode: 'single',
tasks: [{ agentId: 'agent-a', task: 'Inspect', toolProfile: 'read-only' }],
},
new AbortController().signal,
(update: unknown) => updates.push(update),
);
expect(subagentResult).toMatchObject({
details: {
schema: 'subagent.v1', mode: 'single',
tasks: [{ agentId: 'agent-a', status: 'complete', summary: 'done agent-a' }],
},
});
expect(updates.length).toBeGreaterThan(0);
expect(updates.every((update) => (
(update as { details?: { schema?: string } }).details?.schema === 'subagent.v1'
))).toBe(true);
const statuses: Array<string | undefined> = [];
const context = {
@@ -84,12 +121,55 @@ describe('Makelore Pi extension bundle', () => {
await handlers.get('tool_result')?.({ toolCallId: 'write-1' });
expect((await waiting).status).toBe(200);
} finally {
await scheduler.close();
if (previousEnvironment.bridge === undefined) delete process.env.MAKELORE_PI_BRIDGE_URL;
else process.env.MAKELORE_PI_BRIDGE_URL = previousEnvironment.bridge;
if (previousEnvironment.token === undefined) delete process.env.MAKELORE_PI_WORKER_TOKEN;
else process.env.MAKELORE_PI_WORKER_TOKEN = previousEnvironment.token;
if (previousEnvironment.context === undefined) delete process.env.MAKELORE_PI_CONTEXT_FILE;
else process.env.MAKELORE_PI_CONTEXT_FILE = previousEnvironment.context;
if (previousEnvironment.role === undefined) delete process.env.MAKELORE_PI_WORKER_ROLE;
else process.env.MAKELORE_PI_WORKER_ROLE = previousEnvironment.role;
}
});
it('does not expose parent-only tools from a child process', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-child-bundle-'));
roots.push(root);
const host = new PiManagedExtensionHost();
hosts.push(host);
const child = await host.registerWorker({
conversationId: 'conversation-child', generation: 1, projectId: 'project-a',
extensionsDir: root, role: 'child', runId: 'run-parent',
});
const previous = {
bridge: process.env.MAKELORE_PI_BRIDGE_URL,
token: process.env.MAKELORE_PI_WORKER_TOKEN,
context: process.env.MAKELORE_PI_CONTEXT_FILE,
role: process.env.MAKELORE_PI_WORKER_ROLE,
};
Object.assign(process.env, child.env);
try {
const module = await import(
/* @vite-ignore */ `${pathToFileURL(child.extensionPath).href}?child=${Date.now()}`
) as {
default(factory: {
registerTool(tool: ExtensionTool): void;
on(event: string, handler: ExtensionHandler): void;
}): void;
};
const tools: string[] = [];
module.default({ registerTool: (tool) => tools.push(tool.name), on: () => undefined });
expect(tools).toEqual([]);
} finally {
for (const [key, value] of Object.entries(previous)) {
const environmentKey = key === 'bridge' ? 'MAKELORE_PI_BRIDGE_URL'
: key === 'token' ? 'MAKELORE_PI_WORKER_TOKEN'
: key === 'context' ? 'MAKELORE_PI_CONTEXT_FILE'
: 'MAKELORE_PI_WORKER_ROLE';
if (value === undefined) delete process.env[environmentKey];
else process.env[environmentKey] = value;
}
}
});
});

View File

@@ -6,6 +6,11 @@ import path from 'node:path';
import { afterEach, describe, expect, it } from 'vitest';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
import { PiProjectWriteLeaseCoordinator } from '../../electron/coding-runtime/pi/write-lease';
import {
PiSubagentChildError,
PiSubagentScheduler,
} from '../../electron/coding-runtime/pi/subagent';
import { PiProcessBudget } from '../../electron/coding-runtime/pi/worker-pool';
const roots: string[] = [];
const hosts: PiManagedExtensionHost[] = [];
@@ -81,7 +86,7 @@ describe('managed Pi extension bridge', () => {
'utf8',
)) as Record<string, unknown>;
expect(context).toEqual({
conversationId: 'conversation-a', workerGeneration: 2, runId: 'run-a',
conversationId: 'conversation-a', workerGeneration: 2, role: 'parent', runId: 'run-a',
});
await first.dispose();
const response = await post(replacement, {
@@ -91,6 +96,117 @@ describe('managed Pi extension bridge', () => {
expect(response.status).toBe(200);
});
it('streams stable subagent details and rejects recursive child dispatch', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-subagent-'));
roots.push(root);
const scheduler = new PiSubagentScheduler({
processBudget: new PiProcessBudget(8),
openChild: async (input) => ({
id: input.taskId,
async run() { return { summary: `done ${input.agentId}` }; },
async stop() {},
}),
});
const host = new PiManagedExtensionHost();
const tracked: Array<{ kind: string; id: string }> = [];
let untracked = 0;
host.configureSubagents({
scheduler,
trackGenerationResource: (input) => {
tracked.push({ kind: input.kind, id: input.id });
return () => { untracked += 1; };
},
});
hosts.push(host);
const parent = await host.registerWorker({
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
});
const child = await host.registerWorker({
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
role: 'child', runId: 'run-a',
});
await host.bindRun('conversation-a', 1, 'run-a');
const request = {
action: 'subagent.dispatch', conversationId: 'conversation-a', workerGeneration: 1,
runId: 'run-a', resourceId: 'subagent-tool',
request: {
mode: 'single',
tasks: [{ agentId: 'agent-a', task: 'Inspect', toolProfile: 'read-only' }],
},
};
const response = await post(parent, request);
expect(response.status).toBe(200);
const lines = (await response.text()).trim().split('\n').map((line) => JSON.parse(line));
expect(lines.at(-1)).toMatchObject({
done: true,
details: {
schema: 'subagent.v1', mode: 'single',
tasks: [{ agentId: 'agent-a', status: 'complete', summary: 'done agent-a' }],
},
});
expect(tracked).toEqual([{ kind: 'child', id: 'subagent-tool' }]);
expect(untracked).toBe(1);
expect((await post(child, request)).status).toBe(403);
await scheduler.close();
});
it('propagates generation cancellation without an orphan child or leaked budget', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-child-cancel-'));
roots.push(root);
const processBudget = new PiProcessBudget(8);
let stopped = 0;
const scheduler = new PiSubagentScheduler({
processBudget,
openChild: async (input) => ({
id: input.taskId,
async run(_prompt, signal) {
await new Promise<void>((_resolve, reject) => {
const abort = () => reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
if (signal.aborted) abort();
else signal.addEventListener('abort', abort, { once: true });
});
return { summary: 'unreachable' };
},
async stop() { stopped += 1; },
}),
});
let cancelGeneration: (() => void) | undefined;
let untracked = 0;
const host = new PiManagedExtensionHost();
host.configureSubagents({
scheduler,
trackGenerationResource: (input) => {
cancelGeneration = input.cancel;
return () => { untracked += 1; };
},
});
hosts.push(host);
const parent = await host.registerWorker({
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
});
await host.bindRun('conversation-a', 1, 'run-a');
const flight = post(parent, {
action: 'subagent.dispatch', conversationId: 'conversation-a', workerGeneration: 1,
runId: 'run-a', resourceId: 'subagent-tool',
request: {
mode: 'single',
tasks: [{ agentId: 'agent-a', task: 'Wait', toolProfile: 'read-only' }],
},
});
await expect.poll(() => processBudget.activeCount).toBe(1);
cancelGeneration?.();
const response = await flight;
const lines = (await response.text()).trim().split('\n').map((line) => JSON.parse(line));
expect(lines.at(-1)).toMatchObject({
done: true,
details: { tasks: [{ status: 'aborted', errorCode: 'SUBAGENT_ABORTED' }] },
});
expect(stopped).toBe(1);
expect(processBudget.activeCount).toBe(0);
expect(untracked).toBe(1);
await scheduler.close();
});
it('validates worker identity and enforces project-scoped leases over loopback HTTP', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-'));
roots.push(root);
@@ -152,4 +268,38 @@ describe('managed Pi extension bridge', () => {
});
expect(currentWorker.status).toBe(200);
});
it('joins a coding child to the same project write lease as its parent', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-child-lease-'));
roots.push(root);
const leases = new PiProjectWriteLeaseCoordinator();
const host = new PiManagedExtensionHost(leases);
hosts.push(host);
const parent = await host.registerWorker({
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
});
const child = await host.registerWorker({
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
role: 'child', runId: 'run-a',
});
await host.bindRun('conversation-a', 1, 'run-a');
const identity = { conversationId: 'conversation-a', workerGeneration: 1, runId: 'run-a' };
const held = await post(parent, {
...identity, action: 'lease.acquire', resourceId: 'parent-write',
});
const parentLease = await held.json() as { leaseId: string };
let childSettled = false;
const waiting = post(child, {
...identity, action: 'lease.acquire', resourceId: 'child-write',
}).then((response) => {
childSettled = true;
return response;
});
await expect.poll(() => leases.waitingCount('project-a')).toBe(1);
expect(childSettled).toBe(false);
expect((await post(parent, {
...identity, action: 'lease.release', resourceId: 'parent-write', leaseId: parentLease.leaseId,
})).status).toBe(200);
expect((await waiting).status).toBe(200);
});
});

View File

@@ -161,7 +161,7 @@ describe('managed Pi worker opener', () => {
expect(argv).toContain('grilling');
expect(argv).toContain('--session-id');
expect(argv).toContain('--extension');
expect(argv).toContain('makelore-runtime-v1.mjs');
expect(argv).toContain('makelore-runtime-v2.mjs');
expect(options.additionalArgs?.filter((argument) => argument === '--extension')).toHaveLength(1);
expect(argv).not.toContain('PRIVATE MANAGED PROMPT');
expect(argv).not.toContain('provider-secret-value');

View File

@@ -170,6 +170,20 @@ describe('Pi worker process', () => {
'read,bash,edit,write,grep,find,ls,ask_user',
'--model', 'model-a',
]);
expect(buildPiRpcArgs('sessions', ['--no-session'], ['read', 'grep', 'find', 'ls']))
.toEqual([
'--mode', 'rpc',
'--offline',
'--session-dir', 'sessions',
'--no-extensions',
'--no-skills',
'--no-prompt-templates',
'--no-themes',
'--no-context-files',
'--no-approve',
'--tools', 'read,grep,find,ls',
'--no-session',
]);
});
it('correlates out-of-order responses, dispatches events, and reassembles partial lines', async () => {

View File

@@ -33,6 +33,57 @@ function baseSnapshot(): ConversationSnapshot {
const zeroCost = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 };
describe('Pi session projector', () => {
it('hydrates stable subagent nodes and bounds unknown detail versions', async () => {
const snapshot = await projectPiSessionSnapshot({
snapshot: baseSnapshot(),
workerGeneration: 1,
state: { sessionId: 'session-a', isStreaming: false, isCompacting: false },
entries: {
leafId: 'entry-unknown-result',
entries: [
{
type: 'message', id: 'entry-assistant', parentId: null,
message: {
role: 'assistant', stopReason: 'toolUse', usage: { input: 1, output: 1 },
content: [
{ type: 'toolCall', id: 'call-known', name: 'subagent', arguments: {} },
{ type: 'toolCall', id: 'call-unknown', name: 'subagent', arguments: {} },
],
},
},
{
type: 'message', id: 'entry-known-result', parentId: 'entry-assistant',
message: {
role: 'toolResult', toolCallId: 'call-known', toolName: 'subagent', content: [],
details: {
schema: 'subagent.v1', dispatchId: 'dispatch-a', mode: 'single',
tasks: [{
taskId: 'task-a', agentId: 'agent-a', toolProfile: 'coding', status: 'complete',
}],
},
},
},
{
type: 'message', id: 'entry-unknown-result', parentId: 'entry-known-result',
message: {
role: 'toolResult', toolCallId: 'call-unknown', toolName: 'subagent', content: [],
details: { schema: 'subagent.v9', raw: 'RAW_SESSION_SECRET' },
},
},
],
},
});
expect(snapshot.nodes).toContainEqual(expect.objectContaining({
kind: 'subagent', id: 'subagent:dispatch-a',
details: expect.objectContaining({ schema: 'subagent.v1' }),
}));
expect(JSON.stringify(snapshot)).not.toContain('RAW_SESSION_SECRET');
expect(snapshot.nodes).toContainEqual(expect.objectContaining({
kind: 'tool', toolCallId: 'call-unknown',
output: [expect.objectContaining({ text: 'Subagent details are unavailable for this version.' })],
}));
});
it('hydrates only the authoritative active leaf path', async () => {
const snapshot = await projectPiSessionSnapshot({
snapshot: baseSnapshot(),

View File

@@ -0,0 +1,155 @@
// @vitest-environment node
import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';
import { afterEach, describe, expect, it } from 'vitest';
import { createCodingProjectAgent } from '../../electron/coding-projects/project-config';
import {
createCodingProjectStore,
createLocalCodingProject,
createMemoryCodingProjectStorage,
} from '../../electron/coding-projects/project-store';
import type { ProviderAccount } from '../../electron/shared/providers/types';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
import type {
PiRpcCommand,
PiRpcEvent,
PiRpcRequestOptions,
PiRpcResponse,
} from '../../electron/coding-runtime/pi/rpc-client';
import {
createPiManagedSubagentChildOpener,
type PiSubagentProcessAdapter,
} from '../../electron/coding-runtime/pi/subagent-child';
import type { PiWorkerProcessOptions } from '../../electron/coding-runtime/pi/worker-process';
const roots: string[] = [];
const NOW = '2026-08-23T00:00:00.000Z';
class FakeChildProcess implements PiSubagentProcessAdapter {
private readonly listeners = new Set<(event: PiRpcEvent) => void>();
private readonly invalidationListeners = new Set<(error: PiProcessError) => void>();
stopped = false;
async start() { return this; }
async request<T = unknown>(
command: PiRpcCommand,
_options?: PiRpcRequestOptions,
): Promise<PiRpcResponse<T>> {
if (command.type === 'prompt') {
queueMicrotask(() => {
for (const listener of this.listeners) {
listener({
type: 'message_end',
message: { role: 'assistant', usage: { input: 7, output: 11, cacheRead: 2 } },
});
listener({ type: 'agent_settled' });
}
});
}
if (command.type === 'get_last_assistant_text') {
return {
type: 'response', id: 'summary', success: true,
data: { text: 'managed child summary' } as T,
};
}
return { type: 'response', id: command.type, success: true, data: {} as T };
}
subscribe(listener: (event: PiRpcEvent) => void): () => void {
this.listeners.add(listener);
return () => this.listeners.delete(listener);
}
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
this.invalidationListeners.add(listener);
return () => this.invalidationListeners.delete(listener);
}
async stop() {
this.stopped = true;
return { mode: 'stdin-close' as const, code: 0, signal: null };
}
}
afterEach(async () => {
await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
});
describe('managed Pi subagent child opener', () => {
it('opens an ephemeral managed Agent with the exact tool profile and public result', async () => {
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-subagent-child-'));
roots.push(root);
const projectPath = path.join(root, 'project');
const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
createId: () => 'project-a', now: () => NOW,
});
await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
await createCodingProjectAgent(projectPath, {
id: 'agent-a', avatarId: 'avatar-01', roleName: 'Reviewer', name: 'Agent A',
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'high' },
modelResolution: 'resolved',
responsibility: { mission: 'Review', owns: [], boundaries: [], collaborators: [], principles: [] },
prompt: 'PRIVATE CHILD PROMPT', skillIds: ['grilling'],
}, { now: NOW });
const account: ProviderAccount = {
id: 'account-a', vendorId: 'custom', label: 'Account A', authMode: 'api_key',
apiProtocol: 'openai-completions', baseUrl: 'https://provider.example/v1', model: 'model-a',
enabled: true, isDefault: true, createdAt: NOW, updatedAt: NOW,
};
const host = new PiManagedExtensionHost();
const processOptions: PiWorkerProcessOptions[] = [];
const processes: FakeChildProcess[] = [];
const opener = createPiManagedSubagentChildOpener({
projectStore,
executablePath: 'electron.exe',
cliPath: 'pi-cli.js',
userDataDir: path.join(root, 'user-data'),
bundledSkillsDir: path.resolve('resources/coding-skills'),
extensionHost: host,
loadProviderInput: async () => ({ accounts: [account], modelSummaries: [] }),
resolveCredential: async () => 'provider-secret',
getRevision: () => ({ provider: 3, resources: 4 }),
createProcess: (options) => {
processOptions.push(options);
const child = new FakeChildProcess();
processes.push(child);
return child;
},
});
const identity = {
conversationId: 'conversation-a', workerGeneration: 2, runId: 'run-a', projectId: 'project-a',
dispatchId: 'dispatch-a', agentId: 'agent-a',
};
const readOnly = await opener({
...identity, taskId: 'task-read', toolProfile: 'read-only',
});
const coding = await opener({
...identity, taskId: 'task-code', toolProfile: 'coding',
});
expect(processOptions.map(({ tools }) => tools)).toEqual([
['read', 'grep', 'find', 'ls'],
['read', 'bash', 'edit', 'write', 'grep', 'find', 'ls'],
]);
for (const options of processOptions) {
expect(options.additionalArgs).toContain('--no-session');
expect(options.additionalArgs).not.toContain('--session-id');
expect(options.additionalArgs?.some((argument) => argument.includes('grilling'))).toBe(true);
expect(JSON.stringify(options.additionalArgs)).not.toContain('PRIVATE CHILD PROMPT');
expect(JSON.stringify(options.additionalArgs)).not.toContain('provider-secret');
expect(Object.values(options.env ?? {})).toContain('provider-secret');
expect(options.env?.MAKELORE_PI_WORKER_ROLE).toBe('child');
}
await expect(readOnly.run('Inspect', new AbortController().signal)).resolves.toEqual({
summary: 'managed child summary',
usage: { inputTokens: 7, outputTokens: 11, cacheReadTokens: 2 },
});
await readOnly.stop();
await coding.stop();
expect(processes.every(({ stopped }) => stopped)).toBe(true);
await host.close();
});
});

View File

@@ -0,0 +1,241 @@
// @vitest-environment node
import { describe, expect, it } from 'vitest';
import {
PiSubagentChildError,
PiSubagentScheduler,
parsePiSubagentDispatchRequest,
type PiSubagentChild,
type PiSubagentChildOpenInput,
} from '../../electron/coding-runtime/pi/subagent';
import { PiProcessBudget } from '../../electron/coding-runtime/pi/worker-pool';
function parent(runId = 'run-a') {
return {
conversationId: 'conversation-a',
workerGeneration: 1,
runId,
projectId: 'project-a',
};
}
function task(agentId: string, toolProfile: 'read-only' | 'coding' = 'read-only') {
return { agentId, task: `Inspect ${agentId}`, toolProfile };
}
function deferred() {
let resolve!: () => void;
const promise = new Promise<void>((done) => { resolve = done; });
return { promise, resolve };
}
describe('Pi subagent scheduler', () => {
it('validates one bounded dispatch mode and rejects a ninth task', () => {
expect(parsePiSubagentDispatchRequest({ mode: 'single', tasks: [task('one')] })).toEqual({
mode: 'single', tasks: [task('one')],
});
expect(parsePiSubagentDispatchRequest({
mode: 'parallel', tasks: Array.from({ length: 8 }, (_, index) => task(`agent-${index}`)),
}).tasks).toHaveLength(8);
expect(() => parsePiSubagentDispatchRequest({
mode: 'parallel', tasks: Array.from({ length: 9 }, (_, index) => task(`agent-${index}`)),
})).toThrowError('Subagent dispatch accepts at most 8 tasks');
expect(() => parsePiSubagentDispatchRequest({
mode: 'single', tasks: [task('one'), task('two')],
})).toThrowError('Single subagent dispatch requires exactly one task');
expect(() => parsePiSubagentDispatchRequest({ mode: 'unknown', tasks: [task('one')] }))
.toThrowError('Subagent dispatch mode is invalid');
});
it('shares one FIFO four-child cap across two parents and releases the process budget', async () => {
const gate = deferred();
const processBudget = new PiProcessBudget(8);
let running = 0;
let maxRunning = 0;
const opened: PiSubagentChildOpenInput[] = [];
const scheduler = new PiSubagentScheduler({
processBudget,
createId: (() => {
let id = 0;
return (kind) => `${kind}-${++id}`;
})(),
openChild: async (input) => {
opened.push(structuredClone(input));
return {
id: input.taskId,
async run() {
running += 1;
maxRunning = Math.max(maxRunning, running);
await gate.promise;
running -= 1;
return {
summary: `done ${input.agentId}`,
usage: { inputTokens: 1, outputTokens: 2 },
};
},
async stop() {},
} satisfies PiSubagentChild;
},
});
const request = {
mode: 'parallel' as const,
tasks: Array.from({ length: 4 }, (_, index) => task(`agent-${index}`)),
};
const left = scheduler.dispatch({ ...parent('run-left'), request });
const right = scheduler.dispatch({ ...parent('run-right'), request });
await expect.poll(() => running).toBe(4);
expect(maxRunning).toBe(4);
expect(processBudget.activeCount).toBe(4);
gate.resolve();
const [leftResult, rightResult] = await Promise.all([left, right]);
expect(leftResult.details.tasks).toHaveLength(4);
expect(rightResult.details.tasks).toHaveLength(4);
expect(opened).toHaveLength(8);
expect(maxRunning).toBe(4);
expect(processBudget.activeCount).toBe(0);
await scheduler.close();
});
it('preserves parallel sibling results and stops a chain at its first failure', async () => {
const scheduler = new PiSubagentScheduler({
processBudget: new PiProcessBudget(8),
createId: (() => {
let id = 0;
return (kind) => `${kind}-${++id}`;
})(),
openChild: async (input) => ({
id: input.taskId,
async run(prompt) {
if (input.agentId === 'broken') {
throw new PiSubagentChildError('SUBAGENT_CHILD_FAILED');
}
return {
summary: `${input.agentId}:${prompt}`,
usage: { inputTokens: 3, outputTokens: 5 },
};
},
async stop() {},
}),
});
const parallel = await scheduler.dispatch({
...parent('parallel'),
request: {
mode: 'parallel',
tasks: [task('left'), task('broken'), task('right')],
},
});
expect(parallel.details.tasks).toEqual(expect.arrayContaining([
expect.objectContaining({ agentId: 'left', status: 'complete', summary: 'left:Inspect left' }),
expect.objectContaining({ agentId: 'broken', status: 'error', errorCode: 'SUBAGENT_CHILD_FAILED' }),
expect.objectContaining({ agentId: 'right', status: 'complete', usage: { inputTokens: 3, outputTokens: 5 } }),
]));
const chain = await scheduler.dispatch({
...parent('chain'),
request: {
mode: 'chain',
tasks: [
{ agentId: 'first', task: 'first', toolProfile: 'read-only' },
{ agentId: 'broken', task: 'review {previous}', toolProfile: 'read-only' },
{ agentId: 'never', task: 'never', toolProfile: 'coding' },
],
},
});
expect(chain.details.tasks.map(({ status }) => status)).toEqual(['complete', 'error', 'skipped']);
expect(chain.details.tasks[1]).toMatchObject({ errorCode: 'SUBAGENT_CHILD_FAILED' });
await scheduler.close();
});
it('aborts every unfinished child with the parent and leaves no permit or process lease', async () => {
const processBudget = new PiProcessBudget(8);
const stopped: string[] = [];
const scheduler = new PiSubagentScheduler({
processBudget,
openChild: async (input) => ({
id: input.taskId,
async run(_prompt, signal) {
await new Promise<void>((_resolve, reject) => {
const abort = () => reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
if (signal.aborted) abort();
else signal.addEventListener('abort', abort, { once: true });
});
return { summary: 'unreachable' };
},
async stop() { stopped.push(input.taskId); },
}),
});
const controller = new AbortController();
const flight = scheduler.dispatch({
...parent('abort'),
request: { mode: 'parallel', tasks: [task('left'), task('right', 'coding')] },
}, { signal: controller.signal });
await expect.poll(() => processBudget.activeCount).toBe(2);
controller.abort();
const result = await flight;
expect(result.details.tasks.map(({ status }) => status)).toEqual(['aborted', 'aborted']);
expect(stopped).toHaveLength(2);
expect(processBudget.activeCount).toBe(0);
await scheduler.close();
});
it('marks the rest of an in-flight chain aborted when its parent aborts', async () => {
const processBudget = new PiProcessBudget(8);
const scheduler = new PiSubagentScheduler({
processBudget,
openChild: async (input) => ({
id: input.taskId,
async run(_prompt, signal) {
await new Promise<void>((_resolve, reject) => {
const abort = () => reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
if (signal.aborted) abort();
else signal.addEventListener('abort', abort, { once: true });
});
return { summary: 'unreachable' };
},
async stop() {},
}),
});
const controller = new AbortController();
const flight = scheduler.dispatch({
...parent('abort-chain'),
request: { mode: 'chain', tasks: [task('one'), task('two'), task('three')] },
}, { signal: controller.signal });
await expect.poll(() => processBudget.activeCount).toBe(1);
controller.abort();
const result = await flight;
expect(result.details.tasks.map(({ status }) => status))
.toEqual(['aborted', 'aborted', 'aborted']);
expect(processBudget.activeCount).toBe(0);
await scheduler.close();
});
it('reclaims an idle parent process instead of deadlocking on a full global budget', async () => {
const processBudget = new PiProcessBudget(2);
const firstParent = await processBudget.acquire();
const idleParent = await processBudget.acquire();
let reclaimed = 0;
const scheduler = new PiSubagentScheduler({
processBudget,
reclaimProcessCapacity: async () => {
reclaimed += 1;
idleParent.release();
return true;
},
openChild: async (input) => ({
id: input.taskId,
async run() { return { summary: 'done' }; },
async stop() {},
}),
});
await expect(scheduler.dispatch({
...parent('full-budget'),
request: { mode: 'single', tasks: [task('agent-a')] },
})).resolves.toMatchObject({
details: { tasks: [{ status: 'complete' }] },
});
expect(reclaimed).toBe(1);
expect(processBudget.activeCount).toBe(1);
firstParent.release();
await scheduler.close();
});
});

View File

@@ -73,4 +73,53 @@ describe('locked Pi worker process smoke', () => {
await extensionHost.close();
}
}, 15_000);
it('starts a real ephemeral read-only child with the child extension role', async () => {
const requireFromProject = createRequire(resolve('package.json'));
const electronExecutable = requireFromProject('electron') as string;
const packageRoot = realpathSync(resolve(
'node_modules',
'@earendil-works',
'pi-coding-agent',
));
const root = await mkdtemp(join(tmpdir(), 'makelore-pi-real-child-'));
scratchRoots.push(root);
const configDir = join(root, 'config');
const sessionDir = join(root, 'sessions');
const cwd = join(root, 'project');
await Promise.all([mkdir(configDir), mkdir(sessionDir), mkdir(cwd)]);
const extensionHost = new PiManagedExtensionHost();
const extension = await extensionHost.registerWorker({
conversationId: 'real-parent',
generation: 1,
projectId: 'real-project',
extensionsDir: join(root, 'extensions'),
role: 'child',
runId: 'real-parent-run',
});
const worker = await new PiWorkerProcess({
executablePath: electronExecutable,
cliPath: join(packageRoot, 'dist', 'cli.js'),
cwd,
configDir,
sessionDir,
tools: ['read', 'grep', 'find', 'ls'],
additionalArgs: ['--extension', extension.extensionPath, '--no-session'],
env: extension.env,
sensitiveValues: extension.sensitiveValues,
commandTimeoutMs: 5_000,
}).start();
try {
await expect(worker.request({ type: 'get_state' })).resolves.toMatchObject({
type: 'response', command: 'get_state', success: true,
});
expect(worker.stderrDiagnostic).not.toContain('Failed to load extension');
await expect(worker.stop()).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
} finally {
await worker.stop().catch(() => undefined);
await extension.dispose();
await extensionHost.close();
}
}, 15_000);
});