feat(coding-runtime): add managed Pi extension host
This commit is contained in:
177
electron/coding-runtime/pi/interaction.ts
Normal file
177
electron/coding-runtime/pi/interaction.ts
Normal 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}`;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user