feat: add Pi subagent scheduler and child runtime
This commit is contained in:
@@ -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',
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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');
|
||||
|
||||
@@ -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 () => {
|
||||
|
||||
@@ -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(),
|
||||
|
||||
155
tests/unit/pi-subagent-child.test.ts
Normal file
155
tests/unit/pi-subagent-child.test.ts
Normal 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();
|
||||
});
|
||||
});
|
||||
241
tests/unit/pi-subagent.test.ts
Normal file
241
tests/unit/pi-subagent.test.ts
Normal 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();
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user