feat(coding-runtime): add managed Pi extension host
This commit is contained in:
299
electron/coding-runtime/pi/extension-host.ts
Normal file
299
electron/coding-runtime/pi/extension-host.ts
Normal file
@@ -0,0 +1,299 @@
|
||||
import { randomBytes, randomUUID } from 'node:crypto';
|
||||
import { createServer, type IncomingMessage, type Server, type ServerResponse } from 'node:http';
|
||||
import { mkdir, rm } from 'node:fs/promises';
|
||||
import path from 'node:path';
|
||||
import { atomicWriteJson } from '../../coding-projects/atomic-json';
|
||||
import { materializeMakelorePiExtension } from './extensions/makelore-runtime';
|
||||
import {
|
||||
PiProjectWriteLeaseCoordinator,
|
||||
type PiProjectWriteLease,
|
||||
} from './write-lease';
|
||||
|
||||
const MAX_REQUEST_BYTES = 64 * 1024;
|
||||
|
||||
interface WorkerRegistrationRecord {
|
||||
token: string;
|
||||
conversationId: string;
|
||||
generation: number;
|
||||
projectId: string;
|
||||
contextFile: string;
|
||||
runId: string | null;
|
||||
leases: Map<string, PiProjectWriteLease>;
|
||||
waiters: Map<string, AbortController>;
|
||||
}
|
||||
|
||||
export interface PiExtensionWorkerRegistration {
|
||||
extensionPath: string;
|
||||
env: NodeJS.ProcessEnv;
|
||||
sensitiveValues: string[];
|
||||
dispose(): Promise<void>;
|
||||
}
|
||||
|
||||
export interface RegisterPiExtensionWorkerInput {
|
||||
conversationId: string;
|
||||
generation: number;
|
||||
projectId: string;
|
||||
extensionsDir: string;
|
||||
}
|
||||
|
||||
interface BridgeRequest {
|
||||
action: 'lease.acquire' | 'lease.release';
|
||||
conversationId: string;
|
||||
workerGeneration: number;
|
||||
runId: string;
|
||||
resourceId: string;
|
||||
leaseId?: string;
|
||||
}
|
||||
|
||||
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'
|
||||
&& Number.isSafeInteger(value.workerGeneration)
|
||||
&& typeof value.runId === 'string'
|
||||
&& typeof value.resourceId === 'string'
|
||||
&& (value.leaseId === undefined || typeof value.leaseId === 'string');
|
||||
}
|
||||
|
||||
export class PiManagedExtensionHost {
|
||||
private readonly leases: PiProjectWriteLeaseCoordinator;
|
||||
private readonly registrations = new Map<string, WorkerRegistrationRecord>();
|
||||
private readonly runBindings = new Map<string, string>();
|
||||
private server: Server | null = null;
|
||||
private bridgeUrl: string | null = null;
|
||||
private startFlight: Promise<string> | null = null;
|
||||
|
||||
constructor(leases = new PiProjectWriteLeaseCoordinator()) {
|
||||
this.leases = leases;
|
||||
}
|
||||
|
||||
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');
|
||||
}
|
||||
const bridgeUrl = await this.start();
|
||||
await mkdir(input.extensionsDir, { recursive: true });
|
||||
const extensionPath = await materializeMakelorePiExtension(input.extensionsDir);
|
||||
const token = randomBytes(32).toString('base64url');
|
||||
const contextFile = path.join(input.extensionsDir, `worker-${randomUUID()}.json`);
|
||||
const record: WorkerRegistrationRecord = {
|
||||
token,
|
||||
conversationId: input.conversationId,
|
||||
generation: input.generation,
|
||||
projectId: input.projectId,
|
||||
contextFile,
|
||||
runId: this.runBindings.get(input.conversationId) ?? null,
|
||||
leases: new Map(),
|
||||
waiters: new Map(),
|
||||
};
|
||||
this.registrations.set(token, record);
|
||||
await this.writeContext(record);
|
||||
let disposed = false;
|
||||
return {
|
||||
extensionPath,
|
||||
env: {
|
||||
MAKELORE_PI_BRIDGE_URL: bridgeUrl,
|
||||
MAKELORE_PI_WORKER_TOKEN: token,
|
||||
MAKELORE_PI_CONTEXT_FILE: contextFile,
|
||||
},
|
||||
sensitiveValues: [token],
|
||||
dispose: async () => {
|
||||
if (disposed) return;
|
||||
disposed = true;
|
||||
this.disposeRecord(record);
|
||||
await rm(contextFile, { force: true });
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async bindRun(conversationId: string, generation: number, runId: string): Promise<void> {
|
||||
const record = this.findWorker(conversationId, generation);
|
||||
if (!record) throw new Error('Pi extension worker registration is unavailable');
|
||||
this.releaseWorkerResources(record);
|
||||
this.runBindings.set(conversationId, runId);
|
||||
record.runId = runId;
|
||||
await this.writeContext(record);
|
||||
}
|
||||
|
||||
async clearRun(conversationId: string, generation: number, runId?: string): Promise<void> {
|
||||
const record = this.findWorker(conversationId, generation);
|
||||
if (!runId || this.runBindings.get(conversationId) === runId) {
|
||||
this.runBindings.delete(conversationId);
|
||||
}
|
||||
if (!record || (runId && record.runId !== runId)) return;
|
||||
this.releaseWorkerResources(record);
|
||||
record.runId = null;
|
||||
await this.writeContext(record);
|
||||
}
|
||||
|
||||
async close(): Promise<void> {
|
||||
for (const record of [...this.registrations.values()]) this.disposeRecord(record);
|
||||
this.runBindings.clear();
|
||||
const server = this.server;
|
||||
this.server = null;
|
||||
this.bridgeUrl = null;
|
||||
this.startFlight = null;
|
||||
if (!server) return;
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
server.close((error) => error ? reject(error) : resolve());
|
||||
});
|
||||
}
|
||||
|
||||
private start(): Promise<string> {
|
||||
if (this.bridgeUrl) return Promise.resolve(this.bridgeUrl);
|
||||
if (this.startFlight) return this.startFlight;
|
||||
this.startFlight = new Promise<string>((resolve, reject) => {
|
||||
const server = createServer((request, response) => {
|
||||
void this.handle(request, response);
|
||||
});
|
||||
server.once('error', reject);
|
||||
server.listen(0, '127.0.0.1', () => {
|
||||
server.removeListener('error', reject);
|
||||
const address = server.address();
|
||||
if (!address || typeof address === 'string') {
|
||||
reject(new Error('Pi extension bridge did not bind a loopback port'));
|
||||
return;
|
||||
}
|
||||
this.server = server;
|
||||
this.bridgeUrl = `http://127.0.0.1:${address.port}/v1/worker`;
|
||||
resolve(this.bridgeUrl);
|
||||
});
|
||||
}).finally(() => {
|
||||
if (!this.bridgeUrl) this.startFlight = null;
|
||||
});
|
||||
return this.startFlight;
|
||||
}
|
||||
|
||||
private async handle(request: IncomingMessage, response: ServerResponse): Promise<void> {
|
||||
if (request.method !== 'POST' || request.url !== '/v1/worker') {
|
||||
this.respond(response, 404, { error: 'Not found' });
|
||||
return;
|
||||
}
|
||||
const authorization = request.headers.authorization;
|
||||
const token = authorization?.startsWith('Bearer ') ? authorization.slice(7) : '';
|
||||
const record = this.registrations.get(token);
|
||||
if (!record) {
|
||||
this.respond(response, 401, { error: 'Worker token is invalid' });
|
||||
return;
|
||||
}
|
||||
try {
|
||||
const value = await this.readBody(request);
|
||||
if (!bridgeRequest(value)) {
|
||||
this.respond(response, 400, { error: 'Bridge request is invalid' });
|
||||
return;
|
||||
}
|
||||
if (value.conversationId !== record.conversationId
|
||||
|| value.workerGeneration !== record.generation
|
||||
|| value.runId !== record.runId) {
|
||||
this.respond(response, 409, { error: 'Worker run identity is stale' });
|
||||
return;
|
||||
}
|
||||
if (!value.resourceId.trim()) {
|
||||
this.respond(response, 400, { error: 'Bridge resource id is required' });
|
||||
return;
|
||||
}
|
||||
if (value.action === 'lease.release') {
|
||||
const lease = record.leases.get(value.resourceId);
|
||||
if (!lease || !value.leaseId || lease.id !== value.leaseId) {
|
||||
this.respond(response, 409, { error: 'Project write lease is stale' });
|
||||
return;
|
||||
}
|
||||
record.leases.delete(value.resourceId);
|
||||
lease.release();
|
||||
this.respond(response, 200, { released: true });
|
||||
return;
|
||||
}
|
||||
if (record.leases.has(value.resourceId) || record.waiters.has(value.resourceId)) {
|
||||
this.respond(response, 409, { error: 'Project write lease resource already exists' });
|
||||
return;
|
||||
}
|
||||
const controller = new AbortController();
|
||||
record.waiters.set(value.resourceId, controller);
|
||||
const cancel = () => {
|
||||
if (!response.writableEnded) controller.abort();
|
||||
};
|
||||
request.once('aborted', cancel);
|
||||
response.once('close', cancel);
|
||||
try {
|
||||
const lease = await this.leases.acquire(
|
||||
record.projectId,
|
||||
`${record.conversationId}:${record.generation}:${value.resourceId}`,
|
||||
controller.signal,
|
||||
);
|
||||
if (record.runId !== value.runId || this.registrations.get(token) !== record) {
|
||||
lease.release();
|
||||
return;
|
||||
}
|
||||
record.leases.set(value.resourceId, lease);
|
||||
this.respond(response, 200, { leaseId: lease.id });
|
||||
} finally {
|
||||
request.removeListener('aborted', cancel);
|
||||
response.removeListener('close', cancel);
|
||||
record.waiters.delete(value.resourceId);
|
||||
}
|
||||
} catch {
|
||||
if (!response.writableEnded) this.respond(response, 400, { error: 'Bridge request failed' });
|
||||
}
|
||||
}
|
||||
|
||||
private readBody(request: IncomingMessage): Promise<unknown> {
|
||||
return new Promise((resolve, reject) => {
|
||||
const chunks: Buffer[] = [];
|
||||
let bytes = 0;
|
||||
request.on('data', (chunk: Buffer) => {
|
||||
bytes += chunk.length;
|
||||
if (bytes > MAX_REQUEST_BYTES) {
|
||||
reject(new Error('Bridge request is too large'));
|
||||
request.destroy();
|
||||
return;
|
||||
}
|
||||
chunks.push(chunk);
|
||||
});
|
||||
request.once('end', () => {
|
||||
try {
|
||||
resolve(JSON.parse(Buffer.concat(chunks).toString('utf8')));
|
||||
} catch (error) {
|
||||
reject(error);
|
||||
}
|
||||
});
|
||||
request.once('error', reject);
|
||||
});
|
||||
}
|
||||
|
||||
private respond(response: ServerResponse, status: number, body: Record<string, unknown>): void {
|
||||
if (response.writableEnded) return;
|
||||
response.writeHead(status, { 'content-type': 'application/json; charset=utf-8' });
|
||||
response.end(JSON.stringify(body));
|
||||
}
|
||||
|
||||
private findWorker(conversationId: string, generation: number): WorkerRegistrationRecord | undefined {
|
||||
return [...this.registrations.values()].find((record) => (
|
||||
record.conversationId === conversationId && record.generation === generation
|
||||
));
|
||||
}
|
||||
|
||||
private async writeContext(record: WorkerRegistrationRecord): Promise<void> {
|
||||
await atomicWriteJson(record.contextFile, {
|
||||
conversationId: record.conversationId,
|
||||
workerGeneration: record.generation,
|
||||
...(record.runId ? { runId: record.runId } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
private disposeRecord(record: WorkerRegistrationRecord): void {
|
||||
if (this.registrations.get(record.token) !== record) return;
|
||||
this.registrations.delete(record.token);
|
||||
this.releaseWorkerResources(record);
|
||||
}
|
||||
|
||||
private releaseWorkerResources(record: WorkerRegistrationRecord): void {
|
||||
for (const controller of record.waiters.values()) controller.abort();
|
||||
record.waiters.clear();
|
||||
for (const lease of record.leases.values()) lease.release();
|
||||
record.leases.clear();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user