import type { ModelSummary, ProviderAccount } from '../../shared/providers/types'; import { readCodingProjectConfigV2 } from '../../coding-projects/project-config'; import type { CodingProjectStore } from '../../coding-projects/project-store'; import type { PublicUsage } from '../contracts'; import type { PiManagedInputRevision } from './managed-input-revision'; import { buildPiProviderCatalog, buildPiWorkerCredentialProjection, selectPiProviderModel, writePiProviderCatalog, } from './provider-config'; import { buildPiManagedInputArgs, ensurePiManagedPaths, materializePiAgentResources, } from './resource-loader'; import type { PiRpcCommand, PiRpcEvent, PiRpcRequestOptions, PiRpcResponse, } from './rpc-client'; import type { PiSubagentChild, PiSubagentChildOpenInput, PiSubagentChildResult, } from './subagent'; import { PiSubagentChildError } from './subagent'; import type { CodingCapabilityRegistry, ResolvedWorkerResources } from '../../coding-plugins/registry'; import { PiWorkerProcess, type PiWorkerProcessOptions, type PiWorkerStopReason, type PiWorkerStopResult, } from './worker-process'; import type { PiProcessError } from './process-errors'; import type { PiManagedExtensionHost } from './extension-host'; const READ_ONLY_TOOLS = ['read', 'grep', 'find', 'ls'] as const; const CODING_TOOLS = ['read', 'bash', 'edit', 'write', 'grep', 'find', 'ls'] as const; export interface PiSubagentProcessAdapter { start(): Promise; request( command: PiRpcCommand, options?: PiRpcRequestOptions, ): Promise>; subscribe(listener: (event: PiRpcEvent) => void): () => void; subscribeInvalidation(listener: (error: PiProcessError) => void): () => void; stop(reason: PiWorkerStopReason): Promise; } export interface PiManagedSubagentChildOpenerOptions { projectStore: CodingProjectStore; executablePath: string; cliPath: string; userDataDir: string; bundledSkillsDir: string; getSkillRoots?(): readonly string[] | Promise; extensionHost: PiManagedExtensionHost; loadProviderInput(): Promise<{ accounts: ProviderAccount[]; modelSummaries: ModelSummary[] }>; resolveCredential(account: ProviderAccount): Promise; getRevision(): PiManagedInputRevision; getLocalProxyCredential?(): Promise; createProcess?(options: PiWorkerProcessOptions): PiSubagentProcessAdapter; capabilityRegistry?: CodingCapabilityRegistry; } function fallbackWorkerResources(skillIds: readonly string[]): ResolvedWorkerResources { const effectiveSkillIds = [...new Set(skillIds.map((id) => id.trim()).filter(Boolean))]; return { catalogRevision: 0, pluginIds: [], effectiveSkillIds, skillEntries: effectiveSkillIds.map((id) => ({ id, entryPath: `${id}/SKILL.md` })), tools: [], }; } function recordValue(value: unknown): Record | null { return value !== null && typeof value === 'object' && !Array.isArray(value) ? value as Record : null; } function usageOf(value: unknown): PublicUsage | undefined { const usage = recordValue(value); if (!usage || typeof usage.input !== 'number' || typeof usage.output !== 'number') return undefined; return { inputTokens: usage.input, outputTokens: usage.output, ...(typeof usage.cacheRead === 'number' ? { cacheReadTokens: usage.cacheRead } : {}), ...(typeof usage.cacheWrite === 'number' ? { cacheWriteTokens: usage.cacheWrite } : {}), }; } class ManagedPiSubagentChild implements PiSubagentChild { private ran = false; private stopped = false; constructor( readonly id: string, private readonly process: PiSubagentProcessAdapter, private readonly disposeExtension: () => Promise, ) {} async run(prompt: string, signal: AbortSignal): Promise { if (this.ran) throw new PiSubagentChildError('SUBAGENT_ALREADY_RAN'); this.ran = true; if (signal.aborted) throw new PiSubagentChildError('SUBAGENT_ABORTED'); let usage: PublicUsage | undefined; let settle!: () => void; let fail!: (error: PiSubagentChildError) => void; const settled = new Promise((resolve, reject) => { settle = resolve; fail = reject; }); const unsubscribeEvents = this.process.subscribe((event) => { if (event.type === 'message_end') { const message = recordValue(event.message); if (message?.role === 'assistant') usage = usageOf(message.usage) ?? usage; } if (event.type === 'agent_settled') settle(); }); const unsubscribeInvalidation = this.process.subscribeInvalidation(() => { fail(new PiSubagentChildError('SUBAGENT_CHILD_CRASHED')); }); const abort = () => { void this.process.request({ type: 'abort' }).catch(() => undefined); fail(new PiSubagentChildError('SUBAGENT_ABORTED')); }; signal.addEventListener('abort', abort, { once: true }); try { const accepted = await this.process.request({ type: 'prompt', message: prompt }); if (!accepted.success) throw new PiSubagentChildError('SUBAGENT_PROMPT_REJECTED'); await settled; if (signal.aborted) throw new PiSubagentChildError('SUBAGENT_ABORTED'); const response = await this.process.request<{ text?: unknown }>({ type: 'get_last_assistant_text', }, { retry: 'read-only-once' }); if (!response.success) throw new PiSubagentChildError('SUBAGENT_RESULT_UNAVAILABLE'); const summary = typeof response.data?.text === 'string' ? response.data.text : ''; return { summary, ...(usage ? { usage } : {}), }; } finally { signal.removeEventListener('abort', abort); unsubscribeEvents(); unsubscribeInvalidation(); } } async stop(reason: PiWorkerStopReason): Promise { if (this.stopped) return; this.stopped = true; try { await this.process.stop(reason); } finally { await this.disposeExtension(); } } } export function createPiManagedSubagentChildOpener( options: PiManagedSubagentChildOpenerOptions, ): (input: PiSubagentChildOpenInput) => Promise { const createProcess = options.createProcess ?? ((processOptions) => new PiWorkerProcess(processOptions)); return async (input) => { const project = (await options.projectStore.listProjects()) .find((candidate) => candidate.id === input.projectId); if (!project) throw new PiSubagentChildError('SUBAGENT_PROJECT_UNAVAILABLE'); const configRead = await readCodingProjectConfigV2(project.path); if (configRead.status !== 'valid') { throw new PiSubagentChildError('SUBAGENT_PROJECT_CONFIG_UNAVAILABLE'); } const agent = configRead.config.agents.find((candidate) => ( candidate.id === input.agentId && candidate.enabled && candidate.archivedAt === null )); if (!agent) throw new PiSubagentChildError('SUBAGENT_AGENT_UNAVAILABLE'); if (!agent.model || agent.modelResolution !== 'resolved') { throw new PiSubagentChildError('SUBAGENT_MODEL_REQUIRED'); } const providerInput = await options.loadProviderInput(); const catalog = buildPiProviderCatalog(providerInput); const selection = selectPiProviderModel(catalog, agent.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 PiSubagentChildError('SUBAGENT_MODEL_UNAVAILABLE'); const managedPaths = await ensurePiManagedPaths(options.userDataDir); await writePiProviderCatalog(managedPaths.modelsFile, catalog, agent.model); const workerResources = options.capabilityRegistry ? await options.capabilityRegistry.resolveWorkerResources({ projectPath: project.path, projectId: input.projectId, assignedSkillIds: agent.skillIds, role: 'child', }) : fallbackWorkerResources(agent.skillIds); const resources = await materializePiAgentResources({ userDataDir: options.userDataDir, projectId: input.projectId, agentId: agent.id, prompt: agent.prompt, skillEntries: workerResources.skillEntries, catalogRevision: workerResources.catalogRevision, bundledSkillsDir: options.bundledSkillsDir, ...(workerResources.skillRoots ? { skillRoots: workerResources.skillRoots } : options.getSkillRoots ? { skillRoots: await options.getSkillRoots() } : {}), ...(workerResources.effectiveSnapshot ? { effectiveSnapshot: workerResources.effectiveSnapshot } : {}), revision: options.getRevision(), }); const credential = await buildPiWorkerCredentialProjection({ account, descriptor, resolveCredential: options.resolveCredential, ...(options.getLocalProxyCredential ? { localProxyCredential: await options.getLocalProxyCredential() } : {}), }); const extension = await options.extensionHost.registerWorker({ conversationId: input.conversationId, generation: input.workerGeneration, projectId: input.projectId, projectPath: project.path, skillEntries: workerResources.skillEntries, catalogRevision: workerResources.catalogRevision, tools: [], ...(workerResources.effectiveSnapshot ? { effectiveSnapshot: workerResources.effectiveSnapshot } : {}), extensionsDir: managedPaths.extensionsDir, role: 'child', runId: input.runId, }); const process = createProcess({ executablePath: options.executablePath, cliPath: options.cliPath, cwd: project.path, configDir: resources.paths.configDir, sessionDir: resources.projectSessionsDir, tools: input.toolProfile === 'coding' ? CODING_TOOLS : READ_ONLY_TOOLS, additionalArgs: [ ...buildPiManagedInputArgs(selection, resources), '--extension', extension.extensionPath, '--no-session', ], env: { ...credential.env, ...extension.env }, sensitiveValues: [...credential.sensitiveValues, ...extension.sensitiveValues], conversationId: input.conversationId, workerGeneration: input.workerGeneration, }); try { await process.start(); const ready = await process.request({ type: 'get_state' }, { retry: 'read-only-once' }); if (!ready.success) throw new PiSubagentChildError('SUBAGENT_CHILD_START_FAILED'); return new ManagedPiSubagentChild(input.taskId, process, extension.dispose); } catch (error) { await process.stop('subagent_open_failure').catch(() => undefined); await extension.dispose(); throw error; } }; }