feat: add Pi subagent scheduler and child runtime
This commit is contained in:
@@ -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 } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user