feat: add Pi subagent scheduler and child runtime
This commit is contained in:
227
electron/coding-runtime/pi/subagent-child.ts
Normal file
227
electron/coding-runtime/pi/subagent-child.ts
Normal file
@@ -0,0 +1,227 @@
|
||||
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;
|
||||
}
|
||||
};
|
||||
}
|
||||
Reference in New Issue
Block a user