feat(coding-runtime): add managed Pi extension host

This commit is contained in:
2026-08-23 09:39:12 +08:00
parent 47159b3cbd
commit 3861c3286a
22 changed files with 1673 additions and 16 deletions

View 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();
}
}

View File

@@ -0,0 +1,162 @@
import type { PiRpcEvent } from './rpc-client';
const MAX_MESSAGE_LENGTH = 2_000;
const MAX_TITLE_LENGTH = 256;
const MAX_EDITOR_TEXT_LENGTH = 64 * 1024;
const MAX_WIDGET_LINES = 32;
const MAX_WIDGET_LINE_LENGTH = 512;
export type PiExtensionUiProjection =
| { kind: 'notify'; conversationId: string; message: string; level: 'info' | 'warning' | 'error' }
| { kind: 'status'; conversationId: string; key: string; text?: string }
| {
kind: 'widget';
conversationId: string;
key: string;
lines?: string[];
placement?: 'aboveEditor' | 'belowEditor';
}
| { kind: 'title'; conversationId: string; title: string }
| { kind: 'editor-text'; conversationId: string; text: string; draftRevision: number };
export interface PiExtensionUiDiagnostic {
method: string;
reason: 'unsupported-ui-method' | 'invalid-ui-payload' | 'stale-draft-revision';
}
export interface PiExtensionUiProjectorOptions {
getDraftRevision(conversationId: string): number;
knownStatusKeys?: readonly string[];
knownWidgetKeys?: readonly string[];
}
interface RunProjectionState {
runId: string;
generation: number;
draftRevision: number;
}
function bounded(value: string, length: number): string {
return value.length <= length ? value : `${value.slice(0, length - 1)}…`;
}
export class PiExtensionUiProjector {
private readonly getDraftRevision: (conversationId: string) => number;
private readonly knownStatusKeys: Set<string>;
private readonly knownWidgetKeys: Set<string>;
private readonly runs = new Map<string, RunProjectionState>();
private readonly diagnostics: PiExtensionUiDiagnostic[] = [];
constructor(options: PiExtensionUiProjectorOptions) {
this.getDraftRevision = options.getDraftRevision;
this.knownStatusKeys = new Set(options.knownStatusKeys ?? ['makelore.write-lease']);
this.knownWidgetKeys = new Set(options.knownWidgetKeys ?? []);
}
beginRun(conversationId: string, generation: number, runId: string): void {
this.runs.set(conversationId, {
generation,
runId,
draftRevision: this.getDraftRevision(conversationId),
});
}
replaceGeneration(conversationId: string, generation: number, runId: string): void {
const current = this.runs.get(conversationId);
if (current?.runId === runId) current.generation = generation;
}
endRun(conversationId: string, runId: string): void {
if (this.runs.get(conversationId)?.runId === runId) this.runs.delete(conversationId);
}
getDiagnostics(): PiExtensionUiDiagnostic[] {
return structuredClone(this.diagnostics);
}
project(
conversationId: string,
generation: number,
runId: string,
event: PiRpcEvent,
): PiExtensionUiProjection | null {
if (event.type !== 'extension_ui_request' || typeof event.method !== 'string') return null;
const run = this.runs.get(conversationId);
if (!run || run.runId !== runId || run.generation !== generation) return null;
if (event.method === 'notify'
&& typeof event.message === 'string'
&& ['info', 'warning', 'error', undefined].includes(event.notifyType as string | undefined)) {
return {
kind: 'notify',
conversationId,
message: bounded(event.message, MAX_MESSAGE_LENGTH),
level: event.notifyType === 'warning' || event.notifyType === 'error'
? event.notifyType
: 'info',
};
}
if (event.method === 'setStatus'
&& typeof event.statusKey === 'string'
&& (event.statusText === undefined || typeof event.statusText === 'string')) {
if (!this.knownStatusKeys.has(event.statusKey)) return this.unsupported(event.method);
return {
kind: 'status',
conversationId,
key: event.statusKey,
...(typeof event.statusText === 'string'
? { text: bounded(event.statusText, MAX_MESSAGE_LENGTH) }
: {}),
};
}
if (event.method === 'setWidget'
&& typeof event.widgetKey === 'string'
&& (event.widgetLines === undefined || Array.isArray(event.widgetLines))) {
if (!this.knownWidgetKeys.has(event.widgetKey)) return this.unsupported(event.method);
if (Array.isArray(event.widgetLines)
&& !event.widgetLines.every((line) => typeof line === 'string')) {
return this.invalid(event.method);
}
return {
kind: 'widget',
conversationId,
key: event.widgetKey,
...(Array.isArray(event.widgetLines)
? {
lines: event.widgetLines.slice(0, MAX_WIDGET_LINES)
.map((line) => bounded(line as string, MAX_WIDGET_LINE_LENGTH)),
}
: {}),
...(event.widgetPlacement === 'aboveEditor' || event.widgetPlacement === 'belowEditor'
? { placement: event.widgetPlacement }
: {}),
};
}
if (event.method === 'setTitle' && typeof event.title === 'string') {
return { kind: 'title', conversationId, title: bounded(event.title, MAX_TITLE_LENGTH) };
}
if (event.method === 'set_editor_text' && typeof event.text === 'string') {
if (this.getDraftRevision(conversationId) !== run.draftRevision) {
this.diagnostics.push({ method: event.method, reason: 'stale-draft-revision' });
return null;
}
return {
kind: 'editor-text',
conversationId,
text: bounded(event.text, MAX_EDITOR_TEXT_LENGTH),
draftRevision: run.draftRevision,
};
}
if (['select', 'confirm', 'input', 'editor'].includes(event.method)) return null;
return this.invalid(event.method);
}
private unsupported(method: string): null {
this.diagnostics.push({ method: bounded(method, 128), reason: 'unsupported-ui-method' });
return null;
}
private invalid(method: string): null {
this.diagnostics.push({ method: bounded(method, 128), reason: 'invalid-ui-payload' });
return null;
}
}

