feat(coding-runtime): add managed Pi extension host
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user