feat: add Pi subagent scheduler and child runtime
This commit is contained in:
@@ -7,6 +7,7 @@ import type {
|
||||
PublicUsage,
|
||||
} from '../contracts';
|
||||
import type { PiRpcEvent } from './rpc-client';
|
||||
import { subagentDetailsOfResult } from '../subagent-protocol';
|
||||
|
||||
export interface PiEventProjectorOptions {
|
||||
createId(): string;
|
||||
@@ -89,6 +90,49 @@ function outputBlocks(toolId: string, value: unknown): ConversationContentBlock[
|
||||
});
|
||||
}
|
||||
|
||||
function projectToolResult(
|
||||
snapshot: ConversationSnapshot,
|
||||
tool: ConversationToolNode,
|
||||
status: ConversationToolNode['status'],
|
||||
value: unknown,
|
||||
): ConversationPatch[] {
|
||||
const details = subagentDetailsOfResult(value);
|
||||
const result = asRecord(value);
|
||||
const unknownSubagentDetails = tool.toolName === 'subagent'
|
||||
&& result?.details !== undefined
|
||||
&& !details;
|
||||
const output = unknownSubagentDetails
|
||||
? [{
|
||||
kind: 'text' as const,
|
||||
id: `${tool.id}:output:unavailable`,
|
||||
text: 'Subagent details are unavailable for this version.',
|
||||
status: 'complete' as const,
|
||||
}]
|
||||
: outputBlocks(tool.id, value);
|
||||
const patches: ConversationPatch[] = [{
|
||||
op: 'tool.upsert',
|
||||
node: {
|
||||
...tool,
|
||||
status,
|
||||
output,
|
||||
...(details ? { details } : {}),
|
||||
},
|
||||
}];
|
||||
const runId = snapshot.run.runId;
|
||||
if (details && runId) {
|
||||
patches.push({
|
||||
op: 'subagent.upsert',
|
||||
node: {
|
||||
kind: 'subagent',
|
||||
id: `subagent:${details.dispatchId}`,
|
||||
runId,
|
||||
details,
|
||||
},
|
||||
});
|
||||
}
|
||||
return patches;
|
||||
}
|
||||
|
||||
function currentTool(snapshot: ConversationSnapshot, toolCallId: string): ConversationToolNode | undefined {
|
||||
return snapshot.nodes.find(
|
||||
(node): node is ConversationToolNode => node.kind === 'tool'
|
||||
@@ -613,28 +657,19 @@ export class PiEventProjector {
|
||||
if (typeof event.toolCallId !== 'string') return [];
|
||||
const tool = currentTool(snapshot, event.toolCallId);
|
||||
if (!tool) return [];
|
||||
return [{
|
||||
op: 'tool.upsert',
|
||||
node: {
|
||||
...tool,
|
||||
status: 'running',
|
||||
output: outputBlocks(tool.id, event.partialResult),
|
||||
},
|
||||
}];
|
||||
return projectToolResult(snapshot, tool, 'running', event.partialResult);
|
||||
}
|
||||
|
||||
if (event.type === 'tool_execution_end') {
|
||||
if (typeof event.toolCallId !== 'string') return [];
|
||||
const tool = currentTool(snapshot, event.toolCallId);
|
||||
if (!tool) return [];
|
||||
return [{
|
||||
op: 'tool.upsert',
|
||||
node: {
|
||||
...tool,
|
||||
status: event.isError === true ? 'error' : 'complete',
|
||||
output: outputBlocks(tool.id, event.result),
|
||||
},
|
||||
}];
|
||||
return projectToolResult(
|
||||
snapshot,
|
||||
tool,
|
||||
event.isError === true ? 'error' : 'complete',
|
||||
event.result,
|
||||
);
|
||||
}
|
||||
|
||||
if (event.type === 'message_end') {
|
||||
@@ -642,14 +677,12 @@ export class PiEventProjector {
|
||||
if (message?.role === 'toolResult' && typeof message.toolCallId === 'string') {
|
||||
const tool = currentTool(snapshot, message.toolCallId);
|
||||
if (!tool) return [];
|
||||
return [{
|
||||
op: 'tool.upsert',
|
||||
node: {
|
||||
...tool,
|
||||
status: message.isError === true ? 'error' : 'complete',
|
||||
output: outputBlocks(tool.id, message),
|
||||
},
|
||||
}];
|
||||
return projectToolResult(
|
||||
snapshot,
|
||||
tool,
|
||||
message.isError === true ? 'error' : 'complete',
|
||||
message,
|
||||
);
|
||||
}
|
||||
if (message?.role !== 'assistant') return [];
|
||||
const current = currentAssistant(snapshot);
|
||||
|
||||
@@ -8,6 +8,11 @@ import {
|
||||
PiProjectWriteLeaseCoordinator,
|
||||
type PiProjectWriteLease,
|
||||
} from './write-lease';
|
||||
import type { PiGenerationResourceInput } from './worker-pool';
|
||||
import {
|
||||
parsePiSubagentDispatchRequest,
|
||||
type PiSubagentScheduler,
|
||||
} from './subagent';
|
||||
|
||||
const MAX_REQUEST_BYTES = 64 * 1024;
|
||||
|
||||
@@ -16,6 +21,7 @@ interface WorkerRegistrationRecord {
|
||||
conversationId: string;
|
||||
generation: number;
|
||||
projectId: string;
|
||||
role: 'parent' | 'child';
|
||||
contextFile: string;
|
||||
runId: string | null;
|
||||
leases: Map<string, PiProjectWriteLease>;
|
||||
@@ -34,9 +40,11 @@ export interface RegisterPiExtensionWorkerInput {
|
||||
generation: number;
|
||||
projectId: string;
|
||||
extensionsDir: string;
|
||||
role?: 'parent' | 'child';
|
||||
runId?: string;
|
||||
}
|
||||
|
||||
interface BridgeRequest {
|
||||
interface LeaseBridgeRequest {
|
||||
action: 'lease.acquire' | 'lease.release';
|
||||
conversationId: string;
|
||||
workerGeneration: number;
|
||||
@@ -45,17 +53,35 @@ interface BridgeRequest {
|
||||
leaseId?: string;
|
||||
}
|
||||
|
||||
interface SubagentBridgeRequest {
|
||||
action: 'subagent.dispatch';
|
||||
conversationId: string;
|
||||
workerGeneration: number;
|
||||
runId: string;
|
||||
resourceId: string;
|
||||
request: unknown;
|
||||
}
|
||||
|
||||
type BridgeRequest = LeaseBridgeRequest | SubagentBridgeRequest;
|
||||
|
||||
export interface PiExtensionSubagentBridge {
|
||||
scheduler: PiSubagentScheduler;
|
||||
trackGenerationResource?(input: PiGenerationResourceInput): () => void;
|
||||
}
|
||||
|
||||
function recordValue(value: unknown): value is Record<string, unknown> {
|
||||
return value !== null && typeof value === 'object' && !Array.isArray(value);
|
||||
}
|
||||
|
||||
function bridgeRequest(value: unknown): value is BridgeRequest {
|
||||
if (!recordValue(value)) return false;
|
||||
return (value.action === 'lease.acquire' || value.action === 'lease.release')
|
||||
&& typeof value.conversationId === 'string'
|
||||
const common = typeof value.conversationId === 'string'
|
||||
&& Number.isSafeInteger(value.workerGeneration)
|
||||
&& typeof value.runId === 'string'
|
||||
&& typeof value.resourceId === 'string'
|
||||
&& typeof value.resourceId === 'string';
|
||||
if (!common) return false;
|
||||
if (value.action === 'subagent.dispatch') return 'request' in value;
|
||||
return (value.action === 'lease.acquire' || value.action === 'lease.release')
|
||||
&& (value.leaseId === undefined || typeof value.leaseId === 'string');
|
||||
}
|
||||
|
||||
@@ -64,6 +90,7 @@ export class PiManagedExtensionHost {
|
||||
private readonly registrations = new Map<string, WorkerRegistrationRecord>();
|
||||
private readonly runBindings = new Map<string, string>();
|
||||
private readonly requestFlights = new Set<Promise<void>>();
|
||||
private subagentBridge: PiExtensionSubagentBridge | undefined;
|
||||
private server: Server | null = null;
|
||||
private bridgeUrl: string | null = null;
|
||||
private startFlight: Promise<string> | null = null;
|
||||
@@ -73,6 +100,10 @@ export class PiManagedExtensionHost {
|
||||
this.leases = leases;
|
||||
}
|
||||
|
||||
configureSubagents(bridge: PiExtensionSubagentBridge): void {
|
||||
this.subagentBridge = bridge;
|
||||
}
|
||||
|
||||
async registerWorker(input: RegisterPiExtensionWorkerInput): Promise<PiExtensionWorkerRegistration> {
|
||||
if (!Number.isSafeInteger(input.generation) || input.generation <= 0) {
|
||||
throw new Error('Worker generation must be a positive safe integer');
|
||||
@@ -80,6 +111,10 @@ export class PiManagedExtensionHost {
|
||||
const bridgeUrl = await this.start();
|
||||
await mkdir(input.extensionsDir, { recursive: true });
|
||||
const extensionPath = await materializeMakelorePiExtension(input.extensionsDir);
|
||||
const role = input.role ?? 'parent';
|
||||
if (role === 'child' && !input.runId?.trim()) {
|
||||
throw new Error('Child extension registration requires a parent run id');
|
||||
}
|
||||
const token = randomBytes(32).toString('base64url');
|
||||
const contextFile = path.join(input.extensionsDir, `worker-${randomUUID()}.json`);
|
||||
const record: WorkerRegistrationRecord = {
|
||||
@@ -87,8 +122,11 @@ export class PiManagedExtensionHost {
|
||||
conversationId: input.conversationId,
|
||||
generation: input.generation,
|
||||
projectId: input.projectId,
|
||||
role,
|
||||
contextFile,
|
||||
runId: this.runBindings.get(input.conversationId) ?? null,
|
||||
runId: role === 'child'
|
||||
? input.runId as string
|
||||
: this.runBindings.get(input.conversationId) ?? null,
|
||||
leases: new Map(),
|
||||
waiters: new Map(),
|
||||
};
|
||||
@@ -101,6 +139,7 @@ export class PiManagedExtensionHost {
|
||||
MAKELORE_PI_BRIDGE_URL: bridgeUrl,
|
||||
MAKELORE_PI_WORKER_TOKEN: token,
|
||||
MAKELORE_PI_CONTEXT_FILE: contextFile,
|
||||
MAKELORE_PI_WORKER_ROLE: role,
|
||||
},
|
||||
sensitiveValues: [token],
|
||||
dispose: async () => {
|
||||
@@ -206,6 +245,10 @@ export class PiManagedExtensionHost {
|
||||
this.respond(response, 400, { error: 'Bridge resource id is required' });
|
||||
return;
|
||||
}
|
||||
if (value.action === 'subagent.dispatch') {
|
||||
await this.dispatchSubagents(request, response, record, value);
|
||||
return;
|
||||
}
|
||||
if (value.action === 'lease.release') {
|
||||
const lease = record.leases.get(value.resourceId);
|
||||
if (!lease || !value.leaseId || lease.id !== value.leaseId) {
|
||||
@@ -251,6 +294,75 @@ export class PiManagedExtensionHost {
|
||||
}
|
||||
}
|
||||
|
||||
private async dispatchSubagents(
|
||||
request: IncomingMessage,
|
||||
response: ServerResponse,
|
||||
record: WorkerRegistrationRecord,
|
||||
value: SubagentBridgeRequest,
|
||||
): Promise<void> {
|
||||
if (record.role !== 'parent') {
|
||||
this.respond(response, 403, { error: 'Child workers cannot dispatch subagents' });
|
||||
return;
|
||||
}
|
||||
if (!this.subagentBridge) {
|
||||
this.respond(response, 503, { error: 'Subagent scheduler is unavailable' });
|
||||
return;
|
||||
}
|
||||
try {
|
||||
parsePiSubagentDispatchRequest(value.request);
|
||||
} catch {
|
||||
this.respond(response, 400, { error: 'Subagent dispatch request is invalid' });
|
||||
return;
|
||||
}
|
||||
if (record.waiters.has(value.resourceId) || record.leases.has(value.resourceId)) {
|
||||
this.respond(response, 409, { error: 'Bridge resource already exists' });
|
||||
return;
|
||||
}
|
||||
const controller = new AbortController();
|
||||
record.waiters.set(value.resourceId, controller);
|
||||
const cancel = () => controller.abort();
|
||||
request.once('aborted', cancel);
|
||||
response.once('close', cancel);
|
||||
let untrack = () => undefined;
|
||||
try {
|
||||
untrack = this.subagentBridge.trackGenerationResource?.({
|
||||
conversationId: record.conversationId,
|
||||
kind: 'child',
|
||||
id: value.resourceId,
|
||||
cancel,
|
||||
}) ?? untrack;
|
||||
response.writeHead(200, {
|
||||
'content-type': 'application/x-ndjson; charset=utf-8',
|
||||
...(this.closing ? { connection: 'close' } : {}),
|
||||
});
|
||||
const result = await this.subagentBridge.scheduler.dispatch({
|
||||
conversationId: record.conversationId,
|
||||
workerGeneration: record.generation,
|
||||
runId: value.runId,
|
||||
projectId: record.projectId,
|
||||
request: value.request,
|
||||
}, {
|
||||
signal: controller.signal,
|
||||
onUpdate: (details) => {
|
||||
if (!response.writableEnded && !response.destroyed) {
|
||||
response.write(`${JSON.stringify({ details })}\n`);
|
||||
}
|
||||
},
|
||||
});
|
||||
if (!response.writableEnded && !response.destroyed) {
|
||||
response.end(`${JSON.stringify({ details: result.details, done: true })}\n`);
|
||||
}
|
||||
} catch {
|
||||
if (!response.headersSent) this.respond(response, 400, { error: 'Subagent dispatch failed' });
|
||||
else if (!response.writableEnded && !response.destroyed) response.end();
|
||||
} finally {
|
||||
request.removeListener('aborted', cancel);
|
||||
response.removeListener('close', cancel);
|
||||
record.waiters.delete(value.resourceId);
|
||||
untrack();
|
||||
}
|
||||
}
|
||||
|
||||
private readBody(request: IncomingMessage): Promise<unknown> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const chunks: Buffer[] = [];
|
||||
@@ -286,7 +398,8 @@ export class PiManagedExtensionHost {
|
||||
|
||||
private findWorker(conversationId: string, generation: number): WorkerRegistrationRecord | undefined {
|
||||
return [...this.registrations.values()].find((record) => (
|
||||
record.conversationId === conversationId && record.generation === generation
|
||||
record.role === 'parent'
|
||||
&& record.conversationId === conversationId && record.generation === generation
|
||||
));
|
||||
}
|
||||
|
||||
@@ -294,6 +407,7 @@ export class PiManagedExtensionHost {
|
||||
await atomicWriteJson(record.contextFile, {
|
||||
conversationId: record.conversationId,
|
||||
workerGeneration: record.generation,
|
||||
role: record.role,
|
||||
...(record.runId ? { runId: record.runId } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
@@ -1,13 +1,14 @@
|
||||
import path from 'node:path';
|
||||
import { atomicWriteText } from '../../../coding-projects/atomic-json';
|
||||
|
||||
export const MAKELORE_PI_EXTENSION_VERSION = 1;
|
||||
export const MAKELORE_PI_EXTENSION_VERSION = 2;
|
||||
export const MAKELORE_PI_EXTENSION_FILENAME = `makelore-runtime-v${MAKELORE_PI_EXTENSION_VERSION}.mjs`;
|
||||
|
||||
const BUNDLE_SOURCE = String.raw`
|
||||
import { readFile } from 'node:fs/promises';
|
||||
|
||||
const MUTATION_TOOLS = new Set(['bash', 'edit', 'write']);
|
||||
const WORKER_ROLE = process.env.MAKELORE_PI_WORKER_ROLE || 'parent';
|
||||
const leases = new Map();
|
||||
|
||||
async function runtimeContext() {
|
||||
@@ -32,6 +33,48 @@ async function bridge(action, body, signal) {
|
||||
return result;
|
||||
}
|
||||
|
||||
async function bridgeStream(action, body, signal, onUpdate) {
|
||||
const context = await runtimeContext();
|
||||
const response = await fetch(process.env.MAKELORE_PI_BRIDGE_URL, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
authorization: 'Bearer ' + process.env.MAKELORE_PI_WORKER_TOKEN,
|
||||
'content-type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify({ ...context, action, ...body }),
|
||||
signal,
|
||||
});
|
||||
if (!response.ok) {
|
||||
const result = await response.json().catch(() => ({}));
|
||||
throw new Error(result.error || 'Makelore runtime bridge rejected the request');
|
||||
}
|
||||
if (!response.body) throw new Error('Makelore runtime bridge returned no stream');
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffered = '';
|
||||
let details;
|
||||
while (true) {
|
||||
const { value, done } = await reader.read();
|
||||
buffered += decoder.decode(value || new Uint8Array(), { stream: !done });
|
||||
let newline = buffered.indexOf('\n');
|
||||
while (newline >= 0) {
|
||||
const line = buffered.slice(0, newline);
|
||||
buffered = buffered.slice(newline + 1);
|
||||
if (line) {
|
||||
const item = JSON.parse(line);
|
||||
if (item.details?.schema === 'subagent.v1') {
|
||||
details = item.details;
|
||||
onUpdate?.({ content: [], details });
|
||||
}
|
||||
}
|
||||
newline = buffered.indexOf('\n');
|
||||
}
|
||||
if (done) break;
|
||||
}
|
||||
if (!details) throw new Error('Makelore subagent stream returned no details');
|
||||
return details;
|
||||
}
|
||||
|
||||
async function releaseLease(toolCallId) {
|
||||
const leaseId = leases.get(toolCallId);
|
||||
if (!leaseId) return;
|
||||
@@ -44,7 +87,7 @@ async function releaseAll() {
|
||||
}
|
||||
|
||||
export default function makeloreRuntime(pi) {
|
||||
pi.registerTool({
|
||||
if (WORKER_ROLE === 'parent') pi.registerTool({
|
||||
name: 'ask_user',
|
||||
label: 'Ask user',
|
||||
description: 'Ask the user for a selection, confirmation, short input, or editor text.',
|
||||
@@ -77,6 +120,52 @@ export default function makeloreRuntime(pi) {
|
||||
},
|
||||
});
|
||||
|
||||
if (WORKER_ROLE === 'parent') pi.registerTool({
|
||||
name: 'subagent',
|
||||
label: 'Subagent',
|
||||
description: 'Dispatch one or more managed project Agents in single, parallel, or chain mode.',
|
||||
parameters: {
|
||||
type: 'object',
|
||||
additionalProperties: false,
|
||||
required: ['mode', 'tasks'],
|
||||
properties: {
|
||||
mode: { type: 'string', enum: ['single', 'parallel', 'chain'] },
|
||||
tasks: {
|
||||
type: 'array', minItems: 1, maxItems: 8,
|
||||
items: {
|
||||
type: 'object',
|
||||
additionalProperties: false,
|
||||
required: ['agentId', 'task', 'toolProfile'],
|
||||
properties: {
|
||||
agentId: { type: 'string' },
|
||||
task: { type: 'string' },
|
||||
toolProfile: { type: 'string', enum: ['read-only', 'coding'] },
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
async execute(toolCallId, params, signal, onUpdate) {
|
||||
const details = await bridgeStream(
|
||||
'subagent.dispatch',
|
||||
{ resourceId: toolCallId, request: params },
|
||||
signal,
|
||||
onUpdate,
|
||||
);
|
||||
const complete = details.tasks.filter((task) => task.status === 'complete').length;
|
||||
const failed = details.tasks.length - complete;
|
||||
return {
|
||||
content: [{
|
||||
type: 'text',
|
||||
text: failed === 0
|
||||
? complete + ' subagent task(s) completed'
|
||||
: complete + ' completed; ' + failed + ' did not complete',
|
||||
}],
|
||||
details,
|
||||
};
|
||||
},
|
||||
});
|
||||
|
||||
pi.on('tool_call', async (event, ctx) => {
|
||||
if (!MUTATION_TOOLS.has(event.toolName)) return;
|
||||
ctx.ui.setStatus('makelore.write-lease', '等待项目写入');
|
||||
|
||||
@@ -75,6 +75,7 @@ import {
|
||||
projectPiSessionSnapshot,
|
||||
} from './session-projector';
|
||||
import { PiManagedExtensionHost } from './extension-host';
|
||||
import type { PiSubagentScheduler } from './subagent';
|
||||
import {
|
||||
PiInteractionStore,
|
||||
type PiInteractionResponse,
|
||||
@@ -98,6 +99,7 @@ export interface PiConversationRuntimeOptions {
|
||||
isAuthenticationError?(error: unknown): boolean;
|
||||
refreshCredential?(accountId: string): Promise<void>;
|
||||
extensionHost?: PiManagedExtensionHost;
|
||||
subagentScheduler?: PiSubagentScheduler;
|
||||
getDraftRevision?(conversationId: string): number;
|
||||
knownExtensionWidgetKeys?: readonly string[];
|
||||
onExtensionUiProjection?(projection: PiExtensionUiProjection): void;
|
||||
@@ -480,6 +482,15 @@ export class PiConversationRuntime implements CodingConversationRuntime {
|
||||
this.isAuthenticationError = options.isAuthenticationError;
|
||||
this.refreshCredential = options.refreshCredential;
|
||||
this.extensionHost = options.extensionHost;
|
||||
if (options.subagentScheduler && !this.extensionHost) {
|
||||
throw new Error('Subagent scheduler requires the managed extension host');
|
||||
}
|
||||
if (options.subagentScheduler) {
|
||||
this.extensionHost?.configureSubagents({
|
||||
scheduler: options.subagentScheduler,
|
||||
trackGenerationResource: (input) => this.pool.trackGenerationResource(input),
|
||||
});
|
||||
}
|
||||
this.interactions = new PiInteractionStore(this.pool, (interaction) => {
|
||||
this.emit(interaction.conversationId, { op: 'interaction.upsert', interaction }, interaction.runId);
|
||||
});
|
||||
|
||||
@@ -12,6 +12,7 @@ import type {
|
||||
PiImageProjectionInput,
|
||||
PiProjectedAttachment,
|
||||
} from './event-projector';
|
||||
import { subagentDetailsOfResult } from '../subagent-protocol';
|
||||
|
||||
export interface PiSessionSnapshotInput {
|
||||
snapshot: ConversationSnapshot;
|
||||
@@ -184,6 +185,7 @@ async function projectEntries(
|
||||
): Promise<ConversationNode[]> {
|
||||
const nodes: ConversationNode[] = [];
|
||||
const tools = new Map<string, ConversationToolNode>();
|
||||
const subagentIds = new Set<string>();
|
||||
for (const entry of path) {
|
||||
if (entry.type === 'compaction') {
|
||||
nodes.push({
|
||||
@@ -211,6 +213,27 @@ async function projectEntries(
|
||||
'output',
|
||||
);
|
||||
tool.status = message.isError === true ? 'error' : 'complete';
|
||||
const details = subagentDetailsOfResult(message);
|
||||
if (details) {
|
||||
tool.details = details;
|
||||
const id = `subagent:${details.dispatchId}`;
|
||||
if (!subagentIds.has(id)) {
|
||||
subagentIds.add(id);
|
||||
nodes.push({
|
||||
kind: 'subagent',
|
||||
id,
|
||||
runId: input.snapshot.run.runId ?? `session:${entry.id}`,
|
||||
details,
|
||||
});
|
||||
}
|
||||
} else if (tool.toolName === 'subagent' && message.details !== undefined) {
|
||||
tool.output = [{
|
||||
kind: 'text',
|
||||
id: `${tool.id}:output:unavailable`,
|
||||
text: 'Subagent details are unavailable for this version.',
|
||||
status: 'complete',
|
||||
}];
|
||||
}
|
||||
continue;
|
||||
}
|
||||
if (message.role !== 'user' && message.role !== 'assistant') continue;
|
||||
|
||||
227
electron/coding-runtime/pi/subagent-child.ts
Normal file
227
electron/coding-runtime/pi/subagent-child.ts
Normal file
@@ -0,0 +1,227 @@
|
||||
import type { ModelSummary, ProviderAccount } from '../../shared/providers/types';
|
||||
import { readCodingProjectConfigV2 } from '../../coding-projects/project-config';
|
||||
import type { CodingProjectStore } from '../../coding-projects/project-store';
|
||||
import type { PublicUsage } from '../contracts';
|
||||
import type { PiManagedInputRevision } from './managed-input-revision';
|
||||
import {
|
||||
buildPiProviderCatalog,
|
||||
buildPiWorkerCredentialProjection,
|
||||
selectPiProviderModel,
|
||||
writePiProviderCatalog,
|
||||
} from './provider-config';
|
||||
import {
|
||||
buildPiManagedInputArgs,
|
||||
ensurePiManagedPaths,
|
||||
materializePiAgentResources,
|
||||
} from './resource-loader';
|
||||
import type {
|
||||
PiRpcCommand,
|
||||
PiRpcEvent,
|
||||
PiRpcRequestOptions,
|
||||
PiRpcResponse,
|
||||
} from './rpc-client';
|
||||
import type {
|
||||
PiSubagentChild,
|
||||
PiSubagentChildOpenInput,
|
||||
PiSubagentChildResult,
|
||||
} from './subagent';
|
||||
import { PiSubagentChildError } from './subagent';
|
||||
import { PiWorkerProcess, type PiWorkerProcessOptions, type PiWorkerStopResult } from './worker-process';
|
||||
import type { PiProcessError } from './process-errors';
|
||||
import type { PiManagedExtensionHost } from './extension-host';
|
||||
|
||||
const READ_ONLY_TOOLS = ['read', 'grep', 'find', 'ls'] as const;
|
||||
const CODING_TOOLS = ['read', 'bash', 'edit', 'write', 'grep', 'find', 'ls'] as const;
|
||||
|
||||
export interface PiSubagentProcessAdapter {
|
||||
start(): Promise<unknown>;
|
||||
request<T = unknown>(
|
||||
command: PiRpcCommand,
|
||||
options?: PiRpcRequestOptions,
|
||||
): Promise<PiRpcResponse<T>>;
|
||||
subscribe(listener: (event: PiRpcEvent) => void): () => void;
|
||||
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void;
|
||||
stop(): Promise<PiWorkerStopResult>;
|
||||
}
|
||||
|
||||
export interface PiManagedSubagentChildOpenerOptions {
|
||||
projectStore: CodingProjectStore;
|
||||
executablePath: string;
|
||||
cliPath: string;
|
||||
userDataDir: string;
|
||||
bundledSkillsDir: string;
|
||||
extensionHost: PiManagedExtensionHost;
|
||||
loadProviderInput(): Promise<{ accounts: ProviderAccount[]; modelSummaries: ModelSummary[] }>;
|
||||
resolveCredential(account: ProviderAccount): Promise<string | null>;
|
||||
getRevision(): PiManagedInputRevision;
|
||||
getLocalProxyCredential?(): Promise<string | undefined>;
|
||||
createProcess?(options: PiWorkerProcessOptions): PiSubagentProcessAdapter;
|
||||
}
|
||||
|
||||
function recordValue(value: unknown): Record<string, unknown> | null {
|
||||
return value !== null && typeof value === 'object' && !Array.isArray(value)
|
||||
? value as Record<string, unknown>
|
||||
: null;
|
||||
}
|
||||
|
||||
function usageOf(value: unknown): PublicUsage | undefined {
|
||||
const usage = recordValue(value);
|
||||
if (!usage || typeof usage.input !== 'number' || typeof usage.output !== 'number') return undefined;
|
||||
return {
|
||||
inputTokens: usage.input,
|
||||
outputTokens: usage.output,
|
||||
...(typeof usage.cacheRead === 'number' ? { cacheReadTokens: usage.cacheRead } : {}),
|
||||
...(typeof usage.cacheWrite === 'number' ? { cacheWriteTokens: usage.cacheWrite } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
class ManagedPiSubagentChild implements PiSubagentChild {
|
||||
private ran = false;
|
||||
private stopped = false;
|
||||
|
||||
constructor(
|
||||
readonly id: string,
|
||||
private readonly process: PiSubagentProcessAdapter,
|
||||
private readonly disposeExtension: () => Promise<void>,
|
||||
) {}
|
||||
|
||||
async run(prompt: string, signal: AbortSignal): Promise<PiSubagentChildResult> {
|
||||
if (this.ran) throw new PiSubagentChildError('SUBAGENT_ALREADY_RAN');
|
||||
this.ran = true;
|
||||
if (signal.aborted) throw new PiSubagentChildError('SUBAGENT_ABORTED');
|
||||
let usage: PublicUsage | undefined;
|
||||
let settle!: () => void;
|
||||
let fail!: (error: PiSubagentChildError) => void;
|
||||
const settled = new Promise<void>((resolve, reject) => {
|
||||
settle = resolve;
|
||||
fail = reject;
|
||||
});
|
||||
const unsubscribeEvents = this.process.subscribe((event) => {
|
||||
if (event.type === 'message_end') {
|
||||
const message = recordValue(event.message);
|
||||
if (message?.role === 'assistant') usage = usageOf(message.usage) ?? usage;
|
||||
}
|
||||
if (event.type === 'agent_settled') settle();
|
||||
});
|
||||
const unsubscribeInvalidation = this.process.subscribeInvalidation(() => {
|
||||
fail(new PiSubagentChildError('SUBAGENT_CHILD_CRASHED'));
|
||||
});
|
||||
const abort = () => {
|
||||
void this.process.request({ type: 'abort' }).catch(() => undefined);
|
||||
fail(new PiSubagentChildError('SUBAGENT_ABORTED'));
|
||||
};
|
||||
signal.addEventListener('abort', abort, { once: true });
|
||||
try {
|
||||
const accepted = await this.process.request({ type: 'prompt', message: prompt });
|
||||
if (!accepted.success) throw new PiSubagentChildError('SUBAGENT_PROMPT_REJECTED');
|
||||
await settled;
|
||||
if (signal.aborted) throw new PiSubagentChildError('SUBAGENT_ABORTED');
|
||||
const response = await this.process.request<{ text?: unknown }>({
|
||||
type: 'get_last_assistant_text',
|
||||
}, { retry: 'read-only-once' });
|
||||
if (!response.success) throw new PiSubagentChildError('SUBAGENT_RESULT_UNAVAILABLE');
|
||||
const summary = typeof response.data?.text === 'string'
|
||||
? response.data.text
|
||||
: '';
|
||||
return {
|
||||
summary,
|
||||
...(usage ? { usage } : {}),
|
||||
};
|
||||
} finally {
|
||||
signal.removeEventListener('abort', abort);
|
||||
unsubscribeEvents();
|
||||
unsubscribeInvalidation();
|
||||
}
|
||||
}
|
||||
|
||||
async stop(): Promise<void> {
|
||||
if (this.stopped) return;
|
||||
this.stopped = true;
|
||||
try {
|
||||
await this.process.stop();
|
||||
} finally {
|
||||
await this.disposeExtension();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export function createPiManagedSubagentChildOpener(
|
||||
options: PiManagedSubagentChildOpenerOptions,
|
||||
): (input: PiSubagentChildOpenInput) => Promise<PiSubagentChild> {
|
||||
const createProcess = options.createProcess ?? ((processOptions) => new PiWorkerProcess(processOptions));
|
||||
return async (input) => {
|
||||
const project = (await options.projectStore.listProjects())
|
||||
.find((candidate) => candidate.id === input.projectId);
|
||||
if (!project) throw new PiSubagentChildError('SUBAGENT_PROJECT_UNAVAILABLE');
|
||||
const configRead = await readCodingProjectConfigV2(project.path);
|
||||
if (configRead.status !== 'valid') {
|
||||
throw new PiSubagentChildError('SUBAGENT_PROJECT_CONFIG_UNAVAILABLE');
|
||||
}
|
||||
const agent = configRead.config.agents.find((candidate) => (
|
||||
candidate.id === input.agentId && candidate.enabled && candidate.archivedAt === null
|
||||
));
|
||||
if (!agent) throw new PiSubagentChildError('SUBAGENT_AGENT_UNAVAILABLE');
|
||||
if (!agent.model || agent.modelResolution !== 'resolved') {
|
||||
throw new PiSubagentChildError('SUBAGENT_MODEL_REQUIRED');
|
||||
}
|
||||
|
||||
const providerInput = await options.loadProviderInput();
|
||||
const catalog = buildPiProviderCatalog(providerInput);
|
||||
const selection = selectPiProviderModel(catalog, agent.model);
|
||||
const account = providerInput.accounts.find((candidate) => candidate.id === selection.accountId);
|
||||
const descriptor = catalog.descriptors.find((candidate) => candidate.accountId === selection.accountId);
|
||||
if (!account || !descriptor) throw new PiSubagentChildError('SUBAGENT_MODEL_UNAVAILABLE');
|
||||
const managedPaths = await ensurePiManagedPaths(options.userDataDir);
|
||||
await writePiProviderCatalog(managedPaths.modelsFile, catalog, agent.model);
|
||||
const resources = await materializePiAgentResources({
|
||||
userDataDir: options.userDataDir,
|
||||
projectId: input.projectId,
|
||||
agentId: agent.id,
|
||||
prompt: agent.prompt,
|
||||
skillIds: agent.skillIds,
|
||||
bundledSkillsDir: options.bundledSkillsDir,
|
||||
revision: options.getRevision(),
|
||||
});
|
||||
const credential = await buildPiWorkerCredentialProjection({
|
||||
account,
|
||||
descriptor,
|
||||
resolveCredential: options.resolveCredential,
|
||||
...(options.getLocalProxyCredential
|
||||
? { localProxyCredential: await options.getLocalProxyCredential() }
|
||||
: {}),
|
||||
});
|
||||
const extension = await options.extensionHost.registerWorker({
|
||||
conversationId: input.conversationId,
|
||||
generation: input.workerGeneration,
|
||||
projectId: input.projectId,
|
||||
extensionsDir: managedPaths.extensionsDir,
|
||||
role: 'child',
|
||||
runId: input.runId,
|
||||
});
|
||||
const process = createProcess({
|
||||
executablePath: options.executablePath,
|
||||
cliPath: options.cliPath,
|
||||
cwd: project.path,
|
||||
configDir: resources.paths.configDir,
|
||||
sessionDir: resources.projectSessionsDir,
|
||||
tools: input.toolProfile === 'coding' ? CODING_TOOLS : READ_ONLY_TOOLS,
|
||||
additionalArgs: [
|
||||
...buildPiManagedInputArgs(selection, resources),
|
||||
'--extension', extension.extensionPath,
|
||||
'--no-session',
|
||||
],
|
||||
env: { ...credential.env, ...extension.env },
|
||||
sensitiveValues: [...credential.sensitiveValues, ...extension.sensitiveValues],
|
||||
});
|
||||
try {
|
||||
await process.start();
|
||||
const ready = await process.request({ type: 'get_state' }, { retry: 'read-only-once' });
|
||||
if (!ready.success) throw new PiSubagentChildError('SUBAGENT_CHILD_START_FAILED');
|
||||
return new ManagedPiSubagentChild(input.taskId, process, extension.dispose);
|
||||
} catch (error) {
|
||||
await process.stop().catch(() => undefined);
|
||||
await extension.dispose();
|
||||
throw error;
|
||||
}
|
||||
};
|
||||
}
|
||||
428
electron/coding-runtime/pi/subagent.ts
Normal file
428
electron/coding-runtime/pi/subagent.ts
Normal file
@@ -0,0 +1,428 @@
|
||||
import { randomUUID } from 'node:crypto';
|
||||
|
||||
import type { PublicUsage, SubagentDetailsV1 } from '../contracts';
|
||||
import { PiProcessBudget, type PiProcessLease } from './worker-pool';
|
||||
|
||||
export type PiSubagentMode = SubagentDetailsV1['mode'];
|
||||
export type PiSubagentToolProfile = SubagentDetailsV1['tasks'][number]['toolProfile'];
|
||||
|
||||
export interface PiSubagentTaskRequest {
|
||||
agentId: string;
|
||||
task: string;
|
||||
toolProfile: PiSubagentToolProfile;
|
||||
}
|
||||
|
||||
export interface PiSubagentDispatchRequest {
|
||||
mode: PiSubagentMode;
|
||||
tasks: PiSubagentTaskRequest[];
|
||||
}
|
||||
|
||||
export interface PiSubagentParentIdentity {
|
||||
conversationId: string;
|
||||
workerGeneration: number;
|
||||
runId: string;
|
||||
projectId: string;
|
||||
}
|
||||
|
||||
export interface PiSubagentChildOpenInput extends PiSubagentParentIdentity {
|
||||
dispatchId: string;
|
||||
taskId: string;
|
||||
agentId: string;
|
||||
toolProfile: PiSubagentToolProfile;
|
||||
}
|
||||
|
||||
export interface PiSubagentChildResult {
|
||||
summary: string;
|
||||
usage?: PublicUsage;
|
||||
}
|
||||
|
||||
export interface PiSubagentChild {
|
||||
readonly id: string;
|
||||
run(
|
||||
prompt: string,
|
||||
signal: AbortSignal,
|
||||
onProgress?: (summary: string) => void,
|
||||
): Promise<PiSubagentChildResult>;
|
||||
stop(): Promise<void>;
|
||||
}
|
||||
|
||||
export interface PiSubagentDispatchResult {
|
||||
details: SubagentDetailsV1;
|
||||
}
|
||||
|
||||
export interface PiSubagentDispatchOptions {
|
||||
signal?: AbortSignal;
|
||||
onUpdate?: (details: SubagentDetailsV1) => void;
|
||||
}
|
||||
|
||||
export interface PiSubagentSchedulerOptions {
|
||||
openChild(input: PiSubagentChildOpenInput): Promise<PiSubagentChild>;
|
||||
processBudget: PiProcessBudget;
|
||||
reclaimProcessCapacity?(): Promise<boolean>;
|
||||
createId?: (kind: 'dispatch' | 'task') => string;
|
||||
}
|
||||
|
||||
interface DispatchRecord {
|
||||
identity: PiSubagentParentIdentity;
|
||||
controller: AbortController;
|
||||
children: Set<PiSubagentChild>;
|
||||
flight: Promise<PiSubagentDispatchResult>;
|
||||
}
|
||||
|
||||
interface SemaphoreWaiter {
|
||||
signal: AbortSignal;
|
||||
resolve(release: () => void): void;
|
||||
reject(error: Error): void;
|
||||
abort(): void;
|
||||
}
|
||||
|
||||
const MAX_TASKS_PER_DISPATCH = 8;
|
||||
const MAX_AGENT_ID_LENGTH = 128;
|
||||
const MAX_TASK_LENGTH = 6_000;
|
||||
const MAX_SUMMARY_LENGTH = 4_000;
|
||||
|
||||
export class PiSubagentChildError extends Error {
|
||||
constructor(readonly code: string) {
|
||||
super(code);
|
||||
this.name = 'PiSubagentChildError';
|
||||
}
|
||||
}
|
||||
|
||||
class FifoSemaphore {
|
||||
private readonly waiters: SemaphoreWaiter[] = [];
|
||||
private active = 0;
|
||||
|
||||
constructor(readonly maximum: number) {
|
||||
if (!Number.isSafeInteger(maximum) || maximum <= 0) {
|
||||
throw new Error('Subagent concurrency must be a positive safe integer');
|
||||
}
|
||||
}
|
||||
|
||||
acquire(signal: AbortSignal): Promise<() => void> {
|
||||
if (signal.aborted) return Promise.reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
|
||||
if (this.active < this.maximum) return Promise.resolve(this.issuePermit());
|
||||
return new Promise<() => void>((resolve, reject) => {
|
||||
const waiter: SemaphoreWaiter = {
|
||||
signal,
|
||||
resolve,
|
||||
reject,
|
||||
abort: () => {
|
||||
const index = this.waiters.indexOf(waiter);
|
||||
if (index >= 0) this.waiters.splice(index, 1);
|
||||
reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
|
||||
},
|
||||
};
|
||||
signal.addEventListener('abort', waiter.abort, { once: true });
|
||||
this.waiters.push(waiter);
|
||||
});
|
||||
}
|
||||
|
||||
private issuePermit(): () => void {
|
||||
this.active += 1;
|
||||
let released = false;
|
||||
return () => {
|
||||
if (released) return;
|
||||
released = true;
|
||||
this.active -= 1;
|
||||
this.advance();
|
||||
};
|
||||
}
|
||||
|
||||
private advance(): void {
|
||||
while (this.active < this.maximum && this.waiters.length > 0) {
|
||||
const waiter = this.waiters.shift() as SemaphoreWaiter;
|
||||
waiter.signal.removeEventListener('abort', waiter.abort);
|
||||
if (waiter.signal.aborted) {
|
||||
waiter.reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
|
||||
continue;
|
||||
}
|
||||
waiter.resolve(this.issuePermit());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function asRecord(value: unknown): Record<string, unknown> | null {
|
||||
return value !== null && typeof value === 'object' && !Array.isArray(value)
|
||||
? value as Record<string, unknown>
|
||||
: null;
|
||||
}
|
||||
|
||||
function boundedText(value: unknown, label: string, maximum: number): string {
|
||||
if (typeof value !== 'string') throw new Error(`${label} must be a string`);
|
||||
const normalized = value.trim();
|
||||
if (!normalized) throw new Error(`${label} must not be empty`);
|
||||
if (normalized.length > maximum) throw new Error(`${label} is too long`);
|
||||
return normalized;
|
||||
}
|
||||
|
||||
export function parsePiSubagentDispatchRequest(value: unknown): PiSubagentDispatchRequest {
|
||||
const record = asRecord(value);
|
||||
if (!record || !['single', 'parallel', 'chain'].includes(String(record.mode))) {
|
||||
throw new Error('Subagent dispatch mode is invalid');
|
||||
}
|
||||
if (!Array.isArray(record.tasks) || record.tasks.length === 0) {
|
||||
throw new Error('Subagent dispatch requires at least one task');
|
||||
}
|
||||
if (record.tasks.length > MAX_TASKS_PER_DISPATCH) {
|
||||
throw new Error('Subagent dispatch accepts at most 8 tasks');
|
||||
}
|
||||
if (record.mode === 'single' && record.tasks.length !== 1) {
|
||||
throw new Error('Single subagent dispatch requires exactly one task');
|
||||
}
|
||||
const tasks = record.tasks.map((candidate, index) => {
|
||||
const task = asRecord(candidate);
|
||||
if (!task) throw new Error(`Subagent task ${index + 1} is invalid`);
|
||||
if (task.toolProfile !== 'read-only' && task.toolProfile !== 'coding') {
|
||||
throw new Error(`Subagent task ${index + 1} tool profile is invalid`);
|
||||
}
|
||||
return {
|
||||
agentId: boundedText(task.agentId, `Subagent task ${index + 1} agent id`, MAX_AGENT_ID_LENGTH),
|
||||
task: boundedText(task.task, `Subagent task ${index + 1} prompt`, MAX_TASK_LENGTH),
|
||||
toolProfile: task.toolProfile,
|
||||
};
|
||||
});
|
||||
return { mode: record.mode as PiSubagentMode, tasks };
|
||||
}
|
||||
|
||||
function safeSummary(value: string): string {
|
||||
return value.length <= MAX_SUMMARY_LENGTH ? value : `${value.slice(0, MAX_SUMMARY_LENGTH - 1)}…`;
|
||||
}
|
||||
|
||||
function safeUsage(value: PublicUsage | undefined): PublicUsage | undefined {
|
||||
if (!value) return undefined;
|
||||
const fields = [
|
||||
value.inputTokens,
|
||||
value.outputTokens,
|
||||
value.cacheReadTokens,
|
||||
value.cacheWriteTokens,
|
||||
].filter((field) => field !== undefined);
|
||||
if (fields.some((field) => !Number.isFinite(field) || field < 0)) return undefined;
|
||||
return structuredClone(value);
|
||||
}
|
||||
|
||||
function publicErrorCode(error: unknown, aborted: boolean): string {
|
||||
if (aborted) return 'SUBAGENT_ABORTED';
|
||||
if (error instanceof PiSubagentChildError && /^[A-Z][A-Z0-9_]{0,63}$/.test(error.code)) {
|
||||
return error.code;
|
||||
}
|
||||
return 'SUBAGENT_CHILD_FAILED';
|
||||
}
|
||||
|
||||
function parentKey(identity: PiSubagentParentIdentity): string {
|
||||
return `${identity.conversationId}:${identity.workerGeneration}:${identity.runId}`;
|
||||
}
|
||||
|
||||
export class PiSubagentScheduler {
|
||||
private readonly openChild: PiSubagentSchedulerOptions['openChild'];
|
||||
private readonly processBudget: PiProcessBudget;
|
||||
private readonly childPermits: FifoSemaphore;
|
||||
private readonly reclaimProcessCapacity: (() => Promise<boolean>) | undefined;
|
||||
private readonly createId: NonNullable<PiSubagentSchedulerOptions['createId']>;
|
||||
private readonly dispatches = new Map<string, DispatchRecord>();
|
||||
private readonly parentDispatches = new Map<string, Set<string>>();
|
||||
private closing = false;
|
||||
|
||||
constructor(options: PiSubagentSchedulerOptions) {
|
||||
this.openChild = options.openChild;
|
||||
this.processBudget = options.processBudget;
|
||||
this.childPermits = new FifoSemaphore(4);
|
||||
this.reclaimProcessCapacity = options.reclaimProcessCapacity;
|
||||
this.createId = options.createId ?? ((kind) => `${kind}-${randomUUID()}`);
|
||||
}
|
||||
|
||||
dispatch(
|
||||
input: PiSubagentParentIdentity & { request: unknown },
|
||||
options: PiSubagentDispatchOptions = {},
|
||||
): Promise<PiSubagentDispatchResult> {
|
||||
if (this.closing) return Promise.reject(new Error('Subagent scheduler is shutting down'));
|
||||
const request = parsePiSubagentDispatchRequest(input.request);
|
||||
const dispatchId = this.createId('dispatch');
|
||||
if (this.dispatches.has(dispatchId)) {
|
||||
return Promise.reject(new Error(`Duplicate subagent dispatch id: ${dispatchId}`));
|
||||
}
|
||||
const identity: PiSubagentParentIdentity = {
|
||||
conversationId: input.conversationId,
|
||||
workerGeneration: input.workerGeneration,
|
||||
runId: input.runId,
|
||||
projectId: input.projectId,
|
||||
};
|
||||
const controller = new AbortController();
|
||||
const externalAbort = () => controller.abort();
|
||||
options.signal?.addEventListener('abort', externalAbort, { once: true });
|
||||
if (options.signal?.aborted) controller.abort();
|
||||
const tasks: SubagentDetailsV1['tasks'] = request.tasks.map((task) => ({
|
||||
taskId: this.createId('task'),
|
||||
agentId: task.agentId,
|
||||
toolProfile: task.toolProfile,
|
||||
status: 'queued',
|
||||
}));
|
||||
const details: SubagentDetailsV1 = {
|
||||
schema: 'subagent.v1',
|
||||
dispatchId,
|
||||
mode: request.mode,
|
||||
tasks,
|
||||
};
|
||||
const record = {
|
||||
identity,
|
||||
controller,
|
||||
children: new Set<PiSubagentChild>(),
|
||||
} as DispatchRecord;
|
||||
const flight = this.runDispatch(record, request, details, options.onUpdate)
|
||||
.finally(() => {
|
||||
options.signal?.removeEventListener('abort', externalAbort);
|
||||
this.dispatches.delete(dispatchId);
|
||||
const key = parentKey(identity);
|
||||
const ids = this.parentDispatches.get(key);
|
||||
ids?.delete(dispatchId);
|
||||
if (ids?.size === 0) this.parentDispatches.delete(key);
|
||||
});
|
||||
record.flight = flight;
|
||||
this.dispatches.set(dispatchId, record);
|
||||
const key = parentKey(identity);
|
||||
const ids = this.parentDispatches.get(key) ?? new Set<string>();
|
||||
ids.add(dispatchId);
|
||||
this.parentDispatches.set(key, ids);
|
||||
this.emit(details, options.onUpdate);
|
||||
return flight;
|
||||
}
|
||||
|
||||
abortParent(identity: PiSubagentParentIdentity): void {
|
||||
for (const dispatchId of this.parentDispatches.get(parentKey(identity)) ?? []) {
|
||||
this.dispatches.get(dispatchId)?.controller.abort();
|
||||
}
|
||||
}
|
||||
|
||||
async close(): Promise<void> {
|
||||
if (this.closing) {
|
||||
await Promise.allSettled([...this.dispatches.values()].map((record) => record.flight));
|
||||
return;
|
||||
}
|
||||
this.closing = true;
|
||||
for (const record of this.dispatches.values()) record.controller.abort();
|
||||
await Promise.allSettled([...this.dispatches.values()].map((record) => record.flight));
|
||||
}
|
||||
|
||||
private async runDispatch(
|
||||
record: DispatchRecord,
|
||||
request: PiSubagentDispatchRequest,
|
||||
details: SubagentDetailsV1,
|
||||
onUpdate: PiSubagentDispatchOptions['onUpdate'],
|
||||
): Promise<PiSubagentDispatchResult> {
|
||||
if (request.mode === 'chain') {
|
||||
let previous = '';
|
||||
for (let index = 0; index < request.tasks.length; index += 1) {
|
||||
const task = request.tasks[index] as PiSubagentTaskRequest;
|
||||
if (record.controller.signal.aborted) {
|
||||
this.markRemaining(details, index, 'aborted', onUpdate);
|
||||
break;
|
||||
}
|
||||
const prompt = task.task.split('{previous}').join(previous);
|
||||
await this.runTask(record, task, details, index, prompt, onUpdate);
|
||||
const projected = details.tasks[index];
|
||||
if (projected?.status !== 'complete') {
|
||||
this.markRemaining(
|
||||
details,
|
||||
index + 1,
|
||||
record.controller.signal.aborted ? 'aborted' : 'skipped',
|
||||
onUpdate,
|
||||
);
|
||||
break;
|
||||
}
|
||||
previous = projected.summary ?? '';
|
||||
}
|
||||
} else {
|
||||
await Promise.all(request.tasks.map((task, index) => (
|
||||
this.runTask(record, task, details, index, task.task, onUpdate)
|
||||
)));
|
||||
}
|
||||
return { details: structuredClone(details) };
|
||||
}
|
||||
|
||||
private async runTask(
|
||||
record: DispatchRecord,
|
||||
task: PiSubagentTaskRequest,
|
||||
details: SubagentDetailsV1,
|
||||
index: number,
|
||||
prompt: string,
|
||||
onUpdate: PiSubagentDispatchOptions['onUpdate'],
|
||||
): Promise<void> {
|
||||
const projected = details.tasks[index];
|
||||
if (!projected) return;
|
||||
let releaseChild: (() => void) | undefined;
|
||||
let processLease: PiProcessLease | undefined;
|
||||
let child: PiSubagentChild | undefined;
|
||||
try {
|
||||
releaseChild = await this.childPermits.acquire(record.controller.signal);
|
||||
await this.reclaimIdleCapacity(record.controller.signal);
|
||||
processLease = await this.processBudget.acquire(record.controller.signal);
|
||||
if (record.controller.signal.aborted) throw new PiSubagentChildError('SUBAGENT_ABORTED');
|
||||
projected.status = 'running';
|
||||
this.emit(details, onUpdate);
|
||||
child = await this.openChild({
|
||||
...record.identity,
|
||||
dispatchId: details.dispatchId,
|
||||
taskId: projected.taskId,
|
||||
agentId: task.agentId,
|
||||
toolProfile: task.toolProfile,
|
||||
});
|
||||
record.children.add(child);
|
||||
const result = await child.run(prompt, record.controller.signal, (summary) => {
|
||||
projected.summary = safeSummary(summary);
|
||||
this.emit(details, onUpdate);
|
||||
});
|
||||
if (record.controller.signal.aborted) throw new PiSubagentChildError('SUBAGENT_ABORTED');
|
||||
projected.status = 'complete';
|
||||
projected.summary = safeSummary(result.summary);
|
||||
const usage = safeUsage(result.usage);
|
||||
if (usage) projected.usage = usage;
|
||||
} catch (error) {
|
||||
const aborted = record.controller.signal.aborted
|
||||
|| (error instanceof PiSubagentChildError && error.code === 'SUBAGENT_ABORTED');
|
||||
projected.status = aborted ? 'aborted' : 'error';
|
||||
projected.errorCode = publicErrorCode(error, aborted);
|
||||
} finally {
|
||||
if (child) {
|
||||
record.children.delete(child);
|
||||
await child.stop().catch(() => undefined);
|
||||
}
|
||||
processLease?.release();
|
||||
releaseChild?.();
|
||||
this.emit(details, onUpdate);
|
||||
}
|
||||
}
|
||||
|
||||
private async reclaimIdleCapacity(signal: AbortSignal): Promise<void> {
|
||||
while (!signal.aborted
|
||||
&& this.processBudget.activeCount >= this.processBudget.maxProcesses
|
||||
&& this.reclaimProcessCapacity) {
|
||||
if (!await this.reclaimProcessCapacity()) break;
|
||||
}
|
||||
}
|
||||
|
||||
private markRemaining(
|
||||
details: SubagentDetailsV1,
|
||||
from: number,
|
||||
status: 'aborted' | 'skipped',
|
||||
onUpdate: PiSubagentDispatchOptions['onUpdate'],
|
||||
): void {
|
||||
for (let index = from; index < details.tasks.length; index += 1) {
|
||||
const task = details.tasks[index];
|
||||
if (!task || task.status !== 'queued') continue;
|
||||
task.status = status;
|
||||
if (status === 'aborted') task.errorCode = 'SUBAGENT_ABORTED';
|
||||
}
|
||||
this.emit(details, onUpdate);
|
||||
}
|
||||
|
||||
private emit(
|
||||
details: SubagentDetailsV1,
|
||||
onUpdate: PiSubagentDispatchOptions['onUpdate'],
|
||||
): void {
|
||||
if (!onUpdate) return;
|
||||
try {
|
||||
onUpdate(structuredClone(details));
|
||||
} catch {
|
||||
// UI projection observers must not affect child lifecycle.
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -333,6 +333,15 @@ export class PiWorkerPool {
|
||||
return record ? this.publicState(record) : null;
|
||||
}
|
||||
|
||||
async reclaimIdleWorker(): Promise<boolean> {
|
||||
const record = [...this.workers.values()]
|
||||
.filter((candidate) => candidate.state === 'ready' || candidate.state === 'idle')
|
||||
.sort((left, right) => left.lastUsed - right.lastUsed)[0];
|
||||
if (!record) return false;
|
||||
await this.evict(record);
|
||||
return true;
|
||||
}
|
||||
|
||||
subscribe(listener: (event: PiWorkerPoolEvent) => void): () => void {
|
||||
this.listeners.add(listener);
|
||||
return () => this.listeners.delete(listener);
|
||||
|
||||
@@ -52,6 +52,7 @@ export type PiWorkerProcessOptions = {
|
||||
cwd: string;
|
||||
configDir: string;
|
||||
sessionDir: string;
|
||||
tools?: readonly string[];
|
||||
additionalArgs?: readonly string[];
|
||||
env?: NodeJS.ProcessEnv;
|
||||
sensitiveValues?: readonly string[];
|
||||
@@ -64,6 +65,7 @@ export type PiWorkerProcessOptions = {
|
||||
export function buildPiRpcArgs(
|
||||
sessionDir: string,
|
||||
additionalArgs: readonly string[] = [],
|
||||
tools: readonly string[] = ['read', 'bash', 'edit', 'write', 'grep', 'find', 'ls', 'ask_user'],
|
||||
): string[] {
|
||||
return [
|
||||
'--mode', 'rpc',
|
||||
@@ -75,7 +77,7 @@ export function buildPiRpcArgs(
|
||||
'--no-themes',
|
||||
'--no-context-files',
|
||||
'--no-approve',
|
||||
'--tools', 'read,bash,edit,write,grep,find,ls,ask_user',
|
||||
'--tools', tools.join(','),
|
||||
...additionalArgs,
|
||||
];
|
||||
}
|
||||
@@ -217,7 +219,11 @@ export class PiWorkerProcess {
|
||||
async start(): Promise<this> {
|
||||
if (this.child) throw new Error('Pi worker process already started');
|
||||
const generation = this.generationValue;
|
||||
const args = buildPiRpcArgs(this.options.sessionDir, this.options.additionalArgs);
|
||||
const args = buildPiRpcArgs(
|
||||
this.options.sessionDir,
|
||||
this.options.additionalArgs,
|
||||
this.options.tools,
|
||||
);
|
||||
assertSensitiveValuesAbsentFromArgs(args, this.options.sensitiveValues);
|
||||
const child = spawn(
|
||||
this.options.executablePath,
|
||||
|
||||
Reference in New Issue
Block a user