260 lines
10 KiB
TypeScript
260 lines
10 KiB
TypeScript
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<unknown>;
|
|
request<T = unknown>(
|
|
command: PiRpcCommand,
|
|
options?: PiRpcRequestOptions,
|
|
): Promise<PiRpcResponse<T>>;
|
|
subscribe(listener: (event: PiRpcEvent) => void): () => void;
|
|
subscribeInvalidation(listener: (error: PiProcessError) => void): () => void;
|
|
stop(reason: PiWorkerStopReason): 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;
|
|
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<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(reason: PiWorkerStopReason): Promise<void> {
|
|
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<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 workerResources = options.capabilityRegistry
|
|
? await options.capabilityRegistry.resolveWorkerResources({
|
|
projectPath: project.path,
|
|
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,
|
|
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: [],
|
|
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;
|
|
}
|
|
};
|
|
}
|