feat: add Pi conversation worker pool
This commit is contained in:
911
electron/coding-runtime/pi/runtime.ts
Normal file
911
electron/coding-runtime/pi/runtime.ts
Normal file
@@ -0,0 +1,911 @@
|
||||
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 { 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';
|
||||
|
||||
type RuntimeIdKind = 'run' | 'queue';
|
||||
|
||||
export interface PiConversationRuntimeOptions {
|
||||
pool: PiWorkerPool;
|
||||
registry: PiSessionRegistry;
|
||||
resolveModel(model: ProductModelRef): Promise<PiProviderSelection>;
|
||||
resolveImages?(attachments: Array<{ attachmentId: string }>): Promise<unknown[]>;
|
||||
createId?(kind: RuntimeIdKind): string;
|
||||
now?: () => number;
|
||||
providerRefreshCoordinator?: PiProviderRefreshCoordinator;
|
||||
isAuthenticationError?(error: unknown): boolean;
|
||||
refreshCredential?(accountId: string): Promise<void>;
|
||||
}
|
||||
|
||||
export interface PiWorkerProcessAdapter {
|
||||
readonly generation: number;
|
||||
start(): Promise<unknown>;
|
||||
request<T = unknown>(
|
||||
command: PiRpcCommand,
|
||||
options?: PiRpcRequestOptions,
|
||||
): Promise<PiRpcResponse<T>>;
|
||||
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;
|
||||
}
|
||||
|
||||
interface PiRpcSessionStateProjection {
|
||||
sessionId: string;
|
||||
sessionFile?: string;
|
||||
}
|
||||
|
||||
class ManagedPiConversationWorker implements PiConversationWorker {
|
||||
constructor(
|
||||
readonly id: string,
|
||||
readonly generation: number,
|
||||
private readonly process: PiWorkerProcessAdapter,
|
||||
) {}
|
||||
|
||||
request<T = unknown>(command: PiRpcCommand, options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
|
||||
return this.process.request<T>(command, options);
|
||||
}
|
||||
|
||||
subscribe(listener: (event: PiRpcEvent) => void): () => void {
|
||||
return this.process.subscribe(listener);
|
||||
}
|
||||
|
||||
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
|
||||
return this.process.subscribeInvalidation(listener);
|
||||
}
|
||||
|
||||
stop(): Promise<PiWorkerStopResult> {
|
||||
return this.process.stop();
|
||||
}
|
||||
}
|
||||
|
||||
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() }
|
||||
: {}),
|
||||
});
|
||||
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),
|
||||
...(input.fork ? ['--fork', input.fork.sourceSession.piSessionId] : []),
|
||||
'--session-id', sessionKey,
|
||||
],
|
||||
env: credential.env,
|
||||
sensitiveValues: credential.sensitiveValues,
|
||||
});
|
||||
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,
|
||||
),
|
||||
session: clone(bound.session),
|
||||
};
|
||||
} catch (error) {
|
||||
await process.stop().catch(() => undefined);
|
||||
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 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 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 states = new Map<string, ConversationReducerState>();
|
||||
private readonly inputs = new Map<string, PrepareConversationInput>();
|
||||
private readonly listeners = new Set<(patch: ConversationPatchEnvelope) => 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.createRuntimeId = options.createId ?? (() => randomUUID());
|
||||
this.now = options.now ?? Date.now;
|
||||
this.providerRefresh = options.providerRefreshCoordinator ?? new PiProviderRefreshCoordinator();
|
||||
this.isAuthenticationError = options.isAuthenticationError;
|
||||
this.refreshCredential = options.refreshCredential;
|
||||
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));
|
||||
if (!this.states.has(input.conversationId)) {
|
||||
this.states.set(input.conversationId, createConversationReducerState(
|
||||
emptySnapshot(canonicalInput, worker),
|
||||
));
|
||||
} else {
|
||||
this.emit(input.conversationId, { op: 'worker.state', state: publicWorkerState(worker) });
|
||||
}
|
||||
return this.runtimeState(input.conversationId);
|
||||
}
|
||||
|
||||
async getSnapshot(conversationId: string): Promise<ConversationSnapshot> {
|
||||
return clone(this.snapshot(conversationId));
|
||||
}
|
||||
|
||||
async prompt(input: PromptConversationInput): Promise<PromptAcceptance> {
|
||||
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 command: PiRpcCommand = {
|
||||
type: 'prompt',
|
||||
message: input.text,
|
||||
...(images.length > 0 ? { images } : {}),
|
||||
};
|
||||
const ticket = this.pool.startTopLevel({
|
||||
conversationId: input.conversationId,
|
||||
runId,
|
||||
command,
|
||||
});
|
||||
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> {
|
||||
const current = this.snapshot(conversationId).run;
|
||||
this.emit(conversationId, {
|
||||
op: 'run.state',
|
||||
run: { ...current, status: 'aborting' },
|
||||
}, current.runId);
|
||||
try {
|
||||
await this.pool.request(conversationId, { type: 'abort' });
|
||||
} 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> {
|
||||
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.hydrateRecoveredGeneration(input.conversationId, worker, true);
|
||||
} catch (error) {
|
||||
const worker = this.pool.getState(input.conversationId);
|
||||
if (worker) this.replaceWorkerGeneration(input.conversationId, worker, true);
|
||||
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> {
|
||||
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> {
|
||||
const runId = this.id('run');
|
||||
const ticket = this.pool.startTopLevel({
|
||||
conversationId,
|
||||
runId,
|
||||
command: { type: 'compact' },
|
||||
});
|
||||
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> {
|
||||
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));
|
||||
return { conversationId: canonicalInput.conversationId, snapshot: clone(snapshot) };
|
||||
}
|
||||
|
||||
async recover(conversationId: string): Promise<ConversationRuntimeState> {
|
||||
const state = await this.pool.recover(conversationId);
|
||||
await this.hydrateRecoveredGeneration(conversationId, state, false);
|
||||
return this.runtimeState(conversationId);
|
||||
}
|
||||
|
||||
async dispose(conversationId: string): Promise<void> {
|
||||
await this.pool.dispose(conversationId);
|
||||
this.registry.forget(conversationId);
|
||||
this.inputs.delete(conversationId);
|
||||
this.states.delete(conversationId);
|
||||
}
|
||||
|
||||
subscribe(listener: (patch: ConversationPatchEnvelope) => void): () => void {
|
||||
this.listeners.add(listener);
|
||||
return () => this.listeners.delete(listener);
|
||||
}
|
||||
|
||||
async shutdown(): Promise<void> {
|
||||
this.unsubscribePool();
|
||||
await this.pool.shutdown();
|
||||
}
|
||||
|
||||
private async queue(
|
||||
mode: 'steer' | 'follow-up',
|
||||
input: QueueMessageInput,
|
||||
): Promise<QueueAcceptance> {
|
||||
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.hydrateRecoveredGeneration(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.hydrateRecoveredGeneration(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') {
|
||||
this.replaceWorkerGeneration(event.conversationId, event.state, true);
|
||||
return;
|
||||
}
|
||||
if (event.type === 'worker.crashed') {
|
||||
const state = this.pool.getState(event.conversationId);
|
||||
if (state) this.emit(event.conversationId, { op: 'worker.state', state: publicWorkerState(state) });
|
||||
return;
|
||||
}
|
||||
if (event.event.type !== 'agent_settled') return;
|
||||
const current = this.states.get(event.conversationId)?.snapshot?.run;
|
||||
if (!current || current.status === 'idle') return;
|
||||
this.emit(event.conversationId, {
|
||||
op: 'run.state',
|
||||
run: {
|
||||
status: 'idle',
|
||||
...(current.runId ? { runId: current.runId } : {}),
|
||||
settledAt: this.now(),
|
||||
terminalReason: current.status === 'aborting' ? 'aborted' : 'completed',
|
||||
},
|
||||
}, current.runId);
|
||||
this.emit(event.conversationId, { op: 'queue.replace', queue: { items: [] } }, current.runId);
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
private async hydrateRecoveredGeneration(
|
||||
conversationId: string,
|
||||
workerState: PiWorkerPoolState,
|
||||
preserveRun: boolean,
|
||||
): Promise<void> {
|
||||
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.replaceWorkerGeneration(conversationId, workerState, preserveRun);
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user