Files
makelore/electron/coding-runtime/pi/runtime.ts

1277 lines
46 KiB
TypeScript

import { randomUUID } from 'node:crypto';
import path from 'node:path';
import type { ModelSummary, ProviderAccount } from '../../shared/providers/types';
import { validateSessionKey } from '../../coding-projects/conversation-store';
import {
buildPiProviderCatalog,
buildPiWorkerCredentialProjection,
selectPiProviderModel,
writePiProviderCatalog,
type PiProviderSelection,
} from './provider-config';
import { PiProviderRefreshCoordinator } from './provider-refresh';
import type {
CodingConversationRuntime,
CodingRuntimePublicError,
ConversationModelState,
ConversationPatch,
ConversationPatchEnvelope,
ConversationRuntimeState,
ConversationSnapshot,
ForkConversationInput,
ForkResult,
PrepareConversationInput,
ProductModelRef,
PromptAcceptance,
PromptConversationInput,
QueueAcceptance,
QueueMessageInput,
SetConversationModelInput,
SetThinkingLevelInput,
} from '../contracts';
import {
createConversationReducerState,
reduceConversationPatch,
type ConversationReducerState,
} from '../conversation-reducer';
import { CodingRuntimeContractError } from '../in-memory-conversation-runtime';
import {
PiEventProjector,
type PiEventProjectorOptions,
} from './event-projector';
import { PiProcessError } from './process-errors';
import type { PiSessionRegistry } from './session-registry';
import {
buildPiManagedInputArgs,
ensurePiManagedPaths,
materializePiAgentResources,
} from './resource-loader';
import type {
PiRpcCommand,
PiRpcEvent,
PiRpcRequestOptions,
PiRpcResponse,
} from './rpc-client';
import {
PiWorkerProcess,
type PiWorkerProcessOptions,
type PiWorkerStopResult,
} from './worker-process';
import {
createPiRuntimeTelemetryEvent,
type PiRuntimeMilestone,
type PiRuntimeTelemetryEvent,
} from './telemetry';
import {
type PiConversationWorker,
type PiWorkerOpenInput,
type PiWorkerOpenResult,
type PiWorkerPoolEvent,
type PiWorkerPoolState,
PiWorkerPool,
} from './worker-pool';
import {
PiSessionProjectionError,
projectPiSessionSnapshot,
} from './session-projector';
import { PiManagedExtensionHost } from './extension-host';
import type { PiSubagentScheduler } from './subagent';
import {
PiInteractionStore,
type PiInteractionResponse,
} from './interaction';
import {
PiExtensionUiProjector,
type PiExtensionUiProjection,
} from './extension-ui-projector';
type RuntimeIdKind = 'run' | 'queue';
export interface PiConversationRuntimeOptions {
pool: PiWorkerPool;
registry: PiSessionRegistry;
resolveModel(model: ProductModelRef): Promise<PiProviderSelection>;
resolveImages?(attachments: Array<{ attachmentId: string }>): Promise<unknown[]>;
projectImage?: PiEventProjectorOptions['projectImage'];
createId?(kind: RuntimeIdKind): string;
now?: () => number;
providerRefreshCoordinator?: PiProviderRefreshCoordinator;
isAuthenticationError?(error: unknown): boolean;
refreshCredential?(accountId: string): Promise<void>;
extensionHost?: PiManagedExtensionHost;
subagentScheduler?: PiSubagentScheduler;
getDraftRevision?(conversationId: string): number;
knownExtensionWidgetKeys?: readonly string[];
onExtensionUiProjection?(projection: PiExtensionUiProjection): void;
}
export interface PiWorkerProcessAdapter {
readonly generation: number;
start(): Promise<unknown>;
request<T = unknown>(
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>;
}
export interface PiManagedProviderInput {
accounts: ProviderAccount[];
modelSummaries: ModelSummary[];
}
export interface PiManagedWorkerOpenerOptions {
registry: PiSessionRegistry;
executablePath: string;
cliPath: string;
userDataDir: string;
bundledSkillsDir: string;
loadProviderInput(): Promise<PiManagedProviderInput>;
resolveCredential(account: ProviderAccount): Promise<string | null>;
getLocalProxyCredential?(): Promise<string | undefined>;
createSessionKey?: () => string;
createProcess?: (options: PiWorkerProcessOptions) => PiWorkerProcessAdapter;
now?: () => number;
onTelemetry?: (event: PiRuntimeTelemetryEvent) => void;
extensionHost: PiManagedExtensionHost;
}
interface PiRpcSessionStateProjection {
sessionId: string;
sessionFile?: string;
}
class ManagedPiConversationWorker implements PiConversationWorker {
constructor(
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);
}
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
return this.process.subscribeInvalidation(listener);
}
async stop(): Promise<PiWorkerStopResult> {
this.unsubscribeExtensionInvalidation();
try {
return await this.process.stop();
} finally {
await this.disposeExtension();
}
}
}
export function createPiManagedWorkerOpener(
options: PiManagedWorkerOpenerOptions,
): (input: PiWorkerOpenInput) => Promise<PiWorkerOpenResult> {
const createProcess = options.createProcess ?? ((processOptions) => new PiWorkerProcess(processOptions));
const createSessionKey = options.createSessionKey ?? randomUUID;
const now = options.now ?? Date.now;
return async (input) => {
const resourcesStartedAt = now();
const registered = await options.registry.prepare(input.conversation);
const model = registered.conversation.model;
if (!model || registered.conversation.modelResolution !== 'resolved') {
throw new CodingRuntimeContractError(
'CODING_MIGRATION_MODEL_REQUIRED',
'Conversation model must be selected before opening Pi',
true,
);
}
const providerInput = await options.loadProviderInput();
const catalog = buildPiProviderCatalog(providerInput);
const selection = selectPiProviderModel(catalog, model);
const account = providerInput.accounts.find((candidate) => candidate.id === selection.accountId);
const descriptor = catalog.descriptors.find((candidate) => candidate.accountId === selection.accountId);
if (!account || !descriptor) throw new Error('Selected Provider account is unavailable');
const managedPaths = await ensurePiManagedPaths(options.userDataDir);
await writePiProviderCatalog(managedPaths.modelsFile, catalog, model);
const resources = await materializePiAgentResources({
userDataDir: options.userDataDir,
projectId: input.conversation.projectId,
agentId: registered.agent.id,
prompt: registered.agent.prompt,
skillIds: registered.agent.skillIds,
bundledSkillsDir: options.bundledSkillsDir,
revision: input.revision,
});
const credential = await buildPiWorkerCredentialProjection({
account,
descriptor,
resolveCredential: options.resolveCredential,
...(options.getLocalProxyCredential
? { localProxyCredential: await options.getLocalProxyCredential() }
: {}),
});
const extension = await options.extensionHost.registerWorker({
conversationId: input.conversation.conversationId,
generation: input.generation,
projectId: input.conversation.projectId,
projectPath: registered.projectPath,
skillIds: registered.agent.skillIds,
extensionsDir: managedPaths.extensionsDir,
});
recordManagedMilestone(
options.onTelemetry,
input,
'resources.ready',
now() - resourcesStartedAt,
now(),
);
const sessionKey = validateSessionKey(
input.existingSession?.sessionKey ?? createSessionKey(),
);
if (input.existingSession && input.fork) {
throw new Error('Pi worker cannot reopen and fork a session at the same time');
}
const process = createProcess({
executablePath: options.executablePath,
cliPath: options.cliPath,
cwd: registered.projectPath,
configDir: resources.paths.configDir,
sessionDir: resources.projectSessionsDir,
additionalArgs: [
...buildPiManagedInputArgs(selection, resources),
'--extension', extension.extensionPath,
...(input.fork ? ['--fork', input.fork.sourceSession.piSessionId] : []),
'--session-id', sessionKey,
],
env: { ...credential.env, ...extension.env },
sensitiveValues: [...credential.sensitiveValues, ...extension.sensitiveValues],
});
let unsubscribeExtensionInvalidation = process.subscribeInvalidation(() => {
void extension.dispose();
});
try {
const spawnStartedAt = now();
await process.start();
recordManagedMilestone(
options.onTelemetry,
input,
'worker.spawn',
now() - spawnStartedAt,
now(),
);
if (input.fork?.sourceEntryId) {
await process.request({ type: 'fork', entryId: input.fork.sourceEntryId });
}
const readyStartedAt = now();
const response = await process.request<PiRpcSessionStateProjection>(
{ type: 'get_state' },
{ retry: 'read-only-once' },
);
recordManagedMilestone(
options.onTelemetry,
input,
'rpc.ready',
now() - readyStartedAt,
now(),
);
const piSessionId = response.data?.sessionId?.trim();
if (!piSessionId) throw new Error('Pi worker did not return a session id');
if (input.existingSession && input.existingSession.piSessionId !== piSessionId) {
throw new CodingRuntimeContractError(
'CODING_SESSION_UNREADABLE',
'Pi reopened a different Conversation session',
true,
);
}
if (response.data?.sessionFile) {
const sessionsRoot = path.resolve(resources.projectSessionsDir);
const sessionFile = path.resolve(response.data.sessionFile);
const relativeSessionFile = path.relative(sessionsRoot, sessionFile);
if (!relativeSessionFile || relativeSessionFile.startsWith('..') || path.isAbsolute(relativeSessionFile)) {
throw new CodingRuntimeContractError(
'CODING_SESSION_UNREADABLE',
'Pi session file is outside the managed session directory',
true,
);
}
}
const sessionStartedAt = now();
const bound = await options.registry.ensureBinding(input.conversation, async () => ({
piSessionId,
sessionKey,
}));
if (!bound.session || bound.session.piSessionId !== piSessionId || bound.session.sessionKey !== sessionKey) {
throw new CodingRuntimeContractError(
'CODING_SESSION_UNREADABLE',
'Pi session binding does not match the Conversation registry',
true,
);
}
recordManagedMilestone(
options.onTelemetry,
input,
'session.open',
now() - sessionStartedAt,
now(),
);
return {
worker: new ManagedPiConversationWorker(
`${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;
}
};
}
function recordManagedMilestone(
listener: ((event: PiRuntimeTelemetryEvent) => void) | undefined,
input: PiWorkerOpenInput,
milestone: Extract<PiRuntimeMilestone, 'resources.ready' | 'worker.spawn' | 'rpc.ready' | 'session.open'>,
durationMs: number,
at: number,
): void {
listener?.(createPiRuntimeTelemetryEvent({
milestone,
conversationId: input.conversation.conversationId,
workerGeneration: input.generation,
cold: true,
durationMs,
at,
}));
}
function clone<T>(value: T): T {
return structuredClone(value);
}
function publicWorkerState(state: PiWorkerPoolState): ConversationSnapshot['worker'] {
if (state.state === 'spawning') return { status: 'starting', generation: state.generation };
if (state.state === 'crashed') {
return {
status: 'error',
generation: state.generation,
error: {
code: state.failureCode === 'PI_RPC_PROTOCOL_ERROR'
? 'CODING_RUNTIME_PROTOCOL_ERROR'
: 'CODING_RUNTIME_START_FAILED',
message: 'The local Agent worker stopped unexpectedly',
recoverable: true,
},
};
}
return { status: 'ready', generation: state.generation };
}
function runtimeFailure(error: unknown): CodingRuntimePublicError {
if (error instanceof CodingRuntimeContractError) return clone(error.publicError);
if (error instanceof PiSessionProjectionError) {
return { code: error.code, message: error.message, recoverable: error.recoverable };
}
if (error instanceof PiProcessError) {
if (error.code === 'PI_RPC_PROTOCOL_ERROR') {
return {
code: 'CODING_RUNTIME_PROTOCOL_ERROR',
message: 'The local Agent protocol failed',
recoverable: true,
};
}
if (error.code === 'PI_RPC_TIMEOUT') {
return {
code: 'CODING_REQUEST_UNCERTAIN',
message: 'The local Agent did not confirm the request',
recoverable: true,
};
}
}
return {
code: 'CODING_RUNTIME_START_FAILED',
message: 'The local Agent is unavailable',
recoverable: true,
};
}
function emptySnapshot(
input: PrepareConversationInput,
workerState: PiWorkerPoolState,
): ConversationSnapshot {
return {
schemaVersion: 1,
conversation: {
id: input.conversationId,
projectId: input.projectId,
agentId: input.agentId,
title: input.title,
model: clone(input.model),
},
nodes: [],
run: { status: 'idle' },
queue: { items: [] },
context: { usedTokens: 0, contextWindow: 0, compaction: 'idle' },
pendingInteractions: [],
worker: publicWorkerState(workerState),
cursor: { workerGeneration: workerState.generation, seq: 0 },
};
}
export class PiConversationRuntime implements CodingConversationRuntime {
private readonly pool: PiWorkerPool;
private readonly registry: PiSessionRegistry;
private readonly resolveModel: PiConversationRuntimeOptions['resolveModel'];
private readonly resolveImages: NonNullable<PiConversationRuntimeOptions['resolveImages']>;
private readonly projectImage: PiEventProjectorOptions['projectImage'];
private readonly createRuntimeId: (kind: RuntimeIdKind) => string;
private readonly now: () => number;
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>();
private readonly projectors = new Map<string, PiEventProjector>();
private readonly projectionChains = new Map<string, Promise<void>>();
private readonly hydrationFlights = new Map<
string,
{ generation: number; flight: Promise<void> }
>();
private readonly unsubscribePool: () => void;
constructor(options: PiConversationRuntimeOptions) {
this.pool = options.pool;
this.registry = options.registry;
this.resolveModel = options.resolveModel;
this.resolveImages = options.resolveImages ?? (async (attachments) => {
if (attachments.length > 0) {
throw new CodingRuntimeContractError(
'CODING_RUNTIME_START_FAILED',
'Conversation attachments are not ready',
true,
);
}
return [];
});
this.projectImage = options.projectImage;
this.createRuntimeId = options.createId ?? (() => randomUUID());
this.now = options.now ?? Date.now;
this.providerRefresh = options.providerRefreshCoordinator ?? new PiProviderRefreshCoordinator();
this.isAuthenticationError = options.isAuthenticationError;
this.refreshCredential = options.refreshCredential;
this.extensionHost = options.extensionHost;
if (options.subagentScheduler && !this.extensionHost) {
throw new Error('Subagent scheduler requires the managed extension host');
}
if (options.subagentScheduler) {
this.extensionHost?.configureSubagents({
scheduler: options.subagentScheduler,
trackGenerationResource: (input) => this.pool.trackGenerationResource(input),
});
}
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');
}
this.unsubscribePool = this.pool.subscribe((event) => this.onPoolEvent(event));
}
async prepare(input: PrepareConversationInput): Promise<ConversationRuntimeState> {
const registered = await this.registry.prepare(input);
const canonicalInput: PrepareConversationInput = {
conversationId: registered.conversation.id,
projectId: input.projectId,
agentId: registered.conversation.agentId,
title: registered.conversation.title,
model: {
model: registered.conversation.model ? clone(registered.conversation.model) : null,
modelResolution: registered.conversation.modelResolution,
},
};
const worker = await this.prepareWorker(canonicalInput);
await this.registry.ensureBinding(canonicalInput, async () => clone(worker.session));
this.inputs.set(input.conversationId, clone(canonicalInput));
const isNewState = !this.states.has(input.conversationId);
if (isNewState) {
this.states.set(input.conversationId, createConversationReducerState(
emptySnapshot(canonicalInput, worker),
));
this.resetProjector(input.conversationId);
} else {
this.emit(input.conversationId, { op: 'worker.state', state: publicWorkerState(worker) });
}
if (isNewState) {
try {
await this.requestHydration(input.conversationId, worker, false);
} catch (error) {
this.recordProjectionFailure(input.conversationId, worker.generation, error);
throw error;
}
}
return this.runtimeState(input.conversationId);
}
async getSnapshot(conversationId: string): Promise<ConversationSnapshot> {
await this.waitForProjection(conversationId);
return clone(this.snapshot(conversationId));
}
async prompt(input: PromptConversationInput): Promise<PromptAcceptance> {
await this.waitForProjection(input.conversationId);
if (input.mode === 'steer') {
const acceptance = await this.steer(input);
return {
accepted: true,
conversationId: input.conversationId,
clientRequestId: input.clientRequestId,
runId: this.snapshot(input.conversationId).run.runId ?? this.id('run'),
mode: 'steer',
queuePosition: acceptance.queuePosition,
};
}
if (input.mode === 'follow-up') {
const acceptance = await this.followUp(input);
return {
accepted: true,
conversationId: input.conversationId,
clientRequestId: input.clientRequestId,
runId: this.snapshot(input.conversationId).run.runId ?? this.id('run'),
mode: 'follow-up',
queuePosition: acceptance.queuePosition,
};
}
this.snapshot(input.conversationId);
const images = await this.resolveImages(input.attachments);
const runId = this.id('run');
const messageId = `client:${input.clientRequestId}`;
this.emit(input.conversationId, {
op: 'message.upsert',
node: {
kind: 'message',
id: messageId,
clientRequestId: input.clientRequestId,
role: 'user',
status: 'optimistic',
blocks: input.text.length > 0
? [{ kind: 'text', id: `${messageId}:content:0`, text: input.text, status: 'complete' }]
: [],
},
});
const command: PiRpcCommand = {
type: 'prompt',
message: input.text,
...(images.length > 0 ? { images } : {}),
};
const generation = this.pool.getState(input.conversationId)?.generation;
if (generation) this.extensionUi.beginRun(input.conversationId, generation, runId);
let ticket;
try {
if (this.extensionHost && generation) {
await this.extensionHost.bindRun(input.conversationId, generation, runId);
}
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: {
status: ticket.queuePosition ? 'queued' : 'running',
runId,
mode: 'prompt',
startedAt: this.now(),
},
}, runId);
const acceptance = this.acceptPrompt(input.conversationId, runId, command, ticket.accepted);
if (ticket.queuePosition) void acceptance.catch(() => undefined);
else await acceptance;
return {
accepted: true,
conversationId: input.conversationId,
clientRequestId: input.clientRequestId,
runId,
mode: 'prompt',
...(ticket.queuePosition ? { queuePosition: ticket.queuePosition } : {}),
};
}
async steer(input: QueueMessageInput): Promise<QueueAcceptance> {
return await this.queue('steer', input);
}
async followUp(input: QueueMessageInput): Promise<QueueAcceptance> {
return await this.queue('follow-up', input);
}
async abort(conversationId: string): Promise<void> {
await this.waitForProjection(conversationId);
const current = this.snapshot(conversationId).run;
this.emit(conversationId, {
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') {
this.emit(conversationId, { op: 'run.state', run: current }, current.runId);
}
throw error;
}
}
async setModel(input: SetConversationModelInput): Promise<ConversationModelState> {
await this.waitForProjection(input.conversationId);
const snapshot = this.snapshot(input.conversationId);
const thinkingLevel = snapshot.conversation.model.model?.thinkingLevel ?? 'off';
const selection = await this.resolveModel({
accountId: input.accountId,
modelId: input.modelId,
thinkingLevel,
});
const model: ConversationModelState = {
model: {
accountId: selection.accountId,
modelId: selection.modelId,
thinkingLevel,
},
modelResolution: 'resolved',
};
if (snapshot.conversation.model.model?.accountId !== selection.accountId) {
const persisted = await this.registry.setModel(input.conversationId, model);
this.replaceModel(input.conversationId, persisted);
try {
const worker = await this.pool.reconfigureConversationModel(input.conversationId, persisted);
if (worker) await this.requestHydration(input.conversationId, worker, true);
} catch (error) {
const worker = this.pool.getState(input.conversationId);
if (worker) {
if (this.snapshot(input.conversationId).cursor.workerGeneration !== worker.generation) {
this.replaceWorkerGeneration(input.conversationId, worker, true);
}
this.recordProjectionFailure(input.conversationId, worker.generation, error);
}
throw error;
}
return clone(persisted);
}
await this.pool.request(input.conversationId, {
type: 'set_model',
provider: selection.runtimeProviderId,
modelId: selection.modelId,
});
const persisted = await this.registry.setModel(input.conversationId, model);
this.pool.updateConversationModel(input.conversationId, persisted);
this.replaceModel(input.conversationId, persisted);
return clone(persisted);
}
async setThinking(input: SetThinkingLevelInput): Promise<ConversationModelState> {
await this.waitForProjection(input.conversationId);
const current = this.snapshot(input.conversationId).conversation.model;
if (!current.model) {
throw new CodingRuntimeContractError(
'CODING_MIGRATION_MODEL_REQUIRED',
'Conversation model must be selected first',
true,
);
}
await this.pool.request(input.conversationId, {
type: 'set_thinking_level',
level: input.thinkingLevel,
});
const model: ConversationModelState = {
model: { ...current.model, thinkingLevel: input.thinkingLevel },
modelResolution: 'resolved',
};
const persisted = await this.registry.setModel(input.conversationId, model);
this.pool.updateConversationModel(input.conversationId, persisted);
this.replaceModel(input.conversationId, persisted);
return clone(persisted);
}
async compact(conversationId: string): Promise<void> {
await this.waitForProjection(conversationId);
const runId = this.id('run');
const generation = this.pool.getState(conversationId)?.generation;
if (generation) this.extensionUi.beginRun(conversationId, generation, runId);
let ticket;
try {
if (this.extensionHost && generation) {
await this.extensionHost.bindRun(conversationId, generation, runId);
}
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() },
}, runId);
try {
await ticket.accepted;
} catch (error) {
this.failRun(conversationId, runId, error);
throw error;
}
}
async fork(input: ForkConversationInput): Promise<ForkResult> {
await this.waitForProjection(input.sourceConversationId);
this.snapshot(input.sourceConversationId);
const registered = await this.registry.prepare(input.conversation);
const canonicalInput: PrepareConversationInput = {
conversationId: registered.conversation.id,
projectId: input.conversation.projectId,
agentId: registered.conversation.agentId,
title: registered.conversation.title,
model: {
model: registered.conversation.model ? clone(registered.conversation.model) : null,
modelResolution: registered.conversation.modelResolution,
},
};
const worker = await this.pool.fork(
input.sourceConversationId,
canonicalInput,
input.sourceEntryId,
);
await this.registry.ensureBinding(canonicalInput, async () => clone(worker.session));
this.inputs.set(canonicalInput.conversationId, clone(canonicalInput));
const snapshot = emptySnapshot(canonicalInput, worker);
this.states.set(canonicalInput.conversationId, createConversationReducerState(snapshot));
this.resetProjector(canonicalInput.conversationId);
try {
await this.requestHydration(canonicalInput.conversationId, worker, false);
} catch (error) {
this.recordProjectionFailure(canonicalInput.conversationId, worker.generation, error);
throw error;
}
return {
conversationId: canonicalInput.conversationId,
snapshot: clone(this.snapshot(canonicalInput.conversationId)),
};
}
async recover(conversationId: string): Promise<ConversationRuntimeState> {
await this.waitForProjection(conversationId);
const before = this.snapshot(conversationId);
const generation = this.pool.getState(conversationId)?.generation;
if (before.run.runId && generation) {
await this.interactions.cancelRun(conversationId, before.run.runId, true);
this.extensionUi.endRun(conversationId, before.run.runId);
await this.extensionHost?.clearRun(conversationId, generation, before.run.runId);
}
const state = await this.pool.recover(conversationId);
await this.requestHydration(conversationId, state, false);
this.settleRecoveredRun(conversationId);
return this.runtimeState(conversationId);
}
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);
this.states.delete(conversationId);
this.projectors.delete(conversationId);
this.projectionChains.delete(conversationId);
this.hydrationFlights.delete(conversationId);
}
subscribe(listener: (patch: ConversationPatchEnvelope) => void): () => void {
this.listeners.add(listener);
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();
}
private async queue(
mode: 'steer' | 'follow-up',
input: QueueMessageInput,
): Promise<QueueAcceptance> {
await this.waitForProjection(input.conversationId);
const snapshot = this.snapshot(input.conversationId);
const images = await this.resolveImages(input.attachments);
const queuePosition = snapshot.queue.items.length + 1;
const queueId = this.id('queue');
this.emit(input.conversationId, {
op: 'queue.replace',
queue: {
items: [
...snapshot.queue.items,
{
id: queueId,
clientRequestId: input.clientRequestId,
mode,
text: input.text,
attachmentIds: input.attachments.map(({ attachmentId }) => attachmentId),
},
],
},
}, snapshot.run.runId);
try {
await this.pool.request(input.conversationId, {
type: mode === 'steer' ? 'steer' : 'follow_up',
message: input.text,
...(images.length > 0 ? { images } : {}),
});
} catch (error) {
const current = this.snapshot(input.conversationId);
this.emit(input.conversationId, {
op: 'queue.replace',
queue: { items: current.queue.items.filter(({ id }) => id !== queueId) },
}, current.run.runId);
throw error;
}
return {
accepted: true,
conversationId: input.conversationId,
clientRequestId: input.clientRequestId,
mode,
queuePosition,
};
}
private async prepareWorker(input: PrepareConversationInput): Promise<PiWorkerPoolState> {
if (!this.isAuthenticationError || !this.refreshCredential || !input.model.model) {
return await this.pool.prepare(input);
}
return await this.providerRefresh.withSingleAuthRecovery({
accountId: input.model.model.accountId,
operation: async () => await this.pool.prepare(input),
isAuthenticationError: this.isAuthenticationError,
refreshCredential: async () => await this.refreshCredential!(input.model.model!.accountId),
reopenWorker: async () => {
if (!this.pool.getState(input.conversationId)) return;
const recovered = await this.pool.recover(input.conversationId);
if (this.states.has(input.conversationId)) {
await this.requestHydration(input.conversationId, recovered, false);
}
},
});
}
private async acceptPrompt(
conversationId: string,
runId: string,
command: PiRpcCommand,
firstAcceptance: Promise<PiRpcResponse>,
): Promise<void> {
const accountId = this.snapshot(conversationId).conversation.model.model?.accountId;
try {
if (accountId && this.isAuthenticationError && this.refreshCredential) {
let acceptance = firstAcceptance;
await this.providerRefresh.withSingleAuthRecovery({
accountId,
operation: async (attempt) => {
if (attempt === 1) {
acceptance = this.pool.startTopLevel({ conversationId, runId, command }).accepted;
}
return await acceptance;
},
isAuthenticationError: this.isAuthenticationError,
refreshCredential: async () => await this.refreshCredential!(accountId),
reopenWorker: async () => {
const recovered = await this.pool.recover(conversationId);
await this.requestHydration(conversationId, recovered, true);
},
});
} else {
await firstAcceptance;
}
const current = this.states.get(conversationId)?.snapshot?.run;
if (current?.runId === runId && current.status === 'queued') {
this.emit(conversationId, {
op: 'run.state',
run: { ...current, status: 'running' },
}, runId);
}
} catch (error) {
this.pool.failTopLevel(
conversationId,
runId,
error instanceof Error ? error : new Error('Prompt acceptance failed'),
);
this.failRun(conversationId, runId, error);
throw error;
}
}
private snapshot(conversationId: string): ConversationSnapshot {
const snapshot = this.states.get(conversationId)?.snapshot;
if (!snapshot) {
throw new CodingRuntimeContractError(
'CODING_CONVERSATION_NOT_FOUND',
'Conversation is not prepared',
true,
);
}
return snapshot;
}
private runtimeState(conversationId: string): ConversationRuntimeState {
const snapshot = this.snapshot(conversationId);
return {
conversationId,
status: snapshot.worker.status,
workerGeneration: snapshot.worker.generation,
...(snapshot.worker.error ? { error: clone(snapshot.worker.error) } : {}),
};
}
private emit(conversationId: string, patch: ConversationPatch, runId?: string): void {
const state = this.states.get(conversationId);
const snapshot = state?.snapshot;
if (!state || !snapshot) return;
const generation = this.pool.getState(conversationId)?.generation
?? snapshot.cursor.workerGeneration;
const envelope: ConversationPatchEnvelope = {
conversationId,
workerGeneration: generation,
...(runId ? { runId } : {}),
seq: snapshot.cursor.seq + 1,
at: this.now(),
patch: clone(patch),
};
const next = reduceConversationPatch(state, envelope);
if (next.invalidation) {
throw new CodingRuntimeContractError(
'CODING_RUNTIME_PROTOCOL_ERROR',
next.invalidation.reason,
true,
);
}
this.states.set(conversationId, next);
for (const listener of this.listeners) listener(clone(envelope));
}
private onPoolEvent(event: PiWorkerPoolEvent): void {
if (event.type === 'worker.replaced') {
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;
const activeRun = this.pool.getActiveRun(event.conversationId);
const continuesActiveRun = Boolean(runId && activeRun?.runId === runId);
if (runId && continuesActiveRun) {
this.extensionUi.replaceGeneration(event.conversationId, event.generation, runId);
}
if (this.extensionHost && runId && continuesActiveRun) {
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;
}
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) {
this.emit(event.conversationId, patch, this.snapshot(event.conversationId).run.runId);
}
if (event.event.type === 'agent_end' || event.event.type === 'agent_settled') {
const worker = this.pool.getState(event.conversationId);
if (worker?.generation === event.generation) {
await this.hydrateGenerationNow(
event.conversationId,
worker,
event.event.type === 'agent_end',
);
}
}
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);
});
}
private failRun(conversationId: string, runId: string, error: unknown): void {
const current = this.states.get(conversationId)?.snapshot?.run;
if (current?.runId !== runId) return;
this.emit(conversationId, {
op: 'run.state',
run: {
status: 'error',
runId,
settledAt: this.now(),
terminalReason: 'failed',
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(
conversationId: string,
workerState: PiWorkerPoolState,
preserveRun: boolean,
): Promise<void> {
const existing = this.hydrationFlights.get(conversationId);
if (existing?.generation === workerState.generation) return existing.flight;
const hydration = this.enqueueProjection(conversationId, async () => {
await this.hydrateGenerationNow(conversationId, workerState, preserveRun);
});
const flight = hydration.finally(() => {
if (this.hydrationFlights.get(conversationId)?.flight === flight) {
this.hydrationFlights.delete(conversationId);
}
});
this.hydrationFlights.set(conversationId, { generation: workerState.generation, flight });
return flight;
}
private async hydrateGenerationNow(
conversationId: string,
workerState: PiWorkerPoolState,
preserveRun: boolean,
): Promise<void> {
const currentWorker = this.pool.getState(conversationId);
if (!currentWorker || currentWorker.generation !== workerState.generation) return;
if (this.snapshot(conversationId).cursor.workerGeneration !== workerState.generation) {
this.replaceWorkerGeneration(conversationId, workerState, preserveRun);
this.resetProjector(conversationId);
}
const [stateResponse, entriesResponse, statsResponse] = await Promise.all([
this.pool.request(conversationId, { type: 'get_state' }, { retry: 'read-only-once' }),
this.pool.request(conversationId, { type: 'get_entries' }, { retry: 'read-only-once' }),
this.pool.request(conversationId, { type: 'get_session_stats' }, { retry: 'read-only-once' }),
]);
const before = this.snapshot(conversationId);
let projected = await projectPiSessionSnapshot({
snapshot: before,
workerGeneration: workerState.generation,
state: stateResponse.data,
entries: entriesResponse.data,
stats: statsResponse.data,
...(this.projectImage ? { projectImage: this.projectImage } : {}),
});
if (preserveRun) {
projected = {
...projected,
run: clone(before.run),
queue: clone(before.queue),
};
}
this.states.set(conversationId, createConversationReducerState(projected));
}
private enqueueProjection(conversationId: string, action: () => Promise<void>): Promise<void> {
const previous = this.projectionChains.get(conversationId) ?? Promise.resolve();
const flight = previous.then(action);
const tail = flight.catch(() => undefined);
this.projectionChains.set(conversationId, tail);
void tail.finally(() => {
if (this.projectionChains.get(conversationId) === tail) {
this.projectionChains.delete(conversationId);
}
});
return flight;
}
private async waitForProjection(conversationId: string): Promise<void> {
await this.projectionChains.get(conversationId);
}
private projector(conversationId: string): PiEventProjector {
const existing = this.projectors.get(conversationId);
if (existing) return existing;
return this.resetProjector(conversationId);
}
private resetProjector(conversationId: string): PiEventProjector {
const projector = new PiEventProjector({
createId: randomUUID,
now: this.now,
...(this.projectImage ? { projectImage: this.projectImage } : {}),
});
this.projectors.set(conversationId, projector);
return projector;
}
private recordProjectionFailure(
conversationId: string,
generation: number,
error: unknown,
): void {
const state = this.states.get(conversationId);
const snapshot = state?.snapshot;
if (!state || !snapshot || snapshot.cursor.workerGeneration !== generation) return;
const publicError = runtimeFailure(error);
this.states.set(conversationId, createConversationReducerState({
...clone(snapshot),
worker: { status: 'error', generation, error: publicError },
run: {
...clone(snapshot.run),
status: 'error',
error: publicError,
},
}));
}
private settleRecoveredRun(conversationId: string): void {
const current = this.snapshot(conversationId);
this.states.set(conversationId, createConversationReducerState({
...clone(current),
run: { status: 'idle' },
queue: { items: [] },
pendingInteractions: [],
}));
}
private replaceWorkerGeneration(
conversationId: string,
workerState: PiWorkerPoolState,
preserveRun: boolean,
): void {
const current = this.snapshot(conversationId);
const next: ConversationSnapshot = {
...clone(current),
run: preserveRun ? clone(current.run) : { status: 'idle' },
queue: preserveRun ? clone(current.queue) : { items: [] },
pendingInteractions: [],
worker: publicWorkerState(workerState),
cursor: {
workerGeneration: workerState.generation,
seq: 0,
...(current.cursor.leafEntryId ? { leafEntryId: current.cursor.leafEntryId } : {}),
},
};
this.states.set(conversationId, createConversationReducerState(next));
}
private replaceModel(conversationId: string, model: ConversationModelState): void {
const state = this.states.get(conversationId);
const snapshot = state?.snapshot;
if (!state || !snapshot) return;
this.states.set(conversationId, createConversationReducerState({
...snapshot,
conversation: { ...snapshot.conversation, model: clone(model) },
}));
const input = this.inputs.get(conversationId);
if (input) input.model = clone(model);
}
private id(kind: RuntimeIdKind): string {
return this.createRuntimeId(kind);
}
}