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

@@ -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(