From 9f55eec9aa52f6acdb295b40e791a302532dfb1d Mon Sep 17 00:00:00 2001 From: brother7 <7brother7@gmail.com> Date: Sat, 22 Aug 2026 23:21:46 +0800 Subject: [PATCH] feat: add Pi conversation worker pool --- .../tasks/20260822-pi-worker-pool-8a4e2c91.md | 78 ++ .../coding-projects/conversation-store.ts | 22 + electron/coding-runtime/pi/runtime.ts | 911 ++++++++++++++++++ .../coding-runtime/pi/session-registry.ts | 139 +++ electron/coding-runtime/pi/telemetry.ts | 55 ++ electron/coding-runtime/pi/worker-pool.ts | 898 +++++++++++++++++ electron/coding-runtime/pi/worker-process.ts | 15 + tests/fixtures/fake-pi-pool-child.mjs | 66 ++ tests/unit/pi-conversation-runtime.test.ts | 321 ++++++ tests/unit/pi-managed-worker-opener.test.ts | 173 ++++ tests/unit/pi-rpc-foundation.test.ts | 3 + tests/unit/pi-runtime-auth-recovery.test.ts | 130 +++ tests/unit/pi-session-registry.test.ts | 95 ++ ...pi-worker-pool-process-integration.test.ts | 163 ++++ tests/unit/pi-worker-pool.test.ts | 531 ++++++++++ 15 files changed, 3600 insertions(+) create mode 100644 .project-docs/30-worklog/tasks/20260822-pi-worker-pool-8a4e2c91.md create mode 100644 electron/coding-runtime/pi/runtime.ts create mode 100644 electron/coding-runtime/pi/session-registry.ts create mode 100644 electron/coding-runtime/pi/telemetry.ts create mode 100644 electron/coding-runtime/pi/worker-pool.ts create mode 100644 tests/fixtures/fake-pi-pool-child.mjs create mode 100644 tests/unit/pi-conversation-runtime.test.ts create mode 100644 tests/unit/pi-managed-worker-opener.test.ts create mode 100644 tests/unit/pi-runtime-auth-recovery.test.ts create mode 100644 tests/unit/pi-session-registry.test.ts create mode 100644 tests/unit/pi-worker-pool-process-integration.test.ts create mode 100644 tests/unit/pi-worker-pool.test.ts diff --git a/.project-docs/30-worklog/tasks/20260822-pi-worker-pool-8a4e2c91.md b/.project-docs/30-worklog/tasks/20260822-pi-worker-pool-8a4e2c91.md new file mode 100644 index 0000000..43ee361 --- /dev/null +++ b/.project-docs/30-worklog/tasks/20260822-pi-worker-pool-8a4e2c91.md @@ -0,0 +1,78 @@ +# Task: Implement PI-050 worker pool and session registry + +## Identity + +- Task ID: 20260822-pi-worker-pool-8a4e2c91 +- Mode: Feature +- Branch: codex/20260822-pi-worker-pool-8a4e2c91-pi-worker-pool +- Worktree: D:\Datas\OthersProjects\makelore-pi-worker-pool-8a4e2c91 +- Base commit: 161f3f471bfdb66fc8d9bb941d9e5e78abe13ea7 +- Owner: codex +- Status: Completed + +## Scope + +- Implement PI-050 on top of PI-040 commit `161f3f471bfdb66fc8d9bb941d9e5e78abe13ea7`. +- Own `electron/coding-runtime/pi/runtime.ts`, worker-pool/session-registry/permit-queue/idle-LRU orchestration, and focused tests. +- Provide per-Conversation worker ownership, lazy prepare/session hydration, caps, fair queueing, stale rebuild, crash/recover/dispose/quit cleanup, and privacy-safe milestone correlation. +- Keep Host API, Renderer, event/session projection, extensions/subagents/write leases, and migration cutover in their later PI tickets. + +## Intent And Constraints + +- Follow planner task `20260822-pi-runtime-spec-b6e2c9a4`, ticket `PI-050 — Worker pool and session registry`. +- A worker may own exactly one product Conversation/Pi session; running workers are never evicted. +- Top-level running cap is 4, warm idle cap is 4, total parent/child process soft-cap integration remains compatible with the later child-worker ticket. +- Prompt/steer/follow-up resolve on RPC acceptance, not `agent_settled`; only `agent_settled` releases the top-level permit. +- Reuse PI-040 credential projection, managed-input args, refresh coordinator, and revision coordinator. Secrets stay in child environment/redaction only. +- Real external Provider validation remains **Explicitly Waived / Accepted Risk** with `realTurnVerified=false`; provider-shaped packaged loopback remains required evidence. macOS validation remains deferred to PI-150. +- The implementation must stay surgical: no Host API/Renderer wiring, no speculative feature flags or compatibility layer, and no incidental OpenCode/template refactor. + +## Outcome + +- Implemented the PI-050 runtime/pool/registry seam with one persistent Pi worker session per Conversation, generation-scoped resources, cap-4 fair top-level scheduling, idle LRU, and a shared default-eight process budget for later child workers. +- Wired the PI-040 managed resource, credential projection, revision, and single-auth-refresh contracts into worker open/reopen. Cross-account model changes rebuild only the target worker so the new Provider credential remains generation-isolated; active runs settle before that rebuild. +- Added targeted recover/dispose/fork/settings/queue/compact/abort orchestration, new-generation `get_state/get_entries` fetch plus snapshot replacement, privacy-safe milestone correlation, and bounded pool shutdown. PI-060 retains ownership of projecting the returned session tree into transcript nodes. +- Prompt/steer/follow-up now distinguish RPC acceptance from authoritative `agent_settled`; rejected queue/abort/compact mutations do not leave false accepted state. Worker crash, stale rebuild, cross-account reconfiguration, and explicit recover publish isolated generations without replaying an uncertain prompt. +- Shutdown rejects new/queued work, waits for in-flight prepare/fork/rebuild openings, cancels generation-scoped command/interaction/child references, and stops all parent workers. Running workers are excluded from LRU eviction. +- Real external Provider validation remains **Explicitly Waived / Accepted Risk** and was not run; `realTurnVerified=false`. macOS remains deferred to PI-150. + +## Verification + +- Final focused suite: 7 files / 34 tests passed, covering prepare/binding single-flight, cap-4 FIFO permits, default-eight shared process budget, running-safe idle LRU, generation resource cleanup, stale revisions, in-flight shutdown, RPC acceptance/rejection, model/thinking target isolation, one-refresh auth recovery, recover/fork/dispose, managed credentials/resources, and a two-real-child-process abort-isolation integration. +- `corepack pnpm run typecheck`: passed. +- `corepack pnpm run lint:check`: passed with 0 errors and 6 pre-existing warnings outside PI-050 files. +- `corepack pnpm test`: 193 files / 2169 tests passed. +- `corepack pnpm run build:vite`: passed for Renderer, Main, Preload, and release utility output; existing dynamic-import and chunk-size warnings remain unchanged. +- `corepack pnpm run test:electron:windows`: 1 file / 3 tests passed. No Host API or Renderer behavior is wired in PI-050, so there was no applicable user-visible Playwright spec to add or run. +- `node scripts/probe-pi-provider-contracts.mjs --timeout-ms 30000`: all four local provider-shaped contracts passed with distinct sessions, overlapping two-worker image turns, target-only abort, environment credential references, and clean stdin-close. The report explicitly retained `realTurnVerified=false` and `realProviderDecision=explicitly-waived-accepted-risk`. +- `node scripts/probe-pi-packaged-runtime.mjs --samples 1 --timeout-ms 30000`: Windows x64 production-shaped closure contained all 130 expected packages and 6 runtime assets; session/reopen/abort/settle/two-worker isolation and all four packaged Provider-shaped contracts passed. The aggregate result intentionally remains `partial-pass` / `incomplete`: Linux aggregation is outside this Windows task, macOS is deferred, and real Provider evidence is waived rather than passed. +- Existing Pi RPC process-tree test passed in the full suite, including forced Windows descendant cleanup; the PI-050 process integration additionally ended both parent workers and returned the shared process budget to zero. +- `git diff --check`: passed before task documentation completion. + +## Follow-ups + +- `PI-060` owns durable `get_entries` active-branch projection, live-event reduction, settled reconciliation, and full recovery hydration. It should consume the generation replacement/fetch seam without moving Pi wire fields into Renderer contracts. +- `PI-080` should receive the same `PiProcessBudget` instance for child workers and register child cancellation through generation resources, preserving the global parent/child soft cap of 8. +- `PI-100` may expose only sanitized runtime/telemetry summaries through Host API; PI-050 worker/session/credential objects remain Main-private. +- Real Provider authentication, protocol/account variation, rate limiting, true-provider concurrency, credential isolation, abort/session/model/image behavior remain explicitly accepted unverified risk with `realTurnVerified=false`. macOS x64/arm64 remains deferred to PI-150. + +## Promotion Candidates + +- Target: Pi runtime orchestration guidance and `PI-080` child scheduler. + Proposal: use one shared FIFO `PiProcessBudget` (default 8) across parent and child worker owners; every successful process open owns one idempotent lease until bounded stop, while pool-specific shutdown cancellation removes only that owner's waiters. + Evidence: focused tests hold a third parent until a lease is released, cancel in-flight opens during shutdown, and verify two real child processes return the budget to zero. + Future impact: PI-080 can enforce the locked global process soft cap without adding a second counter or coupling child scheduling to the parent pool internals. + Semantic conflicts: none found. + Human confirmation required: no. +- Target: Pi credential isolation and model-switch architecture. + Proposal: same-account model changes may use Pi `set_model`; cross-account changes must persist the target model then rebuild only that Conversation worker, deferring rebuild until `agent_settled` when a run is active so each generation receives only its selected account credential environment. + Evidence: runtime tests prove target-only generation replacement and sibling preservation; managed opener tests prove credentials remain in worker environment/redaction rather than argv/catalog/telemetry. + Future impact: preserves PI-040 credential projection when Host API/Renderer model switching arrives, without broadening worker environments to every account secret. + Semantic conflicts: none found. + Human confirmation required: no. +- Target: `PI-060` generation/recovery architecture. + Proposal: every worker replacement publishes a fresh generation Snapshot seam; explicit recover and auth reopen fetch `get_state/get_entries`, clear generation-scoped interactions, and never replay an uncertain prompt. + Evidence: generation-gap regression, one-refresh/two-auth-failure, target recover, stale rebuild, and crash cleanup tests pass. + Future impact: lets the projector own transcript reconstruction while the pool remains authoritative only for process/session lifecycle. + Semantic conflicts: none found. + Human confirmation required: no. diff --git a/electron/coding-projects/conversation-store.ts b/electron/coding-projects/conversation-store.ts index ce86057..31ed3cd 100644 --- a/electron/coding-projects/conversation-store.ts +++ b/electron/coding-projects/conversation-store.ts @@ -213,6 +213,28 @@ export function createCodingConversationStore( }); }, + async setModelState( + conversationId: string, + modelState: ConversationModelState, + ): Promise { + return await mutate(async (file) => { + const current = file.conversations.find((item) => item.id === conversationId); + if (!current) throw new Error('Conversation does not exist'); + const updated = normalizeConversation({ + ...current, + ...modelState, + updatedAt: now(), + }); + return { + result: updated, + file: { + schemaVersion: 2, + conversations: file.conversations.map((item) => item.id === conversationId ? updated : item), + }, + }; + }); + }, + async ensureSessionBinding( conversationId: string, createBinding: () => Promise, diff --git a/electron/coding-runtime/pi/runtime.ts b/electron/coding-runtime/pi/runtime.ts new file mode 100644 index 0000000..a12cd32 --- /dev/null +++ b/electron/coding-runtime/pi/runtime.ts @@ -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; + resolveImages?(attachments: Array<{ attachmentId: string }>): Promise; + createId?(kind: RuntimeIdKind): string; + now?: () => number; + providerRefreshCoordinator?: PiProviderRefreshCoordinator; + isAuthenticationError?(error: unknown): boolean; + refreshCredential?(accountId: string): Promise; +} + +export interface PiWorkerProcessAdapter { + readonly generation: number; + start(): Promise; + request( + command: PiRpcCommand, + options?: PiRpcRequestOptions, + ): Promise>; + subscribe(listener: (event: PiRpcEvent) => void): () => void; + subscribeInvalidation(listener: (error: PiProcessError) => void): () => void; + stop(): Promise; +} + +export interface PiManagedProviderInput { + accounts: ProviderAccount[]; + modelSummaries: ModelSummary[]; +} + +export interface PiManagedWorkerOpenerOptions { + registry: PiSessionRegistry; + executablePath: string; + cliPath: string; + userDataDir: string; + bundledSkillsDir: string; + loadProviderInput(): Promise; + resolveCredential(account: ProviderAccount): Promise; + getLocalProxyCredential?(): Promise; + 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(command: PiRpcCommand, options?: PiRpcRequestOptions): Promise> { + return this.process.request(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 { + return this.process.stop(); + } +} + +export function createPiManagedWorkerOpener( + options: PiManagedWorkerOpenerOptions, +): (input: PiWorkerOpenInput) => Promise { + 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( + { 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, + durationMs: number, + at: number, +): void { + listener?.(createPiRuntimeTelemetryEvent({ + milestone, + conversationId: input.conversation.conversationId, + workerGeneration: input.generation, + cold: true, + durationMs, + at, + })); +} + +function clone(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; + 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) | undefined; + private readonly states = new Map(); + private readonly inputs = new Map(); + 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 { + 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 { + return clone(this.snapshot(conversationId)); + } + + async prompt(input: PromptConversationInput): Promise { + 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 { + return await this.queue('steer', input); + } + + async followUp(input: QueueMessageInput): Promise { + return await this.queue('follow-up', input); + } + + async abort(conversationId: string): Promise { + 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 { + 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 { + 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 { + 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 { + 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 { + const state = await this.pool.recover(conversationId); + await this.hydrateRecoveredGeneration(conversationId, state, false); + return this.runtimeState(conversationId); + } + + async dispose(conversationId: string): Promise { + 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 { + this.unsubscribePool(); + await this.pool.shutdown(); + } + + private async queue( + mode: 'steer' | 'follow-up', + input: QueueMessageInput, + ): Promise { + 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 { + 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, + ): Promise { + 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 { + 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); + } +} diff --git a/electron/coding-runtime/pi/session-registry.ts b/electron/coding-runtime/pi/session-registry.ts new file mode 100644 index 0000000..0042499 --- /dev/null +++ b/electron/coding-runtime/pi/session-registry.ts @@ -0,0 +1,139 @@ +import { + createCodingConversationStore, + type CodingConversationV2, + type PiSessionBinding, +} from '../../coding-projects/conversation-store'; +import { + readCodingProjectConfigV2, + type CodingProjectAgentV2, +} from '../../coding-projects/project-config'; +import type { CodingProjectStore } from '../../coding-projects/project-store'; +import type { + ConversationModelState, + PrepareConversationInput, +} from '../contracts'; + +export interface PiRegisteredConversation { + projectPath: string; + conversation: CodingConversationV2; + agent: CodingProjectAgentV2; + session: PiSessionBinding | null; +} + +export interface PiSessionRegistryOptions { + projectStore: CodingProjectStore; +} + +interface RegistryRecord extends PiRegisteredConversation { + store: ReturnType; +} + +function publicRecord(record: RegistryRecord): PiRegisteredConversation { + return { + projectPath: record.projectPath, + conversation: structuredClone(record.conversation), + agent: structuredClone(record.agent), + session: record.session ? structuredClone(record.session) : null, + }; +} + +function modelStateOf(conversation: CodingConversationV2): ConversationModelState { + return { + model: conversation.model ? structuredClone(conversation.model) : null, + modelResolution: conversation.modelResolution, + }; +} + +function sameModelState(left: ConversationModelState, right: ConversationModelState): boolean { + return left.modelResolution === right.modelResolution + && (left.model === null || right.model === null + ? left.model === right.model + : left.model.accountId === right.model.accountId + && left.model.modelId === right.model.modelId + && left.model.thinkingLevel === right.model.thinkingLevel); +} + +export class PiSessionRegistry { + private readonly projectStore: CodingProjectStore; + private readonly records = new Map(); + private readonly prepareFlights = new Map>(); + + constructor(options: PiSessionRegistryOptions) { + this.projectStore = options.projectStore; + } + + async prepare(input: PrepareConversationInput): Promise { + return publicRecord(await this.prepareRecord(input)); + } + + async ensureBinding( + input: PrepareConversationInput, + createBinding: () => Promise, + ): Promise { + const record = await this.prepareRecord(input); + const conversation = await record.store.ensureSessionBinding(input.conversationId, createBinding); + record.conversation = conversation; + record.session = { + piSessionId: conversation.piSessionId as string, + sessionKey: conversation.sessionKey as string, + }; + return publicRecord(record); + } + + async setModel( + conversationId: string, + model: ConversationModelState, + ): Promise { + const record = this.records.get(conversationId); + if (!record) throw new Error('Conversation is not registered'); + record.conversation = await record.store.setModelState(conversationId, model); + return modelStateOf(record.conversation); + } + + forget(conversationId: string): void { + this.records.delete(conversationId); + } + + private prepareRecord(input: PrepareConversationInput): Promise { + const existing = this.records.get(input.conversationId); + if (existing) return Promise.resolve(existing); + const pending = this.prepareFlights.get(input.conversationId); + if (pending) return pending; + const flight = this.load(input).finally(() => { + if (this.prepareFlights.get(input.conversationId) === flight) { + this.prepareFlights.delete(input.conversationId); + } + }); + this.prepareFlights.set(input.conversationId, flight); + return flight; + } + + private async load(input: PrepareConversationInput): Promise { + const project = (await this.projectStore.listProjects()) + .find((candidate) => candidate.id === input.projectId); + if (!project) throw new Error('Coding project does not exist'); + const configRead = await readCodingProjectConfigV2(project.path); + if (configRead.status !== 'valid') throw new Error('Coding project configuration is unavailable'); + const agent = configRead.config.agents.find((candidate) => candidate.id === input.agentId); + if (!agent) throw new Error('Coding Agent does not exist'); + const store = createCodingConversationStore(project.path); + const conversation = await store.get(input.conversationId); + if (!conversation || conversation.agentId !== input.agentId) { + throw new Error('Coding Conversation does not exist for the selected Agent'); + } + if (!sameModelState(modelStateOf(conversation), input.model)) { + throw new Error('Coding Conversation model metadata is stale'); + } + const record: RegistryRecord = { + projectPath: project.path, + conversation, + agent, + session: conversation.piSessionId && conversation.sessionKey + ? { piSessionId: conversation.piSessionId, sessionKey: conversation.sessionKey } + : null, + store, + }; + this.records.set(input.conversationId, record); + return record; + } +} diff --git a/electron/coding-runtime/pi/telemetry.ts b/electron/coding-runtime/pi/telemetry.ts new file mode 100644 index 0000000..2d0c738 --- /dev/null +++ b/electron/coding-runtime/pi/telemetry.ts @@ -0,0 +1,55 @@ +export const PI_RUNTIME_MILESTONES = [ + 'worker.queue_wait', + 'worker.spawn', + 'rpc.ready', + 'session.open', + 'resources.ready', + 'prompt.accepted', +] as const; + +export type PiRuntimeMilestone = (typeof PI_RUNTIME_MILESTONES)[number]; + +export interface PiRuntimeTelemetryEvent { + milestone: PiRuntimeMilestone; + conversationRef: string; + workerGeneration: number; + spanId: string; + runRef?: string; + cold: boolean; + durationMs: number; + at: number; +} + +export interface PiRuntimeMilestoneInput { + milestone: PiRuntimeMilestone; + conversationId: string; + workerGeneration: number; + runId?: string; + cold: boolean; + durationMs: number; + at: number; +} + +function shortRef(value: string): string { + const normalized = value.replace(/[^A-Za-z0-9]/g, ''); + return normalized.slice(-8).toLowerCase() || 'unknown'; +} + +export function createPiRuntimeTelemetryEvent( + input: PiRuntimeMilestoneInput, +): PiRuntimeTelemetryEvent { + const conversationRef = shortRef(input.conversationId); + const runRef = input.runId ? shortRef(input.runId) : undefined; + return { + milestone: input.milestone, + conversationRef, + workerGeneration: input.workerGeneration, + spanId: runRef + ? `${conversationRef}:${runRef}` + : `${conversationRef}:worker-${input.workerGeneration}`, + ...(runRef ? { runRef } : {}), + cold: input.cold, + durationMs: Math.max(0, Math.round(input.durationMs)), + at: input.at, + }; +} diff --git a/electron/coding-runtime/pi/worker-pool.ts b/electron/coding-runtime/pi/worker-pool.ts new file mode 100644 index 0000000..c202ce7 --- /dev/null +++ b/electron/coding-runtime/pi/worker-pool.ts @@ -0,0 +1,898 @@ +import type { PrepareConversationInput } from '../contracts'; +import type { ConversationModelState } from '../contracts'; +import type { PiProcessError, PiProcessErrorCode } from './process-errors'; +import type { + PiRpcCommand, + PiRpcEvent, + PiRpcRequestOptions, + PiRpcResponse, +} from './rpc-client'; +import type { PiWorkerStopResult } from './worker-process'; +import { + PiManagedInputRevisionCoordinator, + type PiManagedInputRevision, +} from './managed-input-revision'; +import { + createPiRuntimeTelemetryEvent, + type PiRuntimeTelemetryEvent, +} from './telemetry'; + +export interface PiConversationWorker { + readonly id: string; + readonly generation: number; + request( + command: PiRpcCommand, + options?: PiRpcRequestOptions, + ): Promise>; + subscribe(listener: (event: PiRpcEvent) => void): () => void; + subscribeInvalidation(listener: (error: PiProcessError) => void): () => void; + stop(): Promise; +} + +export interface PiWorkerSessionBinding { + piSessionId: string; + sessionKey: string; +} + +export interface PiWorkerOpenInput { + conversation: PrepareConversationInput; + generation: number; + revision: PiManagedInputRevision; + existingSession?: PiWorkerSessionBinding; + fork?: { + sourceSession: PiWorkerSessionBinding; + sourceEntryId?: string; + }; +} + +export interface PiWorkerOpenResult { + worker: PiConversationWorker; + session: PiWorkerSessionBinding; +} + +export interface PiWorkerPoolState { + conversationId: string; + workerId: string; + state: 'spawning' | 'ready' | 'queued' | 'running' | 'idle' | 'crashed'; + generation: number; + session: PiWorkerSessionBinding; + failureCode?: PiProcessErrorCode; +} + +export interface PiTopLevelRunInput { + conversationId: string; + runId: string; + command: PiRpcCommand; +} + +export interface PiTopLevelRunTicket { + queuePosition?: number; + accepted: Promise; +} + +export type PiGenerationResourceKind = 'command' | 'interaction' | 'child'; + +export interface PiGenerationResourceInput { + conversationId: string; + kind: PiGenerationResourceKind; + id: string; + cancel(): void; +} + +export type PiWorkerPoolEvent = + | { + type: 'worker.event'; + conversationId: string; + generation: number; + event: PiRpcEvent; + } + | { + type: 'worker.crashed'; + conversationId: string; + generation: number; + error: PiProcessError; + } + | { + type: 'worker.replaced'; + conversationId: string; + generation: number; + state: PiWorkerPoolState; + }; + +export interface PiWorkerPoolOptions { + openWorker(input: PiWorkerOpenInput): Promise; + maxRunning?: number; + maxIdle?: number; + processBudget?: PiProcessBudget; + revisionCoordinator?: PiManagedInputRevisionCoordinator; + now?: () => number; + onTelemetry?: (event: PiRuntimeTelemetryEvent) => void; +} + +export interface PiProcessLease { + release(): void; +} + +interface PiProcessBudgetWaiter { + signal?: AbortSignal; + resolve(lease: PiProcessLease): void; + reject(error: Error): void; + onAbort?: () => void; +} + +export class PiProcessBudget { + private readonly waiters: PiProcessBudgetWaiter[] = []; + private active = 0; + + constructor(readonly maxProcesses = 8) { + if (!Number.isSafeInteger(maxProcesses) || maxProcesses <= 0) { + throw new Error('maxProcesses must be a positive safe integer'); + } + } + + get activeCount(): number { return this.active; } + get waitingCount(): number { return this.waiters.length; } + + acquire(signal?: AbortSignal): Promise { + if (signal?.aborted) return Promise.reject(new Error('Pi process budget acquisition cancelled')); + if (this.active < this.maxProcesses) return Promise.resolve(this.issueLease()); + return new Promise((resolve, reject) => { + const waiter: PiProcessBudgetWaiter = { resolve, reject, ...(signal ? { signal } : {}) }; + if (signal) { + waiter.onAbort = () => { + const index = this.waiters.indexOf(waiter); + if (index >= 0) this.waiters.splice(index, 1); + reject(new Error('Pi process budget acquisition cancelled')); + }; + signal.addEventListener('abort', waiter.onAbort, { once: true }); + } + this.waiters.push(waiter); + }); + } + + private issueLease(): PiProcessLease { + this.active += 1; + let released = false; + return { + release: () => { + if (released) return; + released = true; + this.active -= 1; + this.advance(); + }, + }; + } + + private advance(): void { + while (this.active < this.maxProcesses && this.waiters.length > 0) { + const waiter = this.waiters.shift() as PiProcessBudgetWaiter; + if (waiter.onAbort && waiter.signal) { + waiter.signal.removeEventListener('abort', waiter.onAbort); + } + if (waiter.signal?.aborted) { + waiter.reject(new Error('Pi process budget acquisition cancelled')); + continue; + } + waiter.resolve(this.issueLease()); + } + } +} + +interface WorkerRecord { + conversation: PrepareConversationInput; + worker: PiConversationWorker; + session: PiWorkerSessionBinding; + generation: number; + revisionWorkerId: string; + lastUsed: number; + state: PiWorkerPoolState['state']; + unsubscribeEvent: () => void; + unsubscribeInvalidation: () => void; + failureCode?: PiProcessErrorCode; + generationResources: Record void>>; + rebuildFlight?: Promise; + acceptedPromptCount: number; + processLease: PiProcessLease | null; + processStopFlight?: Promise; + reconfigureAfterSettled: boolean; +} + +interface PendingTopLevelRun extends PiTopLevelRunInput { + resolve(response: PiRpcResponse): void; + reject(error: unknown): void; + queuedAt?: number; +} + +export class PiWorkerPool { + private readonly openWorker: PiWorkerPoolOptions['openWorker']; + private readonly maxRunning: number; + private readonly maxIdle: number; + private readonly processBudget: PiProcessBudget; + private readonly revisions: PiManagedInputRevisionCoordinator; + private readonly now: () => number; + private readonly onTelemetry: ((event: PiRuntimeTelemetryEvent) => void) | undefined; + private readonly workers = new Map(); + private readonly prepareFlights = new Map>(); + private readonly rebuildFlights = new Set>(); + private readonly waitingRuns: PendingTopLevelRun[] = []; + private readonly activeRuns = new Map(); + private readonly generations = new Map(); + private readonly listeners = new Set<(event: PiWorkerPoolEvent) => void>(); + private commandSequence = 0; + private runningCount = 0; + private useSequence = 0; + private shuttingDown = false; + private shutdownFlight: Promise | null = null; + private readonly shutdownController = new AbortController(); + + constructor(options: PiWorkerPoolOptions) { + this.openWorker = options.openWorker; + this.maxRunning = options.maxRunning ?? 4; + this.maxIdle = options.maxIdle ?? 4; + this.processBudget = options.processBudget ?? new PiProcessBudget(); + this.revisions = options.revisionCoordinator ?? new PiManagedInputRevisionCoordinator(); + this.now = options.now ?? Date.now; + this.onTelemetry = options.onTelemetry; + if (!Number.isSafeInteger(this.maxRunning) || this.maxRunning <= 0) { + throw new Error('maxRunning must be a positive safe integer'); + } + if (!Number.isSafeInteger(this.maxIdle) || this.maxIdle < 0) { + throw new Error('maxIdle must be a non-negative safe integer'); + } + } + + prepare(conversation: PrepareConversationInput): Promise { + if (this.shuttingDown) return Promise.reject(new Error('Pi worker pool is shutting down')); + const existing = this.workers.get(conversation.conversationId); + if (existing) { + this.touch(existing); + void this.trimIdleWorkers(); + return Promise.resolve(this.publicState(existing)); + } + const pending = this.prepareFlights.get(conversation.conversationId); + if (pending) return pending; + + const generation = this.nextGeneration(conversation.conversationId); + const revision = this.revisions.current; + const flight = this.openWithLease({ + conversation: structuredClone(conversation), + generation, + revision, + }) + .then(async ({ opened: { worker, session }, lease }) => { + const record = this.createRecord(conversation, worker, session, generation, revision, lease); + this.workers.set(conversation.conversationId, record); + const state = this.publicState(record); + await this.trimIdleWorkers(); + return state; + }) + .finally(() => { + if (this.prepareFlights.get(conversation.conversationId) === flight) { + this.prepareFlights.delete(conversation.conversationId); + } + }); + this.prepareFlights.set(conversation.conversationId, flight); + return flight; + } + + markProviderStale(): PiManagedInputRevision { + return this.revisions.markProviderStale(); + } + + markResourcesStale(): PiManagedInputRevision { + return this.revisions.markResourcesStale(); + } + + async fork( + sourceConversationId: string, + conversation: PrepareConversationInput, + sourceEntryId?: string, + ): Promise { + if (this.shuttingDown) throw new Error('Pi worker pool is shutting down'); + const source = this.workers.get(sourceConversationId); + if (!source || source.state === 'crashed') throw new Error('Source Conversation worker is not available'); + if (this.workers.has(conversation.conversationId) + || this.prepareFlights.has(conversation.conversationId)) { + throw new Error('Fork target Conversation worker already exists'); + } + const generation = this.nextGeneration(conversation.conversationId); + const revision = this.revisions.current; + const flight = (async () => { + const { opened, lease } = await this.openWithLease({ + conversation: structuredClone(conversation), + generation, + revision, + fork: { + sourceSession: structuredClone(source.session), + ...(sourceEntryId ? { sourceEntryId } : {}), + }, + }); + const record = this.createRecord( + conversation, + opened.worker, + opened.session, + generation, + revision, + lease, + ); + this.workers.set(conversation.conversationId, record); + await this.trimIdleWorkers(); + return this.publicState(record); + })().finally(() => { + if (this.prepareFlights.get(conversation.conversationId) === flight) { + this.prepareFlights.delete(conversation.conversationId); + } + }); + this.prepareFlights.set(conversation.conversationId, flight); + return await flight; + } + + getState(conversationId: string): PiWorkerPoolState | null { + const record = this.workers.get(conversationId); + return record ? this.publicState(record) : null; + } + + subscribe(listener: (event: PiWorkerPoolEvent) => void): () => void { + this.listeners.add(listener); + return () => this.listeners.delete(listener); + } + + async request( + conversationId: string, + command: PiRpcCommand, + options: PiRpcRequestOptions = {}, + ): Promise> { + let record = this.workers.get(conversationId); + if (!record || record.state === 'crashed') throw new Error('Conversation worker is not available'); + if (record.rebuildFlight) record = await record.rebuildFlight; + const controller = new AbortController(); + const abort = () => controller.abort(); + options.signal?.addEventListener('abort', abort, { once: true }); + const untrack = this.trackGenerationResource({ + conversationId, + kind: 'command', + id: `command-${++this.commandSequence}`, + cancel: abort, + }); + try { + return await record.worker.request(command, { + ...options, + signal: controller.signal, + }); + } finally { + options.signal?.removeEventListener('abort', abort); + untrack(); + } + } + + failTopLevel(conversationId: string, runId: string, error: Error): void { + const active = this.activeRuns.get(conversationId); + if (active?.runId === runId) { + this.activeRuns.delete(conversationId); + this.runningCount -= 1; + const record = this.workers.get(conversationId); + if (record && record.state !== 'crashed') { + record.state = 'idle'; + try { + this.revisions.settleRun(record.revisionWorkerId); + } catch { + // Recover/crash may already have removed this generation. + } + } + } + for (let index = this.waitingRuns.length - 1; index >= 0; index -= 1) { + const pending = this.waitingRuns[index]; + if (pending?.conversationId !== conversationId || pending.runId !== runId) continue; + this.waitingRuns.splice(index, 1); + pending.reject(error); + } + this.launchWaitingRuns(); + } + + updateConversationModel(conversationId: string, model: ConversationModelState): void { + const record = this.workers.get(conversationId); + if (!record) throw new Error('Conversation worker is not prepared'); + record.conversation = { + ...record.conversation, + model: structuredClone(model), + }; + } + + async reconfigureConversationModel( + conversationId: string, + model: ConversationModelState, + ): Promise { + if (this.shuttingDown) throw new Error('Pi worker pool is shutting down'); + let record = this.workers.get(conversationId); + if (!record) throw new Error('Conversation worker is not prepared'); + if (record.rebuildFlight) record = await record.rebuildFlight; + record.conversation = { + ...record.conversation, + model: structuredClone(model), + }; + if (this.activeRuns.has(conversationId)) { + record.reconfigureAfterSettled = true; + return null; + } + record.rebuildFlight = this.beginRebuild(record, this.revisions.current); + return this.publicState(await record.rebuildFlight); + } + + trackGenerationResource(input: PiGenerationResourceInput): () => void { + const record = this.workers.get(input.conversationId); + if (!record || record.state === 'crashed') { + throw new Error('Conversation worker generation is not available'); + } + const resources = record.generationResources[input.kind]; + if (resources.has(input.id)) throw new Error(`Duplicate ${input.kind} resource: ${input.id}`); + resources.set(input.id, input.cancel); + return () => { + if (this.workers.get(input.conversationId) === record) resources.delete(input.id); + }; + } + + startTopLevel(input: PiTopLevelRunInput): PiTopLevelRunTicket { + if (this.shuttingDown) throw new Error('Pi worker pool is shutting down'); + const record = this.workers.get(input.conversationId); + if (!record) throw new Error('Conversation worker is not prepared'); + if (this.activeRuns.has(input.conversationId) + || this.waitingRuns.some((run) => run.conversationId === input.conversationId)) { + throw new Error('Conversation already has a top-level run'); + } + + let resolve!: (response: PiRpcResponse) => void; + let reject!: (error: unknown) => void; + const accepted = new Promise((done, fail) => { + resolve = done; + reject = fail; + }); + const pending: PendingTopLevelRun = { + ...structuredClone(input), + resolve, + reject, + }; + if (this.runningCount < this.maxRunning) { + this.launchTopLevel(record, pending); + return { accepted }; + } + this.waitingRuns.push(pending); + pending.queuedAt = this.now(); + record.state = 'queued'; + return { queuePosition: this.waitingRuns.length, accepted }; + } + + shutdown(): Promise { + if (!this.shutdownFlight) this.shutdownFlight = this.performShutdown(); + return this.shutdownFlight; + } + + async recover(conversationId: string): Promise { + if (this.shuttingDown) throw new Error('Pi worker pool is shutting down'); + const pendingPrepare = this.prepareFlights.get(conversationId); + if (pendingPrepare) await pendingPrepare; + const record = this.workers.get(conversationId); + if (!record) throw new Error('Conversation worker is not prepared'); + this.cancelConversationRuns(conversationId, new Error('Conversation worker is recovering')); + record.rebuildFlight = this.beginRebuild(record, this.revisions.current); + return this.publicState(await record.rebuildFlight); + } + + async dispose(conversationId: string): Promise { + const pendingPrepare = this.prepareFlights.get(conversationId); + if (pendingPrepare) await pendingPrepare.catch(() => undefined); + let record = this.workers.get(conversationId); + if (!record) return; + if (record.rebuildFlight) { + record = await record.rebuildFlight.catch(() => record as WorkerRecord); + } + this.cancelConversationRuns(conversationId, new Error('Conversation worker was disposed')); + record.unsubscribeEvent(); + record.unsubscribeInvalidation(); + this.cancelGenerationResources(record); + this.revisions.removeWorker(record.revisionWorkerId); + if (this.workers.get(conversationId) === record) this.workers.delete(conversationId); + await this.stopAndRelease(record); + } + + private async performShutdown(): Promise { + this.shuttingDown = true; + this.shutdownController.abort(); + const error = new Error('Pi worker pool is shutting down'); + for (const pending of this.waitingRuns.splice(0)) pending.reject(error); + await Promise.allSettled([...this.prepareFlights.values()]); + await Promise.allSettled([...this.rebuildFlights]); + const records = [...this.workers.values()]; + this.workers.clear(); + this.activeRuns.clear(); + this.runningCount = 0; + await Promise.all(records.map(async (record) => { + record.unsubscribeEvent(); + record.unsubscribeInvalidation(); + this.cancelGenerationResources(record); + this.revisions.removeWorker(record.revisionWorkerId); + await this.stopAndRelease(record); + })); + } + + private launchTopLevel(record: WorkerRecord, run: PendingTopLevelRun): void { + record.state = 'running'; + this.touch(record); + this.runningCount += 1; + this.activeRuns.set(run.conversationId, { + runId: run.runId, + generation: record.generation, + }); + void this.acceptTopLevel(record, run); + } + + private async acceptTopLevel(record: WorkerRecord, run: PendingTopLevelRun): Promise { + let current: WorkerRecord | undefined; + let beganRun = false; + try { + current = await this.ensureFresh(record); + const active = this.activeRuns.get(run.conversationId); + if (!active) throw new Error('Top-level run was cancelled before acceptance'); + active.generation = current.generation; + current.state = 'running'; + this.revisions.beginRun(current.revisionWorkerId); + beganRun = true; + if (run.queuedAt !== undefined) { + this.recordMilestone(current, run, 'worker.queue_wait', this.now() - run.queuedAt); + } + const acceptedAt = this.now(); + const response = await current.worker.request(run.command); + if (run.command.type === 'prompt') { + this.recordMilestone(current, run, 'prompt.accepted', this.now() - acceptedAt); + current.acceptedPromptCount += 1; + } + run.resolve(response); + } catch (error) { + const active = this.activeRuns.get(run.conversationId); + if (active) { + this.activeRuns.delete(run.conversationId); + this.runningCount -= 1; + this.launchWaitingRuns(); + } + if (current && this.workers.get(run.conversationId) === current + && current.state !== 'crashed') { + current.state = 'idle'; + if (beganRun) this.revisions.settleRun(current.revisionWorkerId); + void this.trimIdleWorkers(); + } + run.reject(error); + } + } + + private settleTopLevel(record: WorkerRecord): void { + const conversationId = record.conversation.conversationId; + const active = this.activeRuns.get(conversationId); + if (!active || active.generation !== record.generation) return; + this.activeRuns.delete(conversationId); + this.runningCount -= 1; + record.state = 'idle'; + const revisionAction = this.revisions.settleRun(record.revisionWorkerId); + if (!this.shuttingDown + && (record.reconfigureAfterSettled || revisionAction.action === 'rebuild-after-settled')) { + record.rebuildFlight = this.beginRebuild( + record, + revisionAction.action === 'rebuild-after-settled' + ? revisionAction.revision + : this.revisions.current, + ); + void record.rebuildFlight.catch(() => undefined); + } + this.launchWaitingRuns(); + void this.trimIdleWorkers(); + } + + private async ensureFresh(record: WorkerRecord): Promise { + const current = this.workers.get(record.conversation.conversationId); + if (!current) throw new Error('Conversation worker is no longer available'); + if (current !== record) return await this.ensureFresh(current); + if (record.rebuildFlight) return await record.rebuildFlight; + const action = this.revisions.beforePrompt(record.revisionWorkerId); + if (action.action !== 'rebuild-before-prompt') return record; + record.rebuildFlight = this.beginRebuild(record, action.revision); + return await record.rebuildFlight; + } + + private beginRebuild( + record: WorkerRecord, + revision: PiManagedInputRevision, + ): Promise { + const flight = this.rebuild(record, revision).finally(() => { + this.rebuildFlights.delete(flight); + }); + this.rebuildFlights.add(flight); + return flight; + } + + private async rebuild( + record: WorkerRecord, + revision: PiManagedInputRevision, + ): Promise { + const conversationId = record.conversation.conversationId; + record.state = 'spawning'; + record.unsubscribeEvent(); + record.unsubscribeInvalidation(); + this.cancelGenerationResources(record); + this.revisions.removeWorker(record.revisionWorkerId); + try { + if (record.processStopFlight) await record.processStopFlight; + else await record.worker.stop(); + let lease = record.processLease; + if (!lease) lease = await this.processBudget.acquire(this.shutdownController.signal); + if (this.shuttingDown) { + lease.release(); + if (record.processLease === lease) record.processLease = null; + throw new Error('Pi worker pool is shutting down'); + } + const generation = this.nextGeneration(conversationId); + let opened: PiWorkerOpenResult; + try { + opened = await this.openWorker({ + conversation: structuredClone(record.conversation), + generation, + revision: structuredClone(revision), + existingSession: structuredClone(record.session), + }); + if (this.shuttingDown) { + await opened.worker.stop().catch(() => undefined); + lease.release(); + if (record.processLease === lease) record.processLease = null; + throw new Error('Pi worker pool is shutting down'); + } + } catch (error) { + lease.release(); + if (record.processLease === lease) record.processLease = null; + throw error; + } + if (opened.session.piSessionId !== record.session.piSessionId + || opened.session.sessionKey !== record.session.sessionKey) { + await opened.worker.stop().catch(() => undefined); + lease.release(); + if (record.processLease === lease) record.processLease = null; + throw new Error('Reopened Pi worker returned a different session binding'); + } + record.processLease = null; + const replacement = this.createRecord( + record.conversation, + opened.worker, + opened.session, + generation, + revision, + lease, + ); + replacement.state = record.state === 'spawning' && this.activeRuns.has(conversationId) + ? 'running' + : 'ready'; + this.workers.set(conversationId, replacement); + this.emit({ + type: 'worker.replaced', + conversationId, + generation, + state: this.publicState(replacement), + }); + return replacement; + } catch (error) { + record.state = 'crashed'; + record.failureCode = error && typeof error === 'object' && 'code' in error + ? (error as { code: PiProcessErrorCode }).code + : 'PI_WORKER_START_FAILED'; + throw error; + } + } + + private createRecord( + conversation: PrepareConversationInput, + worker: PiConversationWorker, + session: PiWorkerSessionBinding, + generation: number, + revision: PiManagedInputRevision, + processLease: PiProcessLease, + ): WorkerRecord { + const revisionWorkerId = `${conversation.conversationId}:${generation}`; + this.revisions.registerWorker(revisionWorkerId, revision); + const record: WorkerRecord = { + conversation: structuredClone(conversation), + worker, + session: structuredClone(session), + generation, + revisionWorkerId, + lastUsed: ++this.useSequence, + state: 'ready', + unsubscribeEvent: () => undefined, + unsubscribeInvalidation: () => undefined, + generationResources: { + command: new Map(), + interaction: new Map(), + child: new Map(), + }, + acceptedPromptCount: 0, + processLease, + reconfigureAfterSettled: false, + }; + record.unsubscribeEvent = worker.subscribe((event) => { + if (event.type === 'agent_settled') this.settleTopLevel(record); + this.emit({ + type: 'worker.event', + conversationId: conversation.conversationId, + generation, + event: structuredClone(event), + }); + }); + record.unsubscribeInvalidation = worker.subscribeInvalidation((error) => { + this.handleInvalidation(record, error); + }); + return record; + } + + private nextGeneration(conversationId: string): number { + const generation = (this.generations.get(conversationId) ?? 0) + 1; + this.generations.set(conversationId, generation); + return generation; + } + + private touch(record: WorkerRecord): void { + record.lastUsed = ++this.useSequence; + } + + private async trimIdleWorkers(): Promise { + const idle = [...this.workers.values()] + .filter((record) => record.state === 'ready' || record.state === 'idle') + .sort((left, right) => left.lastUsed - right.lastUsed); + const excess = idle.length - this.maxIdle; + if (excess <= 0) return; + await Promise.all(idle.slice(0, excess).map((record) => this.evict(record))); + } + + private async evict(record: WorkerRecord): Promise { + const conversationId = record.conversation.conversationId; + if (this.workers.get(conversationId) !== record) return; + record.unsubscribeEvent(); + record.unsubscribeInvalidation(); + this.cancelGenerationResources(record); + this.revisions.removeWorker(record.revisionWorkerId); + this.workers.delete(conversationId); + await this.stopAndRelease(record); + } + + private launchWaitingRuns(): void { + while (this.runningCount < this.maxRunning && this.waitingRuns.length > 0) { + const pending = this.waitingRuns.shift() as PendingTopLevelRun; + const record = this.workers.get(pending.conversationId); + if (!record) { + pending.reject(new Error('Conversation worker is no longer available')); + continue; + } + this.launchTopLevel(record, pending); + } + } + + private handleInvalidation(record: WorkerRecord, error: PiProcessError): void { + const conversationId = record.conversation.conversationId; + if (this.workers.get(conversationId) !== record || record.state === 'crashed') return; + record.state = 'crashed'; + record.failureCode = error.code; + record.unsubscribeEvent(); + this.revisions.removeWorker(record.revisionWorkerId); + this.cancelGenerationResources(record); + if (!record.processStopFlight) { + record.processStopFlight = this.stopAndRelease(record).catch(() => undefined); + } + + const active = this.activeRuns.get(conversationId); + if (active?.generation === record.generation) { + this.activeRuns.delete(conversationId); + this.runningCount -= 1; + } + for (let index = this.waitingRuns.length - 1; index >= 0; index -= 1) { + const pending = this.waitingRuns[index]; + if (pending?.conversationId !== conversationId) continue; + this.waitingRuns.splice(index, 1); + pending.reject(error); + } + this.launchWaitingRuns(); + this.emit({ + type: 'worker.crashed', + conversationId, + generation: record.generation, + error, + }); + } + + private cancelConversationRuns(conversationId: string, error: Error): void { + const active = this.activeRuns.get(conversationId); + if (active) { + this.activeRuns.delete(conversationId); + this.runningCount -= 1; + } + for (let index = this.waitingRuns.length - 1; index >= 0; index -= 1) { + const pending = this.waitingRuns[index]; + if (pending?.conversationId !== conversationId) continue; + this.waitingRuns.splice(index, 1); + pending.reject(error); + } + this.launchWaitingRuns(); + } + + private cancelGenerationResources(record: WorkerRecord): void { + for (const resources of Object.values(record.generationResources)) { + for (const cancel of resources.values()) { + try { + cancel(); + } catch { + // A failed cleanup must not prevent sibling resources from being cancelled. + } + } + resources.clear(); + } + } + + private async openWithLease( + input: PiWorkerOpenInput, + ): Promise<{ opened: PiWorkerOpenResult; lease: PiProcessLease }> { + const lease = await this.processBudget.acquire(this.shutdownController.signal); + if (this.shuttingDown) { + lease.release(); + throw new Error('Pi worker pool is shutting down'); + } + try { + const opened = await this.openWorker(input); + if (this.shuttingDown) { + await opened.worker.stop().catch(() => undefined); + lease.release(); + throw new Error('Pi worker pool is shutting down'); + } + return { opened, lease }; + } catch (error) { + lease.release(); + throw error; + } + } + + private async stopAndRelease(record: WorkerRecord): Promise { + await record.worker.stop(); + record.processLease?.release(); + record.processLease = null; + } + + private emit(event: PiWorkerPoolEvent): void { + for (const listener of this.listeners) { + try { + listener(event); + } catch { + // Runtime observers must not affect worker lifecycle. + } + } + } + + private recordMilestone( + record: WorkerRecord, + run: PendingTopLevelRun, + milestone: 'worker.queue_wait' | 'prompt.accepted', + durationMs: number, + ): void { + if (!this.onTelemetry) return; + this.onTelemetry(createPiRuntimeTelemetryEvent({ + milestone, + conversationId: record.conversation.conversationId, + workerGeneration: record.generation, + runId: run.runId, + cold: record.acceptedPromptCount === 0, + durationMs, + at: this.now(), + })); + } + + private publicState(record: WorkerRecord): PiWorkerPoolState { + return { + conversationId: record.conversation.conversationId, + workerId: record.worker.id, + state: record.state, + generation: record.generation, + session: structuredClone(record.session), + ...(record.failureCode ? { failureCode: record.failureCode } : {}), + }; + } +} diff --git a/electron/coding-runtime/pi/worker-process.ts b/electron/coding-runtime/pi/worker-process.ts index 36e1d5b..f7b96d2 100644 --- a/electron/coding-runtime/pi/worker-process.ts +++ b/electron/coding-runtime/pi/worker-process.ts @@ -177,6 +177,7 @@ export class PiWorkerProcess { private stopping = false; private stopPromise: Promise | null = null; private exitResult: Promise<{ code: number | null; signal: NodeJS.Signals | null }> | null = null; + private readonly invalidationListeners = new Set<(error: PiProcessError) => void>(); constructor(options: PiWorkerProcessOptions) { this.options = options; @@ -314,6 +315,11 @@ export class PiWorkerProcess { return this.rpc.subscribe(listener); } + subscribeInvalidation(listener: (error: PiProcessError) => void): () => void { + this.invalidationListeners.add(listener); + return () => this.invalidationListeners.delete(listener); + } + stop(): Promise { if (!this.stopPromise) this.stopPromise = this.performStop(); return this.stopPromise; @@ -345,6 +351,15 @@ export class PiWorkerProcess { this.invalidation = error; this.rpc?.invalidate(error); this.generationValue += 1; + for (const listener of this.invalidationListeners) { + try { + listener(error); + } catch (listenerError) { + this.appendDiagnostic( + `[invalidation-listener] ${listenerError instanceof Error ? listenerError.message : String(listenerError)}\n`, + ); + } + } } private appendDiagnostic(source: string): void { diff --git a/tests/fixtures/fake-pi-pool-child.mjs b/tests/fixtures/fake-pi-pool-child.mjs new file mode 100644 index 0000000..b2adbb4 --- /dev/null +++ b/tests/fixtures/fake-pi-pool-child.mjs @@ -0,0 +1,66 @@ +let input = ''; +let activeTimer = null; +let activeMarker = null; + +function write(record) { + process.stdout.write(`${JSON.stringify(record)}\n`); +} + +function respond(command, data = {}) { + write({ + type: 'response', + id: command.id, + command: command.type, + success: true, + data, + }); +} + +function settle(reason) { + if (!activeMarker) return; + write({ type: 'agent_end', marker: activeMarker, reason }); + write({ type: 'agent_settled', marker: activeMarker, reason }); + activeMarker = null; + activeTimer = null; +} + +function handle(command) { + if (command.type === 'prompt') { + activeMarker = command.message; + respond(command, { accepted: true }); + write({ type: 'agent_start', marker: activeMarker }); + activeTimer = setTimeout(() => settle('completed'), Number(command.delayMs ?? 100)); + return; + } + if (command.type === 'abort') { + if (activeTimer) clearTimeout(activeTimer); + respond(command, { aborted: Boolean(activeMarker) }); + settle('aborted'); + return; + } + if (command.type === 'get_state') { + respond(command, { sessionId: `session-${process.pid}` }); + return; + } + if (command.type === 'get_entries') { + respond(command, { entries: [] }); + return; + } + respond(command); +} + +process.stdin.setEncoding('utf8'); +process.stdin.on('data', (chunk) => { + input += chunk; + while (true) { + const newline = input.indexOf('\n'); + if (newline === -1) break; + const line = input.slice(0, newline).replace(/\r$/, ''); + input = input.slice(newline + 1); + if (line) handle(JSON.parse(line)); + } +}); +process.stdin.on('end', () => { + if (activeTimer) clearTimeout(activeTimer); + process.exit(0); +}); diff --git a/tests/unit/pi-conversation-runtime.test.ts b/tests/unit/pi-conversation-runtime.test.ts new file mode 100644 index 0000000..70aeb94 --- /dev/null +++ b/tests/unit/pi-conversation-runtime.test.ts @@ -0,0 +1,321 @@ +// @vitest-environment node + +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; +import { afterEach, describe, expect, it } from 'vitest'; +import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store'; +import { createCodingProjectAgent } from '../../electron/coding-projects/project-config'; +import { + createCodingProjectStore, + createLocalCodingProject, + createMemoryCodingProjectStorage, +} from '../../electron/coding-projects/project-store'; +import { PiConversationRuntime } from '../../electron/coding-runtime/pi/runtime'; +import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry'; +import { + PiWorkerPool, + type PiConversationWorker, +} from '../../electron/coding-runtime/pi/worker-pool'; +import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors'; +import type { + PiRpcCommand, + PiRpcEvent, + PiRpcRequestOptions, + PiRpcResponse, +} from '../../electron/coding-runtime/pi/rpc-client'; + +const roots: string[] = []; +const NOW = '2026-08-22T15:00:00.000Z'; + +class RuntimeFakeWorker implements PiConversationWorker { + readonly requests: PiRpcCommand[] = []; + private failType: string | null = null; + private readonly responseGates = new Map>(); + private readonly events = new Set<(event: PiRpcEvent) => void>(); + private readonly invalidations = new Set<(error: PiProcessError) => void>(); + + constructor(readonly id: string, readonly generation: number) {} + + async request( + command: PiRpcCommand, + _options?: PiRpcRequestOptions, + ): Promise> { + this.requests.push(structuredClone(command)); + await this.responseGates.get(command.type); + this.responseGates.delete(command.type); + if (command.type === this.failType) { + this.failType = null; + throw new Error(`fake ${command.type} rejection`); + } + return { type: 'response', id: `${this.id}-${this.requests.length}`, success: true }; + } + + failNext(type: string): void { this.failType = type; } + + holdNext(type: string): () => void { + let release!: () => void; + this.responseGates.set(type, new Promise((resolve) => { release = resolve; })); + return release; + } + + subscribe(listener: (event: PiRpcEvent) => void): () => void { + this.events.add(listener); + return () => this.events.delete(listener); + } + + subscribeInvalidation(listener: (error: PiProcessError) => void): () => void { + this.invalidations.add(listener); + return () => this.invalidations.delete(listener); + } + + emit(event: PiRpcEvent): void { + for (const listener of this.events) listener(event); + } + + async stop() { + return { mode: 'stdin-close' as const, code: 0, signal: null }; + } +} + +afterEach(async () => { + await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true }))); +}); + +describe('Pi Conversation runtime', () => { + it('separates RPC acceptance from settle and changes only the target Conversation model', async () => { + const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-runtime-')); + roots.push(projectPath); + const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), { + createId: () => 'project-a', + now: () => NOW, + }); + await createLocalCodingProject({ projectPath, now: NOW }, projectStore); + await createCodingProjectAgent(projectPath, { + id: 'agent-a', + avatarId: 'avatar-01', + roleName: 'Implementer', + name: 'Agent A', + model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' }, + modelResolution: 'resolved', + responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] }, + }, { now: NOW }); + const store = createCodingConversationStore(projectPath, { + createId: (() => { + const ids = [ + 'f47ac10b-58cc-4372-a567-0e02b2c3d479', + '8b1a9953-c461-4d88-9c3e-7e1f8f3f2c11', + '3d594650-3436-4a8c-8b38-7d1c5e3f9a20', + ]; + return () => ids.shift() as string; + })(), + now: () => NOW, + }); + const model = { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' as const }; + const left = await store.create({ agentId: 'agent-a', title: 'Left', model, modelResolution: 'resolved' }); + const right = await store.create({ agentId: 'agent-a', title: 'Right', model, modelResolution: 'resolved' }); + const forkTarget = await store.create({ agentId: 'agent-a', title: 'Fork', model, modelResolution: 'resolved' }); + const inputs = [left, right, forkTarget].map((item) => ({ + conversationId: item.id, + projectId: 'project-a', + agentId: 'agent-a', + title: item.title, + model: { model: item.model, modelResolution: item.modelResolution }, + })); + const workers = new Map(); + const workerHistory = new Map(); + const openInputs: Array<{ conversationId: string; forkSource?: string; sourceEntryId?: string }> = []; + const pool = new PiWorkerPool({ + maxIdle: 4, + openWorker: async ({ conversation, fork, generation, existingSession }) => { + openInputs.push({ + conversationId: conversation.conversationId, + ...(fork ? { + forkSource: fork.sourceSession.piSessionId, + sourceEntryId: fork.sourceEntryId, + } : {}), + }); + const worker = new RuntimeFakeWorker(`worker-${conversation.conversationId}-${generation}`, generation); + workers.set(conversation.conversationId, worker); + workerHistory.set(conversation.conversationId, [ + ...(workerHistory.get(conversation.conversationId) ?? []), + worker, + ]); + return { + worker, + session: existingSession ?? { + piSessionId: `session-${conversation.conversationId}`, + sessionKey: `key-${conversation.conversationId}`, + }, + }; + }, + }); + const runtime = new PiConversationRuntime({ + pool, + registry: new PiSessionRegistry({ projectStore }), + createId: (kind) => `${kind}-fixed`, + resolveModel: async (candidate) => { + if (candidate.accountId !== 'account-b' || candidate.modelId !== 'model-b') { + throw new Error('model unavailable'); + } + return { + accountId: candidate.accountId, + runtimeProviderId: 'runtime-account-b', + modelId: candidate.modelId, + thinkingLevel: candidate.thinkingLevel, + input: ['text'], + }; + }, + }); + await Promise.all(inputs.slice(0, 2).map((input) => runtime.prepare(input))); + + const releasePromptAcceptance = workers.get(left.id)!.holdNext('prompt'); + let acceptanceResolved = false; + const acceptance = runtime.prompt({ + clientRequestId: 'request-left', + conversationId: left.id, + mode: 'prompt', + text: 'Implement the change', + attachments: [], + }).then((value) => { + acceptanceResolved = true; + return value; + }); + await expect.poll(() => workers.get(left.id)!.requests.length).toBe(1); + expect(acceptanceResolved).toBe(false); + releasePromptAcceptance(); + const accepted = await acceptance; + expect(accepted).toMatchObject({ accepted: true, runId: 'run-fixed', mode: 'prompt' }); + expect((await runtime.getSnapshot(left.id)).run.status).toBe('running'); + workers.get(left.id)!.emit({ type: 'agent_end' }); + expect((await runtime.getSnapshot(left.id)).run.status).toBe('running'); + workers.get(left.id)!.emit({ type: 'agent_settled' }); + await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('idle'); + + const changed = await runtime.setModel({ + conversationId: left.id, + accountId: 'account-b', + modelId: 'model-b', + }); + expect(changed.model).toEqual({ accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'medium' }); + expect(workerHistory.get(left.id)).toHaveLength(2); + expect(workerHistory.get(left.id)![0]!.requests).not.toContainEqual(expect.objectContaining({ + type: 'set_model', + })); + expect(workers.get(left.id)!.requests.map(({ type }) => type).sort()).toEqual([ + 'get_entries', + 'get_state', + ]); + expect(pool.getState(left.id)).toMatchObject({ generation: 2, state: 'ready' }); + expect(pool.getState(right.id)).toMatchObject({ generation: 1, state: 'ready' }); + expect((await runtime.getSnapshot(right.id)).conversation.model.model).toEqual(model); + + await expect(runtime.setModel({ + conversationId: left.id, + accountId: 'account-missing', + modelId: 'model-missing', + })).rejects.toThrow('model unavailable'); + expect((await runtime.getSnapshot(left.id)).conversation.model).toEqual(changed); + + await runtime.prompt({ + clientRequestId: 'request-left-active', + conversationId: left.id, + mode: 'prompt', + text: 'Keep working', + attachments: [], + }); + await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('running'); + workers.get(left.id)!.failNext('steer'); + await expect(runtime.steer({ + clientRequestId: 'request-left-rejected-steer', + conversationId: left.id, + text: 'This command is rejected', + attachments: [], + })).rejects.toThrow('fake steer rejection'); + expect((await runtime.getSnapshot(left.id)).queue.items).toEqual([]); + await expect(runtime.steer({ + clientRequestId: 'request-left-steer', + conversationId: left.id, + text: 'Use the smaller seam', + attachments: [], + })).resolves.toMatchObject({ accepted: true, mode: 'steer', queuePosition: 1 }); + await expect(runtime.followUp({ + clientRequestId: 'request-left-follow-up', + conversationId: left.id, + text: 'Then run the focused test', + attachments: [], + })).resolves.toMatchObject({ accepted: true, mode: 'follow-up', queuePosition: 2 }); + await expect(runtime.setThinking({ + conversationId: left.id, + thinkingLevel: 'high', + })).resolves.toMatchObject({ + model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' }, + }); + workers.get(left.id)!.failNext('abort'); + await expect(runtime.abort(left.id)).rejects.toThrow('fake abort rejection'); + expect((await runtime.getSnapshot(left.id)).run.status).toBe('running'); + await runtime.abort(left.id); + expect((await runtime.getSnapshot(left.id)).run.status).toBe('aborting'); + expect(workers.get(left.id)!.requests.slice(-6).map(({ type }) => type)).toEqual([ + 'steer', + 'steer', + 'follow_up', + 'set_thinking_level', + 'abort', + 'abort', + ]); + expect(workers.get(right.id)!.requests).toHaveLength(0); + workers.get(left.id)!.emit({ type: 'agent_settled' }); + await expect.poll(async () => (await runtime.getSnapshot(left.id)).run.status).toBe('idle'); + expect((await runtime.getSnapshot(left.id)).queue.items).toEqual([]); + + workers.get(right.id)!.failNext('compact'); + await expect(runtime.compact(right.id)).rejects.toThrow('fake compact rejection'); + expect((await runtime.getSnapshot(right.id)).run.status).toBe('error'); + await runtime.compact(right.id); + expect((await runtime.getSnapshot(right.id)).run.status).toBe('compacting'); + expect(workers.get(right.id)!.requests.at(-1)).toEqual({ type: 'compact' }); + expect((await runtime.getSnapshot(left.id)).run.status).toBe('idle'); + workers.get(right.id)!.emit({ type: 'agent_end' }); + expect((await runtime.getSnapshot(right.id)).run.status).toBe('compacting'); + workers.get(right.id)!.emit({ type: 'agent_settled' }); + await expect.poll(async () => (await runtime.getSnapshot(right.id)).run.status).toBe('idle'); + + await expect(runtime.recover(right.id)).resolves.toMatchObject({ + conversationId: right.id, + status: 'ready', + workerGeneration: 2, + }); + expect((await runtime.getSnapshot(right.id)).cursor).toMatchObject({ + workerGeneration: 2, + seq: 0, + }); + expect(workers.get(right.id)!.requests.map(({ type }) => type).sort()).toEqual([ + 'get_entries', + 'get_state', + ]); + expect(pool.getState(left.id)).toMatchObject({ generation: 2, state: 'idle' }); + + const sourceRequestCount = workers.get(left.id)!.requests.length; + const forked = await runtime.fork({ + sourceConversationId: left.id, + sourceEntryId: 'entry-a', + conversation: inputs[2]!, + }); + expect(forked.conversationId).toBe(forkTarget.id); + expect(workers.get(forkTarget.id)?.id).toBe(`worker-${forkTarget.id}-1`); + expect(workers.get(left.id)!.requests).toHaveLength(sourceRequestCount); + expect(openInputs.at(-1)).toEqual({ + conversationId: forkTarget.id, + forkSource: `session-${left.id}`, + sourceEntryId: 'entry-a', + }); + + await runtime.dispose(forkTarget.id); + expect(pool.getState(forkTarget.id)).toBeNull(); + expect(pool.getState(left.id)).toMatchObject({ generation: 2, state: 'idle' }); + await expect(runtime.getSnapshot(forkTarget.id)).rejects.toMatchObject({ + publicError: { code: 'CODING_CONVERSATION_NOT_FOUND' }, + }); + }); +}); diff --git a/tests/unit/pi-managed-worker-opener.test.ts b/tests/unit/pi-managed-worker-opener.test.ts new file mode 100644 index 0000000..7d61a08 --- /dev/null +++ b/tests/unit/pi-managed-worker-opener.test.ts @@ -0,0 +1,173 @@ +// @vitest-environment node + +import { mkdtemp, readFile, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; +import { afterEach, describe, expect, it } from 'vitest'; +import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store'; +import { createCodingProjectAgent } from '../../electron/coding-projects/project-config'; +import { + createCodingProjectStore, + createLocalCodingProject, + createMemoryCodingProjectStorage, +} from '../../electron/coding-projects/project-store'; +import type { ProviderAccount } from '../../electron/shared/providers/types'; +import { + createPiManagedWorkerOpener, + type PiWorkerProcessAdapter, +} from '../../electron/coding-runtime/pi/runtime'; +import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry'; +import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors'; +import type { + PiRpcCommand, + PiRpcEvent, + PiRpcRequestOptions, + PiRpcResponse, +} from '../../electron/coding-runtime/pi/rpc-client'; +import type { PiWorkerProcessOptions } from '../../electron/coding-runtime/pi/worker-process'; +import type { PiRuntimeTelemetryEvent } from '../../electron/coding-runtime/pi/telemetry'; + +const roots: string[] = []; +const NOW = '2026-08-22T16:00:00.000Z'; + +class OpenerFakeProcess implements PiWorkerProcessAdapter { + readonly generation = 1; + constructor( + private readonly sessionId: string, + private readonly options: PiWorkerProcessOptions, + ) {} + + async start() { return this; } + + async request(command: PiRpcCommand, _options?: PiRpcRequestOptions): Promise> { + if (command.type !== 'get_state') { + return { type: 'response', id: 'fake', success: true }; + } + return { + type: 'response', + id: 'fake-state', + success: true, + data: { + sessionId: this.sessionId, + sessionFile: path.join(this.options.sessionDir, `${this.sessionId}.jsonl`), + } as T, + }; + } + + subscribe(_listener: (event: PiRpcEvent) => void): () => void { return () => undefined; } + subscribeInvalidation(_listener: (error: PiProcessError) => void): () => void { return () => undefined; } + async stop() { return { mode: 'stdin-close' as const, code: 0, signal: null }; } +} + +afterEach(async () => { + await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true }))); +}); + +describe('managed Pi worker opener', () => { + it('reuses PI-040 credential, managed-resource, and persistent-session projections', async () => { + const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-opener-')); + roots.push(root); + const projectPath = path.join(root, 'project'); + const userDataDir = path.join(root, 'user-data'); + const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), { + createId: () => 'project-a', + now: () => NOW, + }); + await createLocalCodingProject({ projectPath, now: NOW }, projectStore); + await createCodingProjectAgent(projectPath, { + id: 'agent-a', + avatarId: 'avatar-01', + roleName: 'Implementer', + name: 'Agent A', + model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'high' }, + modelResolution: 'resolved', + responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] }, + prompt: 'PRIVATE MANAGED PROMPT', + skillIds: ['grilling'], + }, { now: NOW }); + const conversationStore = createCodingConversationStore(projectPath, { + createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d479', + now: () => NOW, + }); + const created = await conversationStore.create({ + agentId: 'agent-a', + title: 'Conversation A', + model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'high' }, + modelResolution: 'resolved', + }); + const input = { + conversationId: created.id, + projectId: 'project-a', + agentId: 'agent-a', + title: created.title, + model: { model: created.model, modelResolution: created.modelResolution }, + } as const; + const account: ProviderAccount = { + id: 'account-a', + vendorId: 'custom', + label: 'Account A', + authMode: 'api_key', + apiProtocol: 'openai-completions', + baseUrl: 'https://provider.example/v1', + model: 'model-a', + enabled: true, + isDefault: true, + createdAt: NOW, + updatedAt: NOW, + }; + const processOptions: PiWorkerProcessOptions[] = []; + const telemetry: PiRuntimeTelemetryEvent[] = []; + const registry = new PiSessionRegistry({ projectStore }); + const opener = createPiManagedWorkerOpener({ + registry, + executablePath: 'electron.exe', + cliPath: 'pi-cli.js', + userDataDir, + bundledSkillsDir: path.resolve('resources/coding-skills'), + loadProviderInput: async () => ({ accounts: [account], modelSummaries: [] }), + resolveCredential: async () => 'provider-secret-value', + createSessionKey: () => 'session-key-a', + onTelemetry: (event) => telemetry.push(event), + createProcess: (options) => { + processOptions.push(options); + const sessionIndex = options.additionalArgs?.indexOf('--session-id') ?? -1; + return new OpenerFakeProcess(options.additionalArgs?.[sessionIndex + 1] ?? '', options); + }, + }); + + const first = await opener({ + conversation: input, + generation: 1, + revision: { provider: 1, resources: 1 }, + }); + const reopened = await opener({ + conversation: input, + generation: 2, + revision: { provider: 2, resources: 1 }, + existingSession: first.session, + }); + + expect(first.session).toEqual({ piSessionId: 'session-key-a', sessionKey: 'session-key-a' }); + expect(reopened.session).toEqual(first.session); + expect(processOptions).toHaveLength(2); + for (const options of processOptions) { + const argv = JSON.stringify(options.additionalArgs); + expect(argv).toContain('--system-prompt'); + expect(argv).toContain('grilling'); + expect(argv).toContain('--session-id'); + expect(argv).not.toContain('PRIVATE MANAGED PROMPT'); + expect(argv).not.toContain('provider-secret-value'); + expect(Object.values(options.env ?? {})).toContain('provider-secret-value'); + expect(options.sensitiveValues).toContain('provider-secret-value'); + } + const modelsFile = path.join(userDataDir, 'coding-runtime', 'pi', 'config', 'models.json'); + expect(await readFile(modelsFile, 'utf8')).not.toContain('provider-secret-value'); + expect(telemetry.map(({ milestone }) => milestone)).toEqual([ + 'resources.ready', 'worker.spawn', 'rpc.ready', 'session.open', + 'resources.ready', 'worker.spawn', 'rpc.ready', 'session.open', + ]); + expect(JSON.stringify(telemetry)).not.toContain(created.id); + expect(JSON.stringify(telemetry)).not.toContain('PRIVATE MANAGED PROMPT'); + expect(JSON.stringify(telemetry)).not.toContain('provider-secret-value'); + }); +}); diff --git a/tests/unit/pi-rpc-foundation.test.ts b/tests/unit/pi-rpc-foundation.test.ts index 4396cf4..cbdcbab 100644 --- a/tests/unit/pi-rpc-foundation.test.ts +++ b/tests/unit/pi-rpc-foundation.test.ts @@ -233,12 +233,15 @@ describe('Pi worker process', () => { it('settles every pending command after an unexpected exit', async () => { const worker = await makeWorker(); + const invalidations: string[] = []; + worker.subscribeInvalidation((error) => invalidations.push(error.code)); const pending = worker.request({ type: 'no_response' }); const crash = worker.request({ type: 'crash' }); await expect(Promise.all([pending, crash])).rejects.toMatchObject({ code: 'PI_RPC_EXITED' }); expect(worker.pendingCommandCount).toBe(0); expect(worker.generation).toBe(2); + expect(invalidations).toEqual(['PI_RPC_EXITED']); }); it('keeps only bounded redacted stderr diagnostics', async () => { diff --git a/tests/unit/pi-runtime-auth-recovery.test.ts b/tests/unit/pi-runtime-auth-recovery.test.ts new file mode 100644 index 0000000..708f26a --- /dev/null +++ b/tests/unit/pi-runtime-auth-recovery.test.ts @@ -0,0 +1,130 @@ +// @vitest-environment node + +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store'; +import { createCodingProjectAgent } from '../../electron/coding-projects/project-config'; +import { + createCodingProjectStore, + createLocalCodingProject, + createMemoryCodingProjectStorage, +} from '../../electron/coding-projects/project-store'; +import { PiProcessError } from '../../electron/coding-runtime/pi/process-errors'; +import { PiConversationRuntime } from '../../electron/coding-runtime/pi/runtime'; +import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry'; +import { PiWorkerPool, type PiConversationWorker } from '../../electron/coding-runtime/pi/worker-pool'; +import type { + PiRpcCommand, + PiRpcEvent, + PiRpcRequestOptions, + PiRpcResponse, +} from '../../electron/coding-runtime/pi/rpc-client'; + +const roots: string[] = []; +const NOW = '2026-08-22T17:00:00.000Z'; + +class AuthFailureWorker implements PiConversationWorker { + readonly generation = 1; + readonly requests: PiRpcCommand[] = []; + constructor(readonly id: string) {} + + async request(command: PiRpcCommand, _options?: PiRpcRequestOptions): Promise> { + this.requests.push(structuredClone(command)); + if (command.type === 'prompt') { + throw new PiProcessError('PI_RPC_RESPONSE_ERROR', '401 provider authentication failed'); + } + return { type: 'response', id: 'fake', success: true }; + } + + subscribe(_listener: (event: PiRpcEvent) => void): () => void { return () => undefined; } + subscribeInvalidation(_listener: (error: PiProcessError) => void): () => void { return () => undefined; } + async stop() { return { mode: 'stdin-close' as const, code: 0, signal: null }; } +} + +afterEach(async () => { + await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true }))); +}); + +describe('Pi runtime Provider authentication recovery', () => { + it('refreshes and reopens once, then exposes the second authentication failure without looping', async () => { + const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-auth-')); + roots.push(projectPath); + const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), { + createId: () => 'project-auth', + now: () => NOW, + }); + await createLocalCodingProject({ projectPath, now: NOW }, projectStore); + await createCodingProjectAgent(projectPath, { + id: 'agent-auth', + avatarId: 'avatar-01', + roleName: 'Implementer', + name: 'Auth Agent', + model: { accountId: 'account-auth', modelId: 'model-auth', thinkingLevel: 'medium' }, + modelResolution: 'resolved', + responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] }, + }, { now: NOW }); + const store = createCodingConversationStore(projectPath, { + createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d479', + now: () => NOW, + }); + const conversation = await store.create({ + agentId: 'agent-auth', + title: 'Auth Conversation', + model: { accountId: 'account-auth', modelId: 'model-auth', thinkingLevel: 'medium' }, + modelResolution: 'resolved', + }); + const workers: AuthFailureWorker[] = []; + const pool = new PiWorkerPool({ + maxIdle: 2, + openWorker: async ({ conversation: input, existingSession }) => { + const worker = new AuthFailureWorker(`worker-${workers.length + 1}`); + workers.push(worker); + return { + worker, + session: existingSession ?? { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + const refreshCredential = vi.fn(async () => undefined); + const runtime = new PiConversationRuntime({ + pool, + registry: new PiSessionRegistry({ projectStore }), + resolveModel: async () => { throw new Error('not used'); }, + refreshCredential, + isAuthenticationError: (error) => ( + error instanceof PiProcessError && error.message.includes('401') + ), + createId: () => 'run-auth', + }); + await runtime.prepare({ + conversationId: conversation.id, + projectId: 'project-auth', + agentId: 'agent-auth', + title: conversation.title, + model: { model: conversation.model, modelResolution: conversation.modelResolution }, + }); + + await expect(runtime.prompt({ + clientRequestId: 'request-auth', + conversationId: conversation.id, + mode: 'prompt', + text: 'Do not leak credentials', + attachments: [], + })).rejects.toMatchObject({ + code: 'PI_RPC_RESPONSE_ERROR', + }); + + await expect.poll(() => workers.length).toBe(2); + await expect.poll(async () => (await runtime.getSnapshot(conversation.id)).run.status).toBe('error'); + expect(refreshCredential).toHaveBeenCalledTimes(1); + expect(workers).toHaveLength(2); + expect(workers.map((worker) => worker.requests.filter(({ type }) => type === 'prompt').length)) + .toEqual([1, 1]); + expect(pool.getState(conversation.id)).toMatchObject({ state: 'idle', generation: 2 }); + }); +}); diff --git a/tests/unit/pi-session-registry.test.ts b/tests/unit/pi-session-registry.test.ts new file mode 100644 index 0000000..6d9419d --- /dev/null +++ b/tests/unit/pi-session-registry.test.ts @@ -0,0 +1,95 @@ +// @vitest-environment node + +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { createCodingConversationStore } from '../../electron/coding-projects/conversation-store'; +import { createCodingProjectAgent } from '../../electron/coding-projects/project-config'; +import { + createCodingProjectStore, + createLocalCodingProject, + createMemoryCodingProjectStorage, +} from '../../electron/coding-projects/project-store'; +import { PiSessionRegistry } from '../../electron/coding-runtime/pi/session-registry'; + +const roots: string[] = []; +const NOW = '2026-08-22T14:00:00.000Z'; + +afterEach(async () => { + await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true }))); +}); + +describe('Pi session registry', () => { + it('persists one binding and target-only model state across registry reopen', async () => { + const projectPath = await mkdtemp(path.join(tmpdir(), 'makelore-pi-registry-')); + roots.push(projectPath); + const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), { + createId: () => 'project-a', + now: () => NOW, + }); + const { project } = await createLocalCodingProject({ projectPath, now: NOW }, projectStore); + await createCodingProjectAgent(projectPath, { + id: 'agent-a', + avatarId: 'avatar-01', + roleName: 'Implementer', + name: 'Agent A', + model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' }, + modelResolution: 'resolved', + responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] }, + prompt: 'Managed prompt', + skillIds: ['tdd'], + }, { now: NOW }); + const conversations = createCodingConversationStore(projectPath, { + createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d479', + now: () => NOW, + }); + const created = await conversations.create({ + agentId: 'agent-a', + title: 'Conversation A', + model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' }, + modelResolution: 'resolved', + }); + const input = { + conversationId: created.id, + projectId: 'project-a', + agentId: 'agent-a', + title: created.title, + model: { model: created.model, modelResolution: created.modelResolution }, + } as const; + const registry = new PiSessionRegistry({ projectStore }); + const createBinding = vi.fn(async () => ({ + piSessionId: 'pi-session-a', + sessionKey: 'session-key-a', + })); + + const [first, duplicate] = await Promise.all([ + registry.ensureBinding(input, createBinding), + registry.ensureBinding(input, createBinding), + ]); + expect(createBinding).toHaveBeenCalledTimes(1); + expect(first.session).toEqual({ piSessionId: 'pi-session-a', sessionKey: 'session-key-a' }); + expect(duplicate).toEqual(first); + + await registry.setModel(created.id, { + model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' }, + modelResolution: 'resolved', + }); + const reopened = await new PiSessionRegistry({ projectStore }).prepare({ + ...input, + model: { + model: { thinkingLevel: 'high', modelId: 'model-b', accountId: 'account-b' }, + modelResolution: 'resolved', + }, + }); + expect(reopened).toMatchObject({ + projectPath: project.path, + agent: { id: 'agent-a', prompt: 'Managed prompt', skillIds: ['tdd'] }, + conversation: { + id: created.id, + model: { accountId: 'account-b', modelId: 'model-b', thinkingLevel: 'high' }, + }, + session: { piSessionId: 'pi-session-a', sessionKey: 'session-key-a' }, + }); + }); +}); diff --git a/tests/unit/pi-worker-pool-process-integration.test.ts b/tests/unit/pi-worker-pool-process-integration.test.ts new file mode 100644 index 0000000..1ec0ba0 --- /dev/null +++ b/tests/unit/pi-worker-pool-process-integration.test.ts @@ -0,0 +1,163 @@ +// @vitest-environment node + +import { mkdir, mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; +import { afterEach, describe, expect, it } from 'vitest'; +import type { PrepareConversationInput } from '../../electron/coding-runtime/contracts'; +import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors'; +import type { + PiRpcCommand, + PiRpcEvent, + PiRpcRequestOptions, + PiRpcResponse, +} from '../../electron/coding-runtime/pi/rpc-client'; +import { + PiProcessBudget, + PiWorkerPool, + type PiConversationWorker, +} from '../../electron/coding-runtime/pi/worker-pool'; +import { + PiWorkerProcess, + type PiWorkerStopResult, +} from '../../electron/coding-runtime/pi/worker-process'; + +const scratchRoots: string[] = []; + +function conversation(conversationId: string): PrepareConversationInput { + return { + conversationId, + projectId: 'project-process', + agentId: 'agent-process', + title: conversationId, + model: { + model: { accountId: 'account-process', modelId: 'model-process', thinkingLevel: 'medium' }, + modelResolution: 'resolved', + }, + }; +} + +class ProcessBackedWorker implements PiConversationWorker { + constructor( + readonly id: string, + readonly generation: number, + private readonly process: PiWorkerProcess, + ) {} + + request(command: PiRpcCommand, options?: PiRpcRequestOptions): Promise> { + return this.process.request(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 { + return this.process.stop(); + } +} + +afterEach(async () => { + await Promise.all(scratchRoots.splice(0).map((root) => rm(root, { + recursive: true, + force: true, + maxRetries: 3, + }))); +}); + +describe('Pi worker pool process integration', () => { + it('runs two child processes concurrently and aborts only the addressed Conversation', async () => { + const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-pool-process-')); + scratchRoots.push(root); + const processBudget = new PiProcessBudget(8); + const events: Array<{ conversationId: string; event: PiRpcEvent }> = []; + const pool = new PiWorkerPool({ + maxRunning: 4, + maxIdle: 4, + processBudget, + openWorker: async ({ conversation: input, generation }) => { + const workerRoot = path.join(root, input.conversationId); + const configDir = path.join(workerRoot, 'config'); + const sessionDir = path.join(workerRoot, 'sessions'); + const cwd = path.join(workerRoot, 'project'); + await Promise.all([ + mkdir(configDir, { recursive: true }), + mkdir(sessionDir, { recursive: true }), + mkdir(cwd, { recursive: true }), + ]); + const child = await new PiWorkerProcess({ + executablePath: process.execPath, + cliPath: path.resolve('tests/fixtures/fake-pi-pool-child.mjs'), + cwd, + configDir, + sessionDir, + commandTimeoutMs: 2_000, + shutdownGraceMs: 500, + }).start(); + return { + worker: new ProcessBackedWorker( + `worker-${input.conversationId}-${generation}`, + generation, + child, + ), + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + const unsubscribe = pool.subscribe((event) => { + if (event.type === 'worker.event') { + events.push({ conversationId: event.conversationId, event: event.event }); + } + }); + + try { + await Promise.all([ + pool.prepare(conversation('conversation-left')), + pool.prepare(conversation('conversation-right')), + ]); + const left = pool.startTopLevel({ + conversationId: 'conversation-left', + runId: 'run-left', + command: { type: 'prompt', message: 'left', delayMs: 1_000 }, + }); + const right = pool.startTopLevel({ + conversationId: 'conversation-right', + runId: 'run-right', + command: { type: 'prompt', message: 'right', delayMs: 150 }, + }); + await Promise.all([left.accepted, right.accepted]); + expect(pool.getState('conversation-left')?.state).toBe('running'); + expect(pool.getState('conversation-right')?.state).toBe('running'); + + await pool.request('conversation-left', { type: 'abort' }); + await expect.poll(() => pool.getState('conversation-left')?.state).toBe('idle'); + expect(pool.getState('conversation-right')?.state).toBe('running'); + await expect.poll(() => pool.getState('conversation-right')?.state).toBe('idle'); + + expect(events).toEqual(expect.arrayContaining([ + expect.objectContaining({ + conversationId: 'conversation-left', + event: expect.objectContaining({ type: 'agent_settled', marker: 'left', reason: 'aborted' }), + }), + expect.objectContaining({ + conversationId: 'conversation-right', + event: expect.objectContaining({ type: 'agent_settled', marker: 'right', reason: 'completed' }), + }), + ])); + expect(events.some(({ conversationId, event }) => ( + conversationId === 'conversation-right' && event.reason === 'aborted' + ))).toBe(false); + } finally { + unsubscribe(); + await pool.shutdown(); + } + expect(processBudget.activeCount).toBe(0); + }, 10_000); +}); diff --git a/tests/unit/pi-worker-pool.test.ts b/tests/unit/pi-worker-pool.test.ts new file mode 100644 index 0000000..9b6ee83 --- /dev/null +++ b/tests/unit/pi-worker-pool.test.ts @@ -0,0 +1,531 @@ +// @vitest-environment node + +import { describe, expect, it } from 'vitest'; +import type { PrepareConversationInput } from '../../electron/coding-runtime/contracts'; +import { + PiProcessBudget, + PiWorkerPool, + type PiConversationWorker, + type PiWorkerOpenResult, +} from '../../electron/coding-runtime/pi/worker-pool'; +import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors'; +import { PiProcessError as PiProcessFailure } from '../../electron/coding-runtime/pi/process-errors'; +import type { PiRpcCommand, PiRpcEvent } from '../../electron/coding-runtime/pi/rpc-client'; +import type { PiRuntimeTelemetryEvent } from '../../electron/coding-runtime/pi/telemetry'; + +const MODEL = { + model: { + accountId: 'account-a', + modelId: 'model-a', + thinkingLevel: 'medium' as const, + }, + modelResolution: 'resolved' as const, +}; + +function conversation(conversationId: string): PrepareConversationInput { + return { + conversationId, + projectId: 'project-a', + agentId: 'agent-a', + title: conversationId, + model: MODEL, + }; +} + +function deferred(): { promise: Promise; resolve(): void } { + let resolve!: () => void; + const promise = new Promise((done) => { resolve = done; }); + return { promise, resolve }; +} + +class FakeWorker implements PiConversationWorker { + readonly generation = 1; + readonly requests: PiRpcCommand[] = []; + stopped = false; + private readonly eventListeners = new Set<(event: PiRpcEvent) => void>(); + private readonly invalidationListeners = new Set<(error: PiProcessError) => void>(); + + constructor(readonly id: string) {} + + async request() { + this.requests.push(arguments[0] as PiRpcCommand); + return { type: 'response' as const, id: 'fake', success: true }; + } + + subscribe(listener: (event: PiRpcEvent) => void): () => void { + this.eventListeners.add(listener); + return () => this.eventListeners.delete(listener); + } + + subscribeInvalidation(listener: (error: PiProcessError) => void): () => void { + this.invalidationListeners.add(listener); + return () => this.invalidationListeners.delete(listener); + } + + emit(event: PiRpcEvent): void { + for (const listener of this.eventListeners) listener(event); + } + + invalidate(error = new PiProcessFailure('PI_RPC_EXITED', 'fake worker crashed')): void { + for (const listener of this.invalidationListeners) listener(error); + } + + async stop() { + this.stopped = true; + return { mode: 'stdin-close' as const, code: 0, signal: null }; + } +} + +describe('Pi worker pool', () => { + it('single-flights prepare per Conversation and never shares its worker with another Conversation', async () => { + const gate = deferred(); + const opened: string[] = []; + const pool = new PiWorkerPool({ + openWorker: async ({ conversation: input }): Promise => { + opened.push(input.conversationId); + await gate.promise; + return { + worker: new FakeWorker(`worker-${input.conversationId}`), + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + maxIdle: 6, + }); + + const first = pool.prepare(conversation('conversation-a')); + const duplicate = pool.prepare(conversation('conversation-a')); + const other = pool.prepare(conversation('conversation-b')); + await expect.poll(() => opened).toEqual(['conversation-a', 'conversation-b']); + gate.resolve(); + + const [firstState, duplicateState, otherState] = await Promise.all([first, duplicate, other]); + expect(firstState).toEqual(duplicateState); + expect(firstState).toMatchObject({ + conversationId: 'conversation-a', + workerId: 'worker-conversation-a', + state: 'ready', + generation: 1, + }); + expect(otherState).toMatchObject({ + conversationId: 'conversation-b', + workerId: 'worker-conversation-b', + state: 'ready', + generation: 1, + }); + }); + + it('starts only four top-level runs and advances the remaining queue fairly on agent_settled', async () => { + const workers = new Map(); + const pool = new PiWorkerPool({ + maxIdle: 6, + openWorker: async ({ conversation: input }) => { + const worker = new FakeWorker(`worker-${input.conversationId}`); + workers.set(input.conversationId, worker); + return { + worker, + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + const ids = ['a', 'b', 'c', 'd', 'e', 'f'].map((id) => `conversation-${id}`); + await Promise.all(ids.map((id) => pool.prepare(conversation(id)))); + + const runs = ids.map((conversationId, index) => pool.startTopLevel({ + conversationId, + runId: `run-${index + 1}`, + command: { type: 'prompt', message: conversationId }, + })); + + expect(runs.map((run) => run.queuePosition)).toEqual([undefined, undefined, undefined, undefined, 1, 2]); + await expect.poll(() => ids.map((id) => workers.get(id)!.requests.length)) + .toEqual([1, 1, 1, 1, 0, 0]); + workers.get('conversation-a')!.emit({ type: 'agent_end' }); + expect(workers.get('conversation-e')!.requests).toHaveLength(0); + + workers.get('conversation-a')!.emit({ type: 'agent_settled' }); + await expect.poll(() => workers.get('conversation-e')!.requests.length).toBe(1); + expect(workers.get('conversation-f')!.requests).toHaveLength(0); + workers.get('conversation-b')!.emit({ type: 'agent_settled' }); + await expect.poll(() => workers.get('conversation-f')!.requests.length).toBe(1); + }); + + it('evicts only the least-recent idle worker and reopens it on demand', async () => { + const workers = new Map(); + const pool = new PiWorkerPool({ + maxIdle: 2, + openWorker: async ({ conversation: input }) => { + const worker = new FakeWorker(`worker-${input.conversationId}-${(workers.get(input.conversationId)?.length ?? 0) + 1}`); + workers.set(input.conversationId, [...(workers.get(input.conversationId) ?? []), worker]); + return { + worker, + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + + await pool.prepare(conversation('conversation-a')); + await pool.prepare(conversation('conversation-b')); + await pool.prepare(conversation('conversation-c')); + await expect.poll(() => workers.get('conversation-a')![0]!.stopped).toBe(true); + expect(pool.getState('conversation-a')).toBeNull(); + expect(pool.getState('conversation-b')?.state).toBe('ready'); + expect(pool.getState('conversation-c')?.state).toBe('ready'); + + const reopened = await pool.prepare(conversation('conversation-a')); + expect(reopened).toMatchObject({ + workerId: 'worker-conversation-a-2', + generation: 2, + state: 'ready', + }); + }); + + it('shares a fair total-process budget and starts the next worker only after a lease is released', async () => { + const processBudget = new PiProcessBudget(2); + const opened: string[] = []; + const pool = new PiWorkerPool({ + maxIdle: 3, + processBudget, + openWorker: async ({ conversation: input }) => { + opened.push(input.conversationId); + return { + worker: new FakeWorker(`worker-${input.conversationId}`), + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + + await Promise.all([ + pool.prepare(conversation('conversation-a')), + pool.prepare(conversation('conversation-b')), + ]); + const third = pool.prepare(conversation('conversation-c')); + await expect.poll(() => processBudget.waitingCount).toBe(1); + expect(opened).toEqual(['conversation-a', 'conversation-b']); + + await pool.dispose('conversation-a'); + await expect(third).resolves.toMatchObject({ conversationId: 'conversation-c', state: 'ready' }); + expect(opened).toEqual(['conversation-a', 'conversation-b', 'conversation-c']); + expect(processBudget.activeCount).toBe(2); + + await pool.shutdown(); + expect(processBudget.activeCount).toBe(0); + }); + + it('never evicts a running worker when the warm-idle LRU exceeds its cap', async () => { + const workers = new Map(); + const pool = new PiWorkerPool({ + maxIdle: 1, + openWorker: async ({ conversation: input }) => { + const worker = new FakeWorker(`worker-${input.conversationId}`); + workers.set(input.conversationId, worker); + return { + worker, + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + await pool.prepare(conversation('conversation-running')); + const running = pool.startTopLevel({ + conversationId: 'conversation-running', + runId: 'run-running', + command: { type: 'prompt', message: 'keep alive' }, + }); + await running.accepted; + await pool.prepare(conversation('conversation-old-idle')); + await pool.prepare(conversation('conversation-new-idle')); + + expect(workers.get('conversation-running')!.stopped).toBe(false); + expect(pool.getState('conversation-running')).toMatchObject({ state: 'running' }); + expect(workers.get('conversation-old-idle')!.stopped).toBe(true); + expect(pool.getState('conversation-new-idle')).toMatchObject({ state: 'ready' }); + }); + + it('cleans only the crashed generation and releases its permit for the next Conversation', async () => { + const workers = new Map(); + const pool = new PiWorkerPool({ + maxRunning: 2, + maxIdle: 3, + openWorker: async ({ conversation: input }) => { + const worker = new FakeWorker(`worker-${input.conversationId}`); + workers.set(input.conversationId, worker); + return { + worker, + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + await Promise.all(['a', 'b', 'c'].map((id) => pool.prepare(conversation(`conversation-${id}`)))); + pool.startTopLevel({ conversationId: 'conversation-a', runId: 'run-a', command: { type: 'prompt', message: 'a' } }); + pool.startTopLevel({ conversationId: 'conversation-b', runId: 'run-b', command: { type: 'prompt', message: 'b' } }); + pool.startTopLevel({ conversationId: 'conversation-c', runId: 'run-c', command: { type: 'prompt', message: 'c' } }); + await expect.poll(() => workers.get('conversation-a')!.requests.length).toBe(1); + + const cancelled: string[] = []; + for (const kind of ['command', 'interaction', 'child'] as const) { + pool.trackGenerationResource({ + conversationId: 'conversation-a', + kind, + id: `${kind}-a`, + cancel: () => { cancelled.push(kind); }, + }); + } + pool.trackGenerationResource({ + conversationId: 'conversation-b', + kind: 'child', + id: 'child-b', + cancel: () => { cancelled.push('other'); }, + }); + + workers.get('conversation-a')!.invalidate(); + + expect(cancelled.sort()).toEqual(['child', 'command', 'interaction']); + expect(pool.getState('conversation-a')).toMatchObject({ state: 'crashed', generation: 1 }); + expect(pool.getState('conversation-b')).toMatchObject({ state: 'running', generation: 1 }); + await expect.poll(() => workers.get('conversation-c')!.requests.length).toBe(1); + expect(cancelled).not.toContain('other'); + }); + + it('rebuilds stale idle workers before prompt and lets running workers settle first', async () => { + const workers = new Map(); + const revisions: Array<{ conversationId: string; provider: number; resources: number }> = []; + const pool = new PiWorkerPool({ + maxIdle: 4, + openWorker: async ({ conversation: input, revision }) => { + const worker = new FakeWorker(`worker-${input.conversationId}-${(workers.get(input.conversationId)?.length ?? 0) + 1}`); + workers.set(input.conversationId, [...(workers.get(input.conversationId) ?? []), worker]); + revisions.push({ conversationId: input.conversationId, ...revision }); + return { + worker, + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + await Promise.all([ + pool.prepare(conversation('conversation-running')), + pool.prepare(conversation('conversation-idle')), + ]); + pool.startTopLevel({ + conversationId: 'conversation-running', + runId: 'run-running', + command: { type: 'prompt', message: 'running' }, + }); + await expect.poll(() => workers.get('conversation-running')![0]!.requests.length).toBe(1); + + pool.markProviderStale(); + pool.startTopLevel({ + conversationId: 'conversation-idle', + runId: 'run-idle', + command: { type: 'prompt', message: 'idle' }, + }); + + await expect.poll(() => workers.get('conversation-idle')?.length).toBe(2); + expect(workers.get('conversation-idle')![0]!.stopped).toBe(true); + await expect.poll(() => workers.get('conversation-idle')![1]!.requests.length).toBe(1); + expect(workers.get('conversation-running')).toHaveLength(1); + expect(workers.get('conversation-running')![0]!.stopped).toBe(false); + + workers.get('conversation-running')![0]!.emit({ type: 'agent_settled' }); + await expect.poll(() => workers.get('conversation-running')?.length).toBe(2); + expect(workers.get('conversation-running')![0]!.stopped).toBe(true); + expect(revisions).toEqual(expect.arrayContaining([ + { conversationId: 'conversation-idle', provider: 2, resources: 1 }, + { conversationId: 'conversation-running', provider: 2, resources: 1 }, + ])); + }); + + it('rejects queued work and stops every parent worker during app shutdown', async () => { + const workers: FakeWorker[] = []; + const pool = new PiWorkerPool({ + maxRunning: 1, + maxIdle: 3, + openWorker: async ({ conversation: input }) => { + const worker = new FakeWorker(`worker-${input.conversationId}`); + workers.push(worker); + return { + worker, + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + await Promise.all(['a', 'b', 'c'].map((id) => pool.prepare(conversation(`conversation-${id}`)))); + pool.startTopLevel({ conversationId: 'conversation-a', runId: 'run-a', command: { type: 'prompt', message: 'a' } }); + const queued = pool.startTopLevel({ + conversationId: 'conversation-b', + runId: 'run-b', + command: { type: 'prompt', message: 'b' }, + }); + let childCancelled = false; + pool.trackGenerationResource({ + conversationId: 'conversation-a', + kind: 'child', + id: 'child-a', + cancel: () => { childCancelled = true; }, + }); + + await pool.shutdown(); + + await expect(queued.accepted).rejects.toThrow('shutting down'); + expect(childCancelled).toBe(true); + expect(workers.every((worker) => worker.stopped)).toBe(true); + expect(pool.getState('conversation-a')).toBeNull(); + await expect(pool.prepare(conversation('conversation-after-quit'))) + .rejects.toThrow('shutting down'); + }); + + it('waits for an in-flight fork open and stops that process before shutdown completes', async () => { + const forkGate = deferred(); + const workers: FakeWorker[] = []; + let forkOpenStarted = false; + const pool = new PiWorkerPool({ + maxIdle: 3, + openWorker: async ({ conversation: input, existingSession }) => { + if (input.conversationId === 'conversation-fork') { + forkOpenStarted = true; + await forkGate.promise; + } + const worker = new FakeWorker(`worker-${input.conversationId}`); + workers.push(worker); + return { + worker, + session: existingSession ?? { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + await pool.prepare(conversation('conversation-source')); + const fork = pool.fork('conversation-source', conversation('conversation-fork')); + await expect.poll(() => forkOpenStarted).toBe(true); + const forkOutcome = fork.then( + () => 'resolved', + (error: unknown) => error instanceof Error ? error.message : String(error), + ); + let shutdownCompleted = false; + const shutdown = pool.shutdown().then(() => { shutdownCompleted = true; }); + await Promise.resolve(); + expect(shutdownCompleted).toBe(false); + + forkGate.resolve(); + await shutdown; + expect(await forkOutcome).toContain('shutting down'); + expect(workers).toHaveLength(2); + expect(workers.every((worker) => worker.stopped)).toBe(true); + }); + + it('recovers the target session with a new generation and disposes no sibling worker', async () => { + const workers = new Map(); + const pool = new PiWorkerPool({ + maxIdle: 4, + openWorker: async ({ conversation: input, existingSession }) => { + const worker = new FakeWorker(`worker-${input.conversationId}-${(workers.get(input.conversationId)?.length ?? 0) + 1}`); + workers.set(input.conversationId, [...(workers.get(input.conversationId) ?? []), worker]); + return { + worker, + session: existingSession ?? { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + await Promise.all([ + pool.prepare(conversation('conversation-a')), + pool.prepare(conversation('conversation-b')), + ]); + workers.get('conversation-a')![0]!.invalidate(); + + const recovered = await pool.recover('conversation-a'); + expect(recovered).toMatchObject({ + workerId: 'worker-conversation-a-2', + generation: 2, + state: 'ready', + session: { piSessionId: 'session-conversation-a', sessionKey: 'key-conversation-a' }, + }); + expect(workers.get('conversation-a')![0]!.stopped).toBe(true); + expect(workers.get('conversation-b')![0]!.stopped).toBe(false); + + await pool.dispose('conversation-a'); + expect(workers.get('conversation-a')![1]!.stopped).toBe(true); + expect(pool.getState('conversation-a')).toBeNull(); + expect(pool.getState('conversation-b')).toMatchObject({ state: 'ready', generation: 1 }); + }); + + it('records privacy-safe queue wait and RPC prompt acceptance spans', async () => { + let now = 0; + const telemetry: PiRuntimeTelemetryEvent[] = []; + const workers = new Map(); + const pool = new PiWorkerPool({ + maxRunning: 1, + maxIdle: 2, + now: () => now, + onTelemetry: (event) => telemetry.push(event), + openWorker: async ({ conversation: input }) => { + const worker = new FakeWorker(`worker-${input.conversationId}`); + workers.set(input.conversationId, worker); + return { + worker, + session: { + piSessionId: `session-${input.conversationId}`, + sessionKey: `key-${input.conversationId}`, + }, + }; + }, + }); + await Promise.all([ + pool.prepare(conversation('f47ac10b-58cc-4372-a567-0e02b2c3d479')), + pool.prepare(conversation('8b1a9953-c461-4d88-9c3e-7e1f8f3f2c11')), + ]); + const first = pool.startTopLevel({ + conversationId: 'f47ac10b-58cc-4372-a567-0e02b2c3d479', + runId: 'run-first-1234567890', + command: { type: 'prompt', message: 'private first prompt' }, + }); + now = 5; + const second = pool.startTopLevel({ + conversationId: '8b1a9953-c461-4d88-9c3e-7e1f8f3f2c11', + runId: 'run-second-1234567890', + command: { type: 'prompt', message: 'private second prompt' }, + }); + await first.accepted; + now = 25; + workers.get('f47ac10b-58cc-4372-a567-0e02b2c3d479')!.emit({ type: 'agent_settled' }); + await second.accepted; + + expect(telemetry.map(({ milestone }) => milestone)).toEqual([ + 'prompt.accepted', + 'worker.queue_wait', + 'prompt.accepted', + ]); + expect(telemetry[1]).toMatchObject({ durationMs: 20, workerGeneration: 1, cold: true }); + const serialized = JSON.stringify(telemetry); + expect(serialized).not.toContain('private first prompt'); + expect(serialized).not.toContain('private second prompt'); + expect(serialized).not.toContain('f47ac10b-58cc-4372-a567-0e02b2c3d479'); + expect(serialized).not.toContain('8b1a9953-c461-4d88-9c3e-7e1f8f3f2c11'); + }); +});