makelore/electron/coding-runtime/pi/subagent-child.ts

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 { PiWorkerProcess, type PiWorkerProcessOptions, 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<unknown>;
  request<T = unknown>(
    command: PiRpcCommand,
    options?: PiRpcRequestOptions,
  ): Promise<PiRpcResponse<T>>;
  subscribe(listener: (event: PiRpcEvent) => void): () => void;
  subscribeInvalidation(listener: (error: PiProcessError) => void): () => void;
  stop(): Promise<PiWorkerStopResult>;
}

export interface PiManagedSubagentChildOpenerOptions {
  projectStore: CodingProjectStore;
  executablePath: string;
  cliPath: string;
  userDataDir: string;
  bundledSkillsDir: string;
  extensionHost: PiManagedExtensionHost;
  loadProviderInput(): Promise<{ accounts: ProviderAccount[]; modelSummaries: ModelSummary[] }>;
  resolveCredential(account: ProviderAccount): Promise<string | null>;
  getRevision(): PiManagedInputRevision;
  getLocalProxyCredential?(): Promise<string | undefined>;
  createProcess?(options: PiWorkerProcessOptions): PiSubagentProcessAdapter;
}

function recordValue(value: unknown): Record<string, unknown> | null {
  return value !== null && typeof value === 'object' && !Array.isArray(value)
    ? value as Record<string, unknown>
    : 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<void>,
  ) {}

  async run(prompt: string, signal: AbortSignal): Promise<PiSubagentChildResult> {
    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<void>((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(): Promise<void> {
    if (this.stopped) return;
    this.stopped = true;
    try {
      await this.process.stop();
    } finally {
      await this.disposeExtension();
    }
  }
}

export function createPiManagedSubagentChildOpener(
  options: PiManagedSubagentChildOpenerOptions,
): (input: PiSubagentChildOpenInput) => Promise<PiSubagentChild> {
  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 resources = await materializePiAgentResources({
      userDataDir: options.userDataDir,
      projectId: input.projectId,
      agentId: agent.id,
      prompt: agent.prompt,
      skillIds: agent.skillIds,
      bundledSkillsDir: options.bundledSkillsDir,
      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,
      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],
    });
    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().catch(() => undefined);
      await extension.dispose();
      throw error;
    }
  };
}