368 lines
16 KiB
TypeScript
368 lines
16 KiB
TypeScript
// @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, 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 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';
|
|
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
|
|
import { PiWorkerPool } from '../../electron/coding-runtime/pi/worker-pool';
|
|
import type { CodingCapabilityRegistry } from '../../electron/coding-plugins/registry';
|
|
import type { EffectivePluginSnapshot } from '../../electron/coding-plugins/effective-resolver';
|
|
import { DATA_SERVICE_PLUGIN_DEFINITION } from '../../shared/coding-plugins';
|
|
|
|
const roots: string[] = [];
|
|
const NOW = '2026-08-22T16:00:00.000Z';
|
|
|
|
function deferred<T>() {
|
|
let resolve!: (value: T | PromiseLike<T>) => void;
|
|
const promise = new Promise<T>((yes) => { resolve = yes; });
|
|
return { promise, resolve };
|
|
}
|
|
|
|
class OpenerFakeProcess implements PiWorkerProcessAdapter {
|
|
readonly generation = 1;
|
|
readonly proofResponseDelays: Array<{ commandType: string; delayMs: number }> = [];
|
|
constructor(
|
|
private readonly sessionId: string,
|
|
private readonly options: PiWorkerProcessOptions,
|
|
) {}
|
|
|
|
async start() { return this; }
|
|
|
|
async request<T = unknown>(command: PiRpcCommand, _options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
|
|
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,
|
|
};
|
|
}
|
|
|
|
async send(_command: PiRpcCommand): Promise<void> {}
|
|
|
|
subscribe(_listener: (event: PiRpcEvent) => void): () => void { return () => undefined; }
|
|
subscribeInvalidation(_listener: (error: PiProcessError) => void): () => void { return () => undefined; }
|
|
delayNextResponseForProof(commandType: string, delayMs: number): void {
|
|
this.proofResponseDelays.push({ commandType, delayMs });
|
|
}
|
|
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 processes: OpenerFakeProcess[] = [];
|
|
const telemetry: PiRuntimeTelemetryEvent[] = [];
|
|
const registry = new PiSessionRegistry({ projectStore });
|
|
const extensionHost = new PiManagedExtensionHost();
|
|
const effectiveSnapshot: EffectivePluginSnapshot = Object.freeze({
|
|
accountSessionId: 'account-a\u00001',
|
|
projectId: 'project-a',
|
|
pluginReleaseIds: Object.freeze(['plugin-release-a']),
|
|
effectiveSkillIds: Object.freeze(['grilling']),
|
|
skillEntries: Object.freeze([{ id: 'grilling', entryPath: 'grilling/SKILL.md' }]),
|
|
toolDefinitions: Object.freeze(DATA_SERVICE_PLUGIN_DEFINITION.tools),
|
|
runtimePolicies: Object.freeze([]),
|
|
unavailableReasons: Object.freeze([]),
|
|
});
|
|
const capabilityRegistry = {
|
|
resolveWorkerResources: vi.fn(async () => ({
|
|
catalogRevision: 1,
|
|
pluginIds: [],
|
|
effectiveSkillIds: ['grilling'],
|
|
skillEntries: [{ id: 'grilling', entryPath: 'grilling/SKILL.md' }],
|
|
tools: DATA_SERVICE_PLUGIN_DEFINITION.tools,
|
|
effectiveSnapshot,
|
|
})),
|
|
} as unknown as CodingCapabilityRegistry;
|
|
const firstReleaseCleanup = deferred<void>();
|
|
const releaseCleanups: Array<ReturnType<typeof vi.fn>> = [];
|
|
const registerActivePluginReleases = vi.fn((_releaseIds: readonly string[]) => {
|
|
const waitForFirst = releaseCleanups.length === 0;
|
|
const cleanup = vi.fn(async () => {
|
|
if (waitForFirst) await firstReleaseCleanup.promise;
|
|
});
|
|
releaseCleanups.push(cleanup);
|
|
return cleanup;
|
|
});
|
|
const createOpener = (
|
|
openerRegistry: PiSessionRegistry,
|
|
createSessionKey: () => string = () => 'session-key-a',
|
|
) => createPiManagedWorkerOpener({
|
|
registry: openerRegistry,
|
|
executablePath: 'electron.exe',
|
|
cliPath: 'pi-cli.js',
|
|
userDataDir,
|
|
bundledSkillsDir: path.resolve('resources/coding-skills'),
|
|
extensionHost,
|
|
capabilityRegistry,
|
|
registerActivePluginReleases,
|
|
loadProviderInput: async () => ({ accounts: [account], modelSummaries: [] }),
|
|
resolveCredential: async () => 'provider-secret-value',
|
|
createSessionKey,
|
|
onTelemetry: (event) => telemetry.push(event),
|
|
createProcess: (options) => {
|
|
processOptions.push(options);
|
|
const sessionIndex = options.additionalArgs?.indexOf('--session-id') ?? -1;
|
|
const process = new OpenerFakeProcess(options.additionalArgs?.[sessionIndex + 1] ?? '', options);
|
|
processes.push(process);
|
|
return process;
|
|
},
|
|
});
|
|
const opener = createOpener(registry);
|
|
|
|
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,
|
|
});
|
|
const restartedCreateSessionKey = vi.fn(() => 'unexpected-new-session-key');
|
|
const restarted = await createOpener(
|
|
new PiSessionRegistry({ projectStore }),
|
|
restartedCreateSessionKey,
|
|
)({
|
|
conversation: input,
|
|
generation: 3,
|
|
revision: { provider: 2, resources: 1 },
|
|
});
|
|
|
|
expect(first.session).toEqual({ piSessionId: 'session-key-a', sessionKey: 'session-key-a' });
|
|
expect(reopened.session).toEqual(first.session);
|
|
expect(restarted.session).toEqual(first.session);
|
|
expect(restartedCreateSessionKey).not.toHaveBeenCalled();
|
|
expect(processOptions).toHaveLength(3);
|
|
expect(processOptions[2]?.additionalArgs).toContain('session-key-a');
|
|
expect(processOptions[2]?.additionalArgs).not.toContain('unexpected-new-session-key');
|
|
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).toContain('--extension');
|
|
expect(argv).toContain('makelore-runtime-v5.mjs');
|
|
expect(options.additionalArgs?.filter((argument) => argument === '--extension')).toHaveLength(1);
|
|
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');
|
|
expect(options.env?.MAKELORE_PI_BRIDGE_URL).toMatch(/^http:\/\/127\.0\.0\.1:/);
|
|
expect(options.env?.MAKELORE_PI_CONTEXT_FILE).toContain('worker-');
|
|
expect(options.tools?.filter((name) => name.startsWith('data_service_')))
|
|
.toEqual(effectiveSnapshot.toolDefinitions.map(({ name }) => name));
|
|
const workerContext = JSON.parse(await readFile(
|
|
options.env?.MAKELORE_PI_CONTEXT_FILE as string,
|
|
'utf8',
|
|
)) as Record<string, unknown>;
|
|
expect(workerContext.effectivePluginSnapshot).toEqual(effectiveSnapshot);
|
|
expect(workerContext.allowedToolNames)
|
|
.toEqual(effectiveSnapshot.toolDefinitions.map(({ name }) => name));
|
|
}
|
|
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',
|
|
'resources.ready', 'worker.spawn', 'rpc.ready', 'session.open',
|
|
]);
|
|
expect(telemetry.slice(0, 4).every(({ cold }) => cold)).toBe(true);
|
|
expect(telemetry.slice(4).every(({ cold }) => !cold)).toBe(true);
|
|
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');
|
|
expect(first.worker.delayNextResponseForProof).toBeTypeOf('function');
|
|
first.worker.delayNextResponseForProof?.('prompt', 12_000);
|
|
expect(processes[0]?.proofResponseDelays).toEqual([{ commandType: 'prompt', delayMs: 12_000 }]);
|
|
expect(registerActivePluginReleases).toHaveBeenCalledTimes(3);
|
|
expect(registerActivePluginReleases).toHaveBeenNthCalledWith(1, ['plugin-release-a']);
|
|
let firstStopSettled = false;
|
|
const firstStop = first.worker.stop('test_injection').then((result) => {
|
|
firstStopSettled = true;
|
|
return result;
|
|
});
|
|
await vi.waitFor(() => expect(releaseCleanups[0]).toHaveBeenCalledOnce());
|
|
expect(firstStopSettled).toBe(false);
|
|
firstReleaseCleanup.resolve();
|
|
await firstStop;
|
|
await reopened.worker.stop('test_injection');
|
|
await restarted.worker.stop('test_injection');
|
|
expect(releaseCleanups).toHaveLength(3);
|
|
expect(releaseCleanups.every((cleanup) => cleanup.mock.calls.length === 1)).toBe(true);
|
|
await extensionHost.close();
|
|
});
|
|
|
|
it('reads the current local-proxy credential for first open, recover, and rebuild', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-proxy-opener-'));
|
|
roots.push(root);
|
|
const projectPath = path.join(root, 'project');
|
|
const userDataDir = path.join(root, 'user-data');
|
|
const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
|
|
createId: () => 'project-proxy',
|
|
now: () => NOW,
|
|
});
|
|
await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
|
|
await createCodingProjectAgent(projectPath, {
|
|
id: 'agent-proxy',
|
|
avatarId: 'avatar-01',
|
|
roleName: 'Implementer',
|
|
name: 'Proxy Agent',
|
|
model: { accountId: 'account-proxy', modelId: 'model-proxy', thinkingLevel: 'medium' },
|
|
modelResolution: 'resolved',
|
|
responsibility: { mission: 'Implement', owns: [], boundaries: [], collaborators: [], principles: [] },
|
|
}, { now: NOW });
|
|
const conversationStore = createCodingConversationStore(projectPath, {
|
|
createId: () => 'f47ac10b-58cc-4372-a567-0e02b2c3d480',
|
|
now: () => NOW,
|
|
});
|
|
const created = await conversationStore.create({
|
|
agentId: 'agent-proxy',
|
|
title: 'Proxy Conversation',
|
|
model: { accountId: 'account-proxy', modelId: 'model-proxy', thinkingLevel: 'medium' },
|
|
modelResolution: 'resolved',
|
|
});
|
|
const input = {
|
|
conversationId: created.id,
|
|
projectId: 'project-proxy',
|
|
agentId: 'agent-proxy',
|
|
title: created.title,
|
|
model: { model: created.model, modelResolution: created.modelResolution },
|
|
} as const;
|
|
const account: ProviderAccount = {
|
|
id: 'account-proxy', vendorId: 'custom', label: 'Proxy account', authMode: 'api_key',
|
|
apiProtocol: 'openai-completions', baseUrl: 'http://127.0.0.1:43123/api/ai-proxy/v1',
|
|
model: 'model-proxy', enabled: true, isDefault: true, createdAt: NOW, updatedAt: NOW,
|
|
metadata: { worksSquareCredentialMode: 'works_square_ai_gateway_proxy' },
|
|
};
|
|
const processOptions: PiWorkerProcessOptions[] = [];
|
|
const extensionHost = new PiManagedExtensionHost();
|
|
let currentToken = 'host-token-first-open';
|
|
const getLocalProxyCredential = vi.fn(async () => currentToken);
|
|
const opener = createPiManagedWorkerOpener({
|
|
registry: new PiSessionRegistry({ projectStore }),
|
|
executablePath: 'electron.exe',
|
|
cliPath: 'pi-cli.js',
|
|
userDataDir,
|
|
bundledSkillsDir: path.resolve('resources/coding-skills'),
|
|
extensionHost,
|
|
loadProviderInput: async () => ({ accounts: [account], modelSummaries: [] }),
|
|
resolveCredential: vi.fn(async () => 'stale-secret-store-token'),
|
|
getLocalProxyCredential,
|
|
createSessionKey: () => 'session-key-proxy',
|
|
createProcess: (options) => {
|
|
processOptions.push(options);
|
|
const sessionIndex = options.additionalArgs?.indexOf('--session-id') ?? -1;
|
|
return new OpenerFakeProcess(options.additionalArgs?.[sessionIndex + 1] ?? '', options);
|
|
},
|
|
});
|
|
const pool = new PiWorkerPool({ openWorker: opener, maxIdle: 2 });
|
|
|
|
await pool.prepare(input);
|
|
currentToken = 'host-token-recover-open';
|
|
await pool.recover(created.id);
|
|
currentToken = 'host-token-rebuild-open';
|
|
await pool.reconfigureConversationModel(created.id, input.model);
|
|
|
|
expect(getLocalProxyCredential).toHaveBeenCalledTimes(3);
|
|
expect(processOptions).toHaveLength(3);
|
|
expect(processOptions.map(({ env }) => Object.values(env ?? {}).find((value) => value.startsWith('host-token-'))))
|
|
.toEqual(['host-token-first-open', 'host-token-recover-open', 'host-token-rebuild-open']);
|
|
for (const token of ['host-token-first-open', 'host-token-recover-open', 'host-token-rebuild-open']) {
|
|
expect(JSON.stringify(processOptions.map(({ additionalArgs }) => additionalArgs))).not.toContain(token);
|
|
expect(JSON.stringify(pool.getDiagnostics())).not.toContain(token);
|
|
}
|
|
const modelsFile = path.join(userDataDir, 'coding-runtime', 'pi', 'config', 'models.json');
|
|
const modelsJson = await readFile(modelsFile, 'utf8');
|
|
expect(modelsJson).not.toContain('host-token-');
|
|
|
|
await pool.shutdown();
|
|
await extensionHost.close();
|
|
});
|
|
});
|