View File

@@ -0,0 +1,100 @@
import path from 'node:path';
import { atomicWriteText } from '../../../coding-projects/atomic-json';
export const MAKELORE_PI_EXTENSION_VERSION = 1;
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 leases = new Map();
async function runtimeContext() {
const value = JSON.parse(await readFile(process.env.MAKELORE_PI_CONTEXT_FILE, 'utf8'));
if (!value.runId) throw new Error('Makelore run context is unavailable');
return value;
}
async function bridge(action, body, signal) {
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,
});
const result = await response.json().catch(() => ({}));
if (!response.ok) throw new Error(result.error || 'Makelore runtime bridge rejected the request');
return result;
}
async function releaseLease(toolCallId) {
const leaseId = leases.get(toolCallId);
if (!leaseId) return;
leases.delete(toolCallId);
await bridge('lease.release', { leaseId, resourceId: toolCallId }).catch(() => undefined);
}
async function releaseAll() {
await Promise.all([...leases.keys()].map(releaseLease));
}
export default function makeloreRuntime(pi) {
pi.registerTool({
name: 'ask_user',
label: 'Ask user',
description: 'Ask the user for a selection, confirmation, short input, or editor text.',
parameters: {
type: 'object',
additionalProperties: false,
required: ['kind', 'title'],
properties: {
kind: { type: 'string', enum: ['select', 'confirm', 'input', 'editor'] },
title: { type: 'string' },
message: { type: 'string' },
options: { type: 'array', items: { type: 'string' } },
},
},
async execute(_toolCallId, params, signal, _onUpdate, ctx) {
let value;
if (params.kind === 'select') {
value = await ctx.ui.select(params.title, Array.isArray(params.options) ? params.options : [], { signal });
} else if (params.kind === 'confirm') {
value = await ctx.ui.confirm(params.title, params.message || '', { signal });
} else if (params.kind === 'editor') {
value = await ctx.ui.editor(params.title, params.message || '', { signal });
} else {
value = await ctx.ui.input(params.title, params.message || '', { signal });
}
return {
content: [{ type: 'text', text: value === undefined ? 'User cancelled' : String(value) }],
details: { kind: params.kind, cancelled: value === undefined },
};
},
});
pi.on('tool_call', async (event, ctx) => {
if (!MUTATION_TOOLS.has(event.toolName)) return;
ctx.ui.setStatus('makelore.write-lease', '等待项目写入');
try {
const result = await bridge('lease.acquire', { resourceId: event.toolCallId }, ctx.signal);
leases.set(event.toolCallId, result.leaseId);
} finally {
ctx.ui.setStatus('makelore.write-lease', undefined);
}
});
pi.on('tool_result', async (event) => releaseLease(event.toolCallId));
pi.on('agent_end', releaseAll);
pi.on('session_shutdown', releaseAll);
}
`;
export async function materializeMakelorePiExtension(extensionsDir: string): Promise<string> {
const extensionPath = path.join(extensionsDir, MAKELORE_PI_EXTENSION_FILENAME);
await atomicWriteText(extensionPath, BUNDLE_SOURCE.trimStart());
return extensionPath;
}

View File

@@ -0,0 +1,177 @@
import type { ConversationInteraction } from '../contracts';
import type { PiRpcCommand, PiRpcEvent } from './rpc-client';
import type { PiGenerationResourceInput, PiWorkerPoolState } from './worker-pool';
interface PiInteractionTransport {
getState(conversationId: string): PiWorkerPoolState | null;
getActiveRun(conversationId: string): { runId: string; generation: number } | null;
send(conversationId: string, command: PiRpcCommand): Promise<void>;
trackGenerationResource(input: PiGenerationResourceInput): () => void;
}
interface StoredInteraction {
interaction: ConversationInteraction;
generation: number;
labels: Map<string, string>;
untrack(): void;
}
export type PiInteractionResponse =
| { interactionId: string; cancelled: true }
| { interactionId: string; optionId: string }
| { interactionId: string; confirmed: boolean }
| { interactionId: string; value: string };
function dialogEvent(event: PiRpcEvent): event is PiRpcEvent & {
id: string;
method: 'select' | 'confirm' | 'input' | 'editor';
title: string;
} {
return event.type === 'extension_ui_request'
&& typeof event.id === 'string'
&& ['select', 'confirm', 'input', 'editor'].includes(String(event.method))
&& typeof event.title === 'string';
}
export class PiInteractionStore {
private readonly pending = new Map<string, StoredInteraction>();
constructor(
private readonly transport: PiInteractionTransport,
private readonly onChange: (interaction: ConversationInteraction) => void,
) {}
list(conversationId?: string): ConversationInteraction[] {
return [...this.pending.values()]
.map(({ interaction }) => interaction)
.filter((interaction) => !conversationId || interaction.conversationId === conversationId)
.map((interaction) => structuredClone(interaction));
}
open(
conversationId: string,
generation: number,
runId: string,
event: PiRpcEvent,
): ConversationInteraction | null {
if (!dialogEvent(event)) return null;
const state = this.transport.getState(conversationId);
const active = this.transport.getActiveRun(conversationId);
if (state?.generation !== generation
|| active?.generation !== generation
|| active.runId !== runId) return null;
const key = this.key(conversationId, event.id);
if (this.pending.has(key)) throw new Error(`Duplicate Pi interaction id: ${event.id}`);
const labels = new Map<string, string>();
const options = event.method === 'select' && Array.isArray(event.options)
? event.options.flatMap((label, index) => {
if (typeof label !== 'string') return [];
const id = `${event.id}:option:${index}`;
labels.set(id, label);
return [{ id, label }];
})
: undefined;
const interaction: ConversationInteraction = {
id: event.id,
conversationId,
runId,
kind: event.method,
title: event.title,
...(typeof event.message === 'string' ? { message: event.message } : {}),
...(options ? { options } : {}),
status: 'pending',
};
const stored: StoredInteraction = {
interaction,
generation,
labels,
untrack: () => undefined,
};
stored.untrack = this.transport.trackGenerationResource({
conversationId,
kind: 'interaction',
id: event.id,
cancel: () => {
void this.cancelStored(stored, false);
},
});
this.pending.set(key, stored);
return structuredClone(interaction);
}
async respond(conversationId: string, response: PiInteractionResponse): Promise<ConversationInteraction> {
const stored = this.pending.get(this.key(conversationId, response.interactionId));
if (!stored) throw new Error('Pi interaction is not pending');
const state = this.transport.getState(conversationId);
const active = this.transport.getActiveRun(conversationId);
if (state?.generation !== stored.generation
|| active?.generation !== stored.generation
|| active.runId !== stored.interaction.runId) {
await this.cancelStored(stored, false);
throw new Error('Pi interaction belongs to a stale worker run');
}
let command: PiRpcCommand;
if ('cancelled' in response) {
command = { type: 'extension_ui_response', id: stored.interaction.id, cancelled: true };
} else if (stored.interaction.kind === 'select' && 'optionId' in response) {
const value = stored.labels.get(response.optionId);
if (value === undefined) throw new Error('Pi interaction option is invalid');
command = { type: 'extension_ui_response', id: stored.interaction.id, value };
} else if (stored.interaction.kind === 'confirm' && 'confirmed' in response) {
command = { type: 'extension_ui_response', id: stored.interaction.id, confirmed: response.confirmed };
} else if ((stored.interaction.kind === 'input' || stored.interaction.kind === 'editor')
&& 'value' in response) {
command = { type: 'extension_ui_response', id: stored.interaction.id, value: response.value };
} else {
throw new Error('Pi interaction response does not match its kind');
}
await this.transport.send(conversationId, command);
return this.finish(stored, 'cancelled' in response
? 'cancelled'
: stored.interaction.kind === 'confirm' && 'confirmed' in response && !response.confirmed
? 'rejected'
: 'answered');
}
async cancelRun(conversationId: string, runId: string, notifyWorker: boolean): Promise<void> {
const targets = [...this.pending.values()].filter(({ interaction }) => (
interaction.conversationId === conversationId && interaction.runId === runId
));
await Promise.all(targets.map((stored) => this.cancelStored(stored, notifyWorker)));
}
async cancelGeneration(conversationId: string, generation: number): Promise<void> {
const targets = [...this.pending.values()].filter((stored) => (
stored.interaction.conversationId === conversationId && stored.generation === generation
));
await Promise.all(targets.map((stored) => this.cancelStored(stored, false)));
}
private async cancelStored(stored: StoredInteraction, notifyWorker: boolean): Promise<void> {
if (!this.pending.has(this.key(stored.interaction.conversationId, stored.interaction.id))) return;
if (notifyWorker) {
await this.transport.send(stored.interaction.conversationId, {
type: 'extension_ui_response',
id: stored.interaction.id,
cancelled: true,
}).catch(() => undefined);
}
this.finish(stored, 'cancelled');
}
private finish(
stored: StoredInteraction,
status: Exclude<ConversationInteraction['status'], 'pending'>,
): ConversationInteraction {
this.pending.delete(this.key(stored.interaction.conversationId, stored.interaction.id));
stored.untrack();
const interaction = { ...stored.interaction, status };
this.onChange(structuredClone(interaction));
return interaction;
}
private key(conversationId: string, interactionId: string): string {
return `${conversationId}\u0000${interactionId}`;
}
}

View File

@@ -127,6 +127,19 @@ export class PiRpcClient {
throw lastError;
}
async send(command: PiRpcCommand): Promise<void> {
if (!command.type) throw new Error('Pi RPC command type is required');
if (this.invalidated) throw this.invalidated;
try {
await this.write(`${JSON.stringify(command)}\n`);
} catch (error) {
throw new PiProcessError('PI_RPC_WRITE_FAILED', `Could not write Pi RPC ${command.type}`, {
cause: error,
generation: this.generation,
});
}
}
accept(value: unknown): void {
if (!recordValue(value) || typeof value.type !== 'string') {
throw new PiProcessError('PI_RPC_PROTOCOL_ERROR', 'Pi RPC record must be an object with a type');

View File

@@ -74,6 +74,15 @@ import {
PiSessionProjectionError,
projectPiSessionSnapshot,
} from './session-projector';
import { PiManagedExtensionHost } from './extension-host';
import {
PiInteractionStore,
type PiInteractionResponse,
} from './interaction';
import {
PiExtensionUiProjector,
type PiExtensionUiProjection,
} from './extension-ui-projector';
type RuntimeIdKind = 'run' | 'queue';
@@ -88,6 +97,10 @@ export interface PiConversationRuntimeOptions {
providerRefreshCoordinator?: PiProviderRefreshCoordinator;
isAuthenticationError?(error: unknown): boolean;
refreshCredential?(accountId: string): Promise<void>;
extensionHost?: PiManagedExtensionHost;
getDraftRevision?(conversationId: string): number;
knownExtensionWidgetKeys?: readonly string[];
onExtensionUiProjection?(projection: PiExtensionUiProjection): void;
}
export interface PiWorkerProcessAdapter {
@@ -97,6 +110,7 @@ export interface PiWorkerProcessAdapter {
command: PiRpcCommand,
options?: PiRpcRequestOptions,
): Promise<PiRpcResponse<T>>;
send(command: PiRpcCommand): Promise<void>;
subscribe(listener: (event: PiRpcEvent) => void): () => void;
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void;
stop(): Promise<PiWorkerStopResult>;
@@ -120,6 +134,7 @@ export interface PiManagedWorkerOpenerOptions {
createProcess?: (options: PiWorkerProcessOptions) => PiWorkerProcessAdapter;
now?: () => number;
onTelemetry?: (event: PiRuntimeTelemetryEvent) => void;
extensionHost: PiManagedExtensionHost;
}
interface PiRpcSessionStateProjection {
@@ -132,12 +147,18 @@ class ManagedPiConversationWorker implements PiConversationWorker {
readonly id: string,
readonly generation: number,
private readonly process: PiWorkerProcessAdapter,
private readonly disposeExtension: () => Promise<void>,
private readonly unsubscribeExtensionInvalidation: () => void,
) {}
request<T = unknown>(command: PiRpcCommand, options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
return this.process.request<T>(command, options);
}
send(command: PiRpcCommand): Promise<void> {
return this.process.send(command);
}
subscribe(listener: (event: PiRpcEvent) => void): () => void {
return this.process.subscribe(listener);
}
@@ -146,8 +167,13 @@ class ManagedPiConversationWorker implements PiConversationWorker {
return this.process.subscribeInvalidation(listener);
}
stop(): Promise<PiWorkerStopResult> {
return this.process.stop();
async stop(): Promise<PiWorkerStopResult> {
this.unsubscribeExtensionInvalidation();
try {
return await this.process.stop();
} finally {
await this.disposeExtension();
}
}
}
@@ -193,6 +219,12 @@ export function createPiManagedWorkerOpener(
? { localProxyCredential: await options.getLocalProxyCredential() }
: {}),
});
const extension = await options.extensionHost.registerWorker({
conversationId: input.conversation.conversationId,
generation: input.generation,
projectId: input.conversation.projectId,
extensionsDir: managedPaths.extensionsDir,
});
recordManagedMilestone(
options.onTelemetry,
input,
@@ -214,11 +246,15 @@ export function createPiManagedWorkerOpener(
sessionDir: resources.projectSessionsDir,
additionalArgs: [
...buildPiManagedInputArgs(selection, resources),
'--extension', extension.extensionPath,
...(input.fork ? ['--fork', input.fork.sourceSession.piSessionId] : []),
'--session-id', sessionKey,
],
env: credential.env,
sensitiveValues: credential.sensitiveValues,
env: { ...credential.env, ...extension.env },
sensitiveValues: [...credential.sensitiveValues, ...extension.sensitiveValues],
});
let unsubscribeExtensionInvalidation = process.subscribeInvalidation(() => {
void extension.dispose();
});
try {
const spawnStartedAt = now();
@@ -290,11 +326,18 @@ export function createPiManagedWorkerOpener(
`${input.conversation.conversationId}:${input.generation}`,
input.generation,
process,
extension.dispose,
() => {
unsubscribeExtensionInvalidation();
unsubscribeExtensionInvalidation = () => undefined;
},
),
session: clone(bound.session),
};
} catch (error) {
unsubscribeExtensionInvalidation();
await process.stop().catch(() => undefined);
await extension.dispose();
throw error;
}
};
@@ -401,6 +444,10 @@ export class PiConversationRuntime implements CodingConversationRuntime {
private readonly providerRefresh: PiProviderRefreshCoordinator;
private readonly isAuthenticationError: ((error: unknown) => boolean) | undefined;
private readonly refreshCredential: ((accountId: string) => Promise<void>) | undefined;
private readonly extensionHost: PiManagedExtensionHost | undefined;
private readonly interactions: PiInteractionStore;
private readonly extensionUi: PiExtensionUiProjector;
private readonly onExtensionUiProjection: ((projection: PiExtensionUiProjection) => void) | undefined;
private readonly states = new Map<string, ConversationReducerState>();
private readonly inputs = new Map<string, PrepareConversationInput>();
private readonly listeners = new Set<(patch: ConversationPatchEnvelope) => void>();
@@ -432,6 +479,17 @@ export class PiConversationRuntime implements CodingConversationRuntime {
this.providerRefresh = options.providerRefreshCoordinator ?? new PiProviderRefreshCoordinator();
this.isAuthenticationError = options.isAuthenticationError;
this.refreshCredential = options.refreshCredential;
this.extensionHost = options.extensionHost;
this.interactions = new PiInteractionStore(this.pool, (interaction) => {
this.emit(interaction.conversationId, { op: 'interaction.upsert', interaction }, interaction.runId);
});
this.extensionUi = new PiExtensionUiProjector({
getDraftRevision: options.getDraftRevision ?? (() => 0),
...(options.knownExtensionWidgetKeys
? { knownWidgetKeys: options.knownExtensionWidgetKeys }
: {}),
});
this.onExtensionUiProjection = options.onExtensionUiProjection;
if (Boolean(this.isAuthenticationError) !== Boolean(this.refreshCredential)) {
throw new Error('Provider authentication detection and refresh must be configured together');
}
@@ -525,11 +583,25 @@ export class PiConversationRuntime implements CodingConversationRuntime {
message: input.text,
...(images.length > 0 ? { images } : {}),
};
const ticket = this.pool.startTopLevel({
conversationId: input.conversationId,
runId,
command,
});
const generation = this.pool.getState(input.conversationId)?.generation;
if (generation) this.extensionUi.beginRun(input.conversationId, generation, runId);
if (this.extensionHost && generation) {
await this.extensionHost.bindRun(input.conversationId, generation, runId);
}
let ticket;
try {
ticket = this.pool.startTopLevel({
conversationId: input.conversationId,
runId,
command,
});
} catch (error) {
this.extensionUi.endRun(input.conversationId, runId);
if (this.extensionHost && generation) {
await this.extensionHost.clearRun(input.conversationId, generation, runId);
}
throw error;
}
this.emit(input.conversationId, {
op: 'run.state',
run: {
@@ -567,8 +639,14 @@ export class PiConversationRuntime implements CodingConversationRuntime {
op: 'run.state',
run: { ...current, status: 'aborting' },
}, current.runId);
if (current.runId) await this.interactions.cancelRun(conversationId, current.runId, true);
try {
await this.pool.request(conversationId, { type: 'abort' });
const generation = this.pool.getState(conversationId)?.generation;
if (current.runId && generation) {
await this.extensionHost?.clearRun(conversationId, generation, current.runId);
this.extensionUi.endRun(conversationId, current.runId);
}
} catch (error) {
const latest = this.snapshot(conversationId).run;
if (latest.runId === current.runId && latest.status === 'aborting') {
@@ -651,11 +729,25 @@ export class PiConversationRuntime implements CodingConversationRuntime {
async compact(conversationId: string): Promise<void> {
await this.waitForProjection(conversationId);
const runId = this.id('run');
const ticket = this.pool.startTopLevel({
conversationId,
runId,
command: { type: 'compact' },
});
const generation = this.pool.getState(conversationId)?.generation;
if (generation) this.extensionUi.beginRun(conversationId, generation, runId);
if (this.extensionHost && generation) {
await this.extensionHost.bindRun(conversationId, generation, runId);
}
let ticket;
try {
ticket = this.pool.startTopLevel({
conversationId,
runId,
command: { type: 'compact' },
});
} catch (error) {
this.extensionUi.endRun(conversationId, runId);
if (this.extensionHost && generation) {
await this.extensionHost.clearRun(conversationId, generation, runId);
}
throw error;
}
this.emit(conversationId, {
op: 'run.state',
run: { status: 'compacting', runId, startedAt: this.now() },
@@ -713,6 +805,8 @@ export class PiConversationRuntime implements CodingConversationRuntime {
}
async dispose(conversationId: string): Promise<void> {
const state = this.pool.getState(conversationId);
if (state) await this.interactions.cancelGeneration(conversationId, state.generation);
await this.pool.dispose(conversationId);
this.registry.forget(conversationId);
this.inputs.delete(conversationId);
@@ -727,9 +821,17 @@ export class PiConversationRuntime implements CodingConversationRuntime {
return () => this.listeners.delete(listener);
}
async respondInteraction(
conversationId: string,
response: PiInteractionResponse,
): Promise<void> {
await this.interactions.respond(conversationId, response);
}
async shutdown(): Promise<void> {
this.unsubscribePool();
await this.pool.shutdown();
await this.extensionHost?.close();
this.projectors.clear();
this.projectionChains.clear();
this.hydrationFlights.clear();
@@ -900,12 +1002,25 @@ export class PiConversationRuntime implements CodingConversationRuntime {
if (!this.states.has(event.conversationId)) return;
this.replaceWorkerGeneration(event.conversationId, event.state, true);
this.resetProjector(event.conversationId);
const runId = this.states.get(event.conversationId)?.snapshot.run.runId;
if (runId) this.extensionUi.replaceGeneration(event.conversationId, event.generation, runId);
if (this.extensionHost && runId) {
void this.extensionHost.bindRun(event.conversationId, event.generation, runId).catch((error) => {
this.recordProjectionFailure(event.conversationId, event.generation, error);
});
}
void this.requestHydration(event.conversationId, event.state, true).catch((error) => {
this.recordProjectionFailure(event.conversationId, event.generation, error);
});
return;
}
if (event.type === 'worker.crashed') {
void this.interactions.cancelGeneration(event.conversationId, event.generation);
const runId = this.states.get(event.conversationId)?.snapshot.run.runId;
if (runId) {
this.extensionUi.endRun(event.conversationId, runId);
void this.extensionHost?.clearRun(event.conversationId, event.generation, runId);
}
const state = this.pool.getState(event.conversationId);
if (state) this.emit(event.conversationId, { op: 'worker.state', state: publicWorkerState(state) });
return;
@@ -913,6 +1028,26 @@ export class PiConversationRuntime implements CodingConversationRuntime {
void this.enqueueProjection(event.conversationId, async () => {
const snapshot = this.states.get(event.conversationId)?.snapshot;
if (!snapshot || snapshot.cursor.workerGeneration !== event.generation) return;
if (event.event.type === 'extension_ui_request' && snapshot.run.runId) {
const interaction = this.interactions.open(
event.conversationId,
event.generation,
snapshot.run.runId,
event.event,
);
if (interaction) {
this.emit(event.conversationId, { op: 'interaction.upsert', interaction }, interaction.runId);
return;
}
const projection = this.extensionUi.project(
event.conversationId,
event.generation,
snapshot.run.runId,
event.event,
);
if (projection) this.onExtensionUiProjection?.(projection);
return;
}
const projector = this.projector(event.conversationId);
const patches = await projector.project(snapshot, event.event);
for (const patch of patches) {
@@ -928,6 +1063,15 @@ export class PiConversationRuntime implements CodingConversationRuntime {
);
}
}
if (event.event.type === 'agent_settled' && snapshot.run.runId) {
await this.interactions.cancelRun(event.conversationId, snapshot.run.runId, true);
await this.extensionHost?.clearRun(
event.conversationId,
event.generation,
snapshot.run.runId,
);
this.extensionUi.endRun(event.conversationId, snapshot.run.runId);
}
}).catch((error) => {
this.recordProjectionFailure(event.conversationId, event.generation, error);
});
@@ -946,6 +1090,10 @@ export class PiConversationRuntime implements CodingConversationRuntime {
error: runtimeFailure(error),
},
}, runId);
void this.interactions.cancelRun(conversationId, runId, true);
this.extensionUi.endRun(conversationId, runId);
const generation = this.pool.getState(conversationId)?.generation;
if (generation) void this.extensionHost?.clearRun(conversationId, generation, runId);
}
private requestHydration(

View File

@@ -24,6 +24,7 @@ export interface PiConversationWorker {
command: PiRpcCommand,
options?: PiRpcRequestOptions,
): Promise<PiRpcResponse<T>>;
send(command: PiRpcCommand): Promise<void>;
subscribe(listener: (event: PiRpcEvent) => void): () => void;
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void;
stop(): Promise<PiWorkerStopResult>;
@@ -365,6 +366,18 @@ export class PiWorkerPool {
}
}
async send(conversationId: string, command: PiRpcCommand): Promise<void> {
let record = this.workers.get(conversationId);
if (!record || record.state === 'crashed') throw new Error('Conversation worker is not available');
if (record.rebuildFlight) record = await record.rebuildFlight;
await record.worker.send(command);
}
getActiveRun(conversationId: string): { runId: string; generation: number } | null {
const active = this.activeRuns.get(conversationId);
return active ? { ...active } : null;
}
failTopLevel(conversationId: string, runId: string, error: Error): void {
const active = this.activeRuns.get(conversationId);
if (active?.runId === runId) {

View File

@@ -75,7 +75,7 @@ export function buildPiRpcArgs(
'--no-themes',
'--no-context-files',
'--no-approve',
'--no-tools',
'--tools', 'read,bash,edit,write,grep,find,ls,ask_user',
...additionalArgs,
];
}
@@ -310,6 +310,17 @@ export class PiWorkerProcess {
return this.rpc.request<T>(command, options);
}
send(command: PiRpcCommand): Promise<void> {
if (!this.rpc) {
return Promise.reject(new PiProcessError(
'PI_WORKER_START_FAILED',
'Pi worker has not started',
{ generation: this.generationValue },
));
}
return this.rpc.send(command);
}
subscribe(listener: (event: PiRpcEvent) => void): () => void {
if (!this.rpc) throw new Error('Pi worker has not started');
return this.rpc.subscribe(listener);

View File

@@ -0,0 +1,103 @@
import { randomUUID } from 'node:crypto';
export interface PiProjectWriteLease {
id: string;
projectId: string;
holderId: string;
release(): void;
}
interface WaitingLease {
projectId: string;
holderId: string;
signal?: AbortSignal;
resolve(lease: PiProjectWriteLease): void;
reject(error: Error): void;
onAbort?: () => void;
}
export class PiProjectWriteLeaseCoordinator {
private readonly active = new Map<string, PiProjectWriteLease>();
private readonly waiting = new Map<string, WaitingLease[]>();
get activeCount(): number { return this.active.size; }
waitingCount(projectId?: string): number {
if (projectId) return this.waiting.get(projectId)?.length ?? 0;
return [...this.waiting.values()].reduce((total, queue) => total + queue.length, 0);
}
acquire(projectId: string, holderId: string, signal?: AbortSignal): Promise<PiProjectWriteLease> {
if (!projectId.trim() || !holderId.trim()) throw new Error('Project and lease holder are required');
if (signal?.aborted) return Promise.reject(new Error('Project write lease wait cancelled'));
if (!this.active.has(projectId)) return Promise.resolve(this.issue(projectId, holderId));
return new Promise<PiProjectWriteLease>((resolve, reject) => {
const waiter: WaitingLease = { projectId, holderId, resolve, reject, ...(signal ? { signal } : {}) };
if (signal) {
waiter.onAbort = () => {
this.removeWaiter(waiter);
reject(new Error('Project write lease wait cancelled'));
};
signal.addEventListener('abort', waiter.onAbort, { once: true });
}
const queue = this.waiting.get(projectId) ?? [];
queue.push(waiter);
this.waiting.set(projectId, queue);
});
}
cancelProject(projectId: string): void {
for (const waiter of this.waiting.get(projectId) ?? []) {
this.detachAbort(waiter);
waiter.reject(new Error('Project write lease wait cancelled'));
}
this.waiting.delete(projectId);
this.active.get(projectId)?.release();
}
private issue(projectId: string, holderId: string): PiProjectWriteLease {
let released = false;
const lease: PiProjectWriteLease = {
id: randomUUID(),
projectId,
holderId,
release: () => {
if (released) return;
released = true;
if (this.active.get(projectId) === lease) this.active.delete(projectId);
this.advance(projectId);
},
};
this.active.set(projectId, lease);
return lease;
}
private advance(projectId: string): void {
const queue = this.waiting.get(projectId);
while (!this.active.has(projectId) && queue?.length) {
const waiter = queue.shift() as WaitingLease;
this.detachAbort(waiter);
if (waiter.signal?.aborted) {
waiter.reject(new Error('Project write lease wait cancelled'));
continue;
}
waiter.resolve(this.issue(projectId, waiter.holderId));
}
if (!queue?.length) this.waiting.delete(projectId);
}
private removeWaiter(waiter: WaitingLease): void {
const queue = this.waiting.get(waiter.projectId);
const index = queue?.indexOf(waiter) ?? -1;
if (index >= 0) queue?.splice(index, 1);
if (!queue?.length) this.waiting.delete(waiter.projectId);
this.detachAbort(waiter);
}
private detachAbort(waiter: WaitingLease): void {
if (waiter.signal && waiter.onAbort) {
waiter.signal.removeEventListener('abort', waiter.onAbort);
}
}
}