512 lines
20 KiB
TypeScript
512 lines
20 KiB
TypeScript
// @vitest-environment node
|
|
|
|
import { copyFile, mkdtemp, mkdir, rm, writeFile } from 'node:fs/promises';
|
|
import { createServer, type ServerResponse } from 'node:http';
|
|
import { tmpdir } from 'node:os';
|
|
import path from 'node:path';
|
|
import { afterEach, describe, expect, it } from 'vitest';
|
|
import { PiAgentServerProcess } from '../../electron/coding-runtime/pi/agent-server-process';
|
|
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
|
|
import { materializeMakelorePiExtension } from '../../electron/coding-runtime/pi/extensions/makelore-runtime';
|
|
import { MAKELORE_DEFAULT_LANGUAGE_PROMPT } from '../../electron/coding-runtime/pi/resource-loader';
|
|
import type { PiRpcEvent } from '../../electron/coding-runtime/pi/rpc-client';
|
|
import type { PiWorkerProcessOptions } from '../../electron/coding-runtime/pi/worker-process';
|
|
|
|
const roots: string[] = [];
|
|
|
|
async function startHeldProvider(): Promise<{
|
|
baseUrl: string;
|
|
requests: Array<{ authorization: string; body: string }>;
|
|
activeCount(): number;
|
|
maxActiveCount(): number;
|
|
release(): void;
|
|
close(): Promise<void>;
|
|
}> {
|
|
const requests: Array<{ authorization: string; body: string }> = [];
|
|
const held = new Set<ServerResponse>();
|
|
let maxActive = 0;
|
|
const provider = createServer((request, response) => {
|
|
const chunks: Buffer[] = [];
|
|
request.on('data', (chunk: Buffer) => chunks.push(chunk));
|
|
request.on('end', () => {
|
|
requests.push({
|
|
authorization: String(request.headers.authorization ?? ''),
|
|
body: Buffer.concat(chunks).toString('utf8'),
|
|
});
|
|
held.add(response);
|
|
maxActive = Math.max(maxActive, held.size);
|
|
response.once('close', () => held.delete(response));
|
|
});
|
|
});
|
|
await new Promise<void>((resolve, reject) => {
|
|
provider.once('error', reject);
|
|
provider.listen(0, '127.0.0.1', resolve);
|
|
});
|
|
const address = provider.address();
|
|
if (!address || typeof address === 'string') throw new Error('Test Provider did not bind');
|
|
const release = () => {
|
|
for (const response of [...held]) {
|
|
held.delete(response);
|
|
response.writeHead(200, { 'content-type': 'text/event-stream' });
|
|
response.write(`data: ${JSON.stringify({
|
|
id: 'chatcmpl-agent-server-test',
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1_000),
|
|
model: 'test-model',
|
|
choices: [{ index: 0, delta: { role: 'assistant', content: 'DONE' }, finish_reason: null }],
|
|
})}\n\n`);
|
|
response.write(`data: ${JSON.stringify({
|
|
id: 'chatcmpl-agent-server-test',
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1_000),
|
|
model: 'test-model',
|
|
choices: [{ index: 0, delta: {}, finish_reason: 'stop' }],
|
|
usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
|
|
})}\n\n`);
|
|
response.end('data: [DONE]\n\n');
|
|
}
|
|
};
|
|
return {
|
|
baseUrl: `http://127.0.0.1:${address.port}/v1`,
|
|
requests,
|
|
activeCount: () => held.size,
|
|
maxActiveCount: () => maxActive,
|
|
release,
|
|
close: async () => {
|
|
release();
|
|
await new Promise<void>((resolve, reject) => {
|
|
provider.close((error) => error ? reject(error) : resolve());
|
|
provider.closeIdleConnections?.();
|
|
});
|
|
},
|
|
};
|
|
}
|
|
|
|
async function startBashBatchProvider(): Promise<{
|
|
baseUrl: string;
|
|
requests: Array<Record<string, unknown>>;
|
|
close(): Promise<void>;
|
|
}> {
|
|
const requests: Array<Record<string, unknown>> = [];
|
|
const provider = createServer((request, response) => {
|
|
const chunks: Buffer[] = [];
|
|
request.on('data', (chunk: Buffer) => chunks.push(chunk));
|
|
request.on('end', () => {
|
|
const body = JSON.parse(Buffer.concat(chunks).toString('utf8')) as Record<string, unknown>;
|
|
requests.push(body);
|
|
const messages = Array.isArray(body.messages) ? body.messages : [];
|
|
const hasToolResult = messages.some((message) => (
|
|
message && typeof message === 'object' && !Array.isArray(message)
|
|
&& (message as { role?: unknown }).role === 'tool'
|
|
));
|
|
response.writeHead(200, { 'content-type': 'text/event-stream' });
|
|
const writeChunk = (delta: Record<string, unknown>, finishReason: string | null) => {
|
|
response.write(`data: ${JSON.stringify({
|
|
id: 'chatcmpl-bash-batch-test',
|
|
object: 'chat.completion.chunk',
|
|
created: Math.floor(Date.now() / 1_000),
|
|
model: 'test-model',
|
|
choices: [{ index: 0, delta, finish_reason: finishReason }],
|
|
...(finishReason ? { usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 } } : {}),
|
|
})}\n\n`);
|
|
};
|
|
if (!hasToolResult) {
|
|
writeChunk({
|
|
role: 'assistant',
|
|
tool_calls: [
|
|
{
|
|
index: 0,
|
|
id: 'bash-timeout',
|
|
type: 'function',
|
|
function: {
|
|
name: 'bash',
|
|
arguments: JSON.stringify({ command: 'sleep 1', timeout: 0.05 }),
|
|
},
|
|
},
|
|
{
|
|
index: 1,
|
|
id: 'bash-success',
|
|
type: 'function',
|
|
function: {
|
|
name: 'bash',
|
|
arguments: JSON.stringify({ command: 'printf SECOND_OK' }),
|
|
},
|
|
},
|
|
],
|
|
}, null);
|
|
writeChunk({}, 'tool_calls');
|
|
} else {
|
|
writeChunk({ role: 'assistant', content: 'BASH_BATCH_DONE' }, null);
|
|
writeChunk({}, 'stop');
|
|
}
|
|
response.end('data: [DONE]\n\n');
|
|
});
|
|
});
|
|
await new Promise<void>((resolve, reject) => {
|
|
provider.once('error', reject);
|
|
provider.listen(0, '127.0.0.1', resolve);
|
|
});
|
|
const address = provider.address();
|
|
if (!address || typeof address === 'string') throw new Error('Test Provider did not bind');
|
|
return {
|
|
baseUrl: `http://127.0.0.1:${address.port}/v1`,
|
|
requests,
|
|
close: async () => {
|
|
await new Promise<void>((resolve, reject) => {
|
|
provider.close((error) => error ? reject(error) : resolve());
|
|
provider.closeIdleConnections?.();
|
|
});
|
|
},
|
|
};
|
|
}
|
|
|
|
afterEach(async () => {
|
|
await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
|
|
});
|
|
|
|
describe('Pi Agent Server real process', () => {
|
|
it('resolves runtime packages when the server resource is outside the project module tree', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-agent-server-packaged-layout-'));
|
|
roots.push(root);
|
|
const resourcesDir = path.join(root, 'product-resources', 'resources');
|
|
const configDir = path.join(root, 'config');
|
|
await Promise.all([
|
|
mkdir(resourcesDir, { recursive: true }),
|
|
mkdir(configDir, { recursive: true }),
|
|
]);
|
|
const serverPath = path.join(resourcesDir, 'pi-agent-server.mjs');
|
|
await copyFile(path.resolve('resources/pi-agent-server.mjs'), serverPath);
|
|
const runtimeRoot = path.resolve('node_modules/@earendil-works/pi-coding-agent');
|
|
const server = new PiAgentServerProcess({
|
|
executablePath: process.execPath,
|
|
serverPath,
|
|
runtimeRoot,
|
|
configDir,
|
|
});
|
|
try {
|
|
await expect(server.start()).resolves.toBeUndefined();
|
|
expect(server.processId).toBeTypeOf('number');
|
|
} finally {
|
|
await server.stop().catch(() => undefined);
|
|
}
|
|
}, 10_000);
|
|
|
|
it('executes two write-leased Bash calls without waiting for the HTTP idle timeout', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-agent-server-bash-batch-'));
|
|
roots.push(root);
|
|
const projectPath = path.join(root, 'project');
|
|
const configDir = path.join(root, 'config');
|
|
const sessionDir = path.join(root, 'sessions');
|
|
const extensionDir = path.join(root, 'extensions');
|
|
await Promise.all([
|
|
mkdir(projectPath, { recursive: true }),
|
|
mkdir(configDir, { recursive: true }),
|
|
mkdir(sessionDir, { recursive: true }),
|
|
mkdir(extensionDir, { recursive: true }),
|
|
]);
|
|
const provider = await startBashBatchProvider();
|
|
await Promise.all([
|
|
writeFile(path.join(configDir, 'settings.json'), JSON.stringify({ httpIdleTimeoutMs: 250 })),
|
|
writeFile(path.join(configDir, 'models.json'), JSON.stringify({
|
|
providers: {
|
|
'makelore-test': {
|
|
baseUrl: provider.baseUrl,
|
|
api: 'openai-completions',
|
|
apiKey: '$MAKELORE_TEST_KEY',
|
|
models: [{
|
|
id: 'test-model',
|
|
name: 'Test model',
|
|
reasoning: false,
|
|
input: ['text'],
|
|
contextWindow: 32_000,
|
|
maxTokens: 4_096,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
}],
|
|
},
|
|
},
|
|
})),
|
|
]);
|
|
const promptPath = path.join(root, 'system.md');
|
|
const languagePromptPath = path.join(root, 'language.md');
|
|
await Promise.all([
|
|
writeFile(promptPath, 'Use the provided Bash calls.'),
|
|
writeFile(languagePromptPath, MAKELORE_DEFAULT_LANGUAGE_PROMPT),
|
|
]);
|
|
const extensionHost = new PiManagedExtensionHost();
|
|
const registration = await extensionHost.registerWorker({
|
|
conversationId: 'bash-batch',
|
|
generation: 1,
|
|
projectId: 'project-a',
|
|
projectPath,
|
|
extensionsDir: extensionDir,
|
|
});
|
|
const runtimeRoot = path.resolve('node_modules/@earendil-works/pi-coding-agent');
|
|
const server = new PiAgentServerProcess({
|
|
executablePath: process.execPath,
|
|
serverPath: path.resolve('resources/pi-agent-server.mjs'),
|
|
runtimeRoot,
|
|
configDir,
|
|
});
|
|
const worker = server.createWorker({
|
|
executablePath: process.execPath,
|
|
cliPath: path.join(runtimeRoot, 'dist', 'cli.js'),
|
|
cwd: projectPath,
|
|
configDir,
|
|
sessionDir,
|
|
conversationId: 'bash-batch',
|
|
workerGeneration: 1,
|
|
tools: ['bash'],
|
|
additionalArgs: [
|
|
'--provider', 'makelore-test',
|
|
'--model', 'test-model',
|
|
'--thinking', 'off',
|
|
'--system-prompt', promptPath,
|
|
'--append-system-prompt', languagePromptPath,
|
|
'--extension', registration.extensionPath,
|
|
'--session-id', 'session-bash-batch',
|
|
],
|
|
env: { MAKELORE_TEST_KEY: 'secret-bash-batch', ...registration.env },
|
|
sensitiveValues: ['secret-bash-batch', ...registration.sensitiveValues],
|
|
});
|
|
try {
|
|
await worker.start();
|
|
await extensionHost.bindRun('bash-batch', 1, 'run-bash-batch');
|
|
const settled = new Promise<void>((resolve) => {
|
|
const unsubscribe = worker.subscribe((event) => {
|
|
if (event.type !== 'agent_settled') return;
|
|
unsubscribe();
|
|
resolve();
|
|
});
|
|
});
|
|
const startedAt = Date.now();
|
|
await worker.request({ type: 'prompt', message: 'RUN_BASH_BATCH' });
|
|
await settled;
|
|
expect(Date.now() - startedAt).toBeLessThan(2_000);
|
|
const messagesResponse = await worker.request<{ messages: Array<Record<string, unknown>> }>({
|
|
type: 'get_messages',
|
|
});
|
|
const toolResults = messagesResponse.data?.messages.filter(({ role }) => role === 'toolResult') ?? [];
|
|
const textFor = (toolCallId: string) => {
|
|
const message = toolResults.find((candidate) => candidate.toolCallId === toolCallId);
|
|
const content = Array.isArray(message?.content) ? message.content : [];
|
|
return content.flatMap((item) => (
|
|
item && typeof item === 'object' && !Array.isArray(item)
|
|
&& typeof (item as { text?: unknown }).text === 'string'
|
|
? [(item as { text: string }).text]
|
|
: []
|
|
)).join('\n');
|
|
};
|
|
expect(textFor('bash-timeout')).toContain('Command timed out after 0.05 seconds');
|
|
expect(textFor('bash-success')).toContain('SECOND_OK');
|
|
expect(JSON.stringify(toolResults)).not.toContain('fetch failed');
|
|
expect(provider.requests).toHaveLength(2);
|
|
} finally {
|
|
await worker.stop('test_injection').catch(() => undefined);
|
|
await server.stop().catch(() => undefined);
|
|
await registration.dispose().catch(() => undefined);
|
|
await extensionHost.close().catch(() => undefined);
|
|
await provider.close().catch(() => undefined);
|
|
}
|
|
}, 10_000);
|
|
|
|
it('hosts isolated Conversation threads in one long-lived process', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-agent-server-'));
|
|
roots.push(root);
|
|
const projectPath = path.join(root, 'project');
|
|
const configDir = path.join(root, 'config');
|
|
const sessionDir = path.join(root, 'sessions');
|
|
const extensionDir = path.join(root, 'extensions');
|
|
await Promise.all([
|
|
mkdir(projectPath, { recursive: true }),
|
|
mkdir(configDir, { recursive: true }),
|
|
mkdir(sessionDir, { recursive: true }),
|
|
mkdir(extensionDir, { recursive: true }),
|
|
]);
|
|
const provider = await startHeldProvider();
|
|
await writeFile(path.join(configDir, 'models.json'), JSON.stringify({
|
|
providers: {
|
|
'makelore-test': {
|
|
baseUrl: provider.baseUrl,
|
|
api: 'openai-completions',
|
|
apiKey: '$MAKELORE_TEST_KEY',
|
|
models: [{
|
|
id: 'test-model',
|
|
name: 'Test model',
|
|
reasoning: true,
|
|
input: ['text'],
|
|
contextWindow: 32_000,
|
|
maxTokens: 4_096,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
}],
|
|
},
|
|
},
|
|
}));
|
|
const promptPath = path.join(root, 'system.md');
|
|
const languagePromptPath = path.join(root, 'language.md');
|
|
await Promise.all([
|
|
writeFile(promptPath, 'You are a test Agent.'),
|
|
writeFile(languagePromptPath, MAKELORE_DEFAULT_LANGUAGE_PROMPT),
|
|
]);
|
|
const extensionPath = await materializeMakelorePiExtension(extensionDir);
|
|
const invalidJsonExtensionPath = path.join(extensionDir, 'invalid-json-extension.mjs');
|
|
await writeFile(invalidJsonExtensionPath, `
|
|
export function createMakeloreRuntime() {
|
|
return async function invalidJsonRuntime(pi) {
|
|
pi.on('message_end', async (event) => ({
|
|
message: { ...event.message, makeloreInvalidJson: 1n },
|
|
}));
|
|
};
|
|
}
|
|
`);
|
|
const runtimeRoot = path.resolve('node_modules/@earendil-works/pi-coding-agent');
|
|
const server = new PiAgentServerProcess({
|
|
executablePath: process.execPath,
|
|
serverPath: path.resolve('resources/pi-agent-server.mjs'),
|
|
runtimeRoot,
|
|
configDir,
|
|
});
|
|
|
|
const workerOptions = (
|
|
conversationId: string,
|
|
generation: number,
|
|
selectedExtensionPath = extensionPath,
|
|
): PiWorkerProcessOptions => {
|
|
const contextFile = path.join(root, `${conversationId}.json`);
|
|
return {
|
|
executablePath: process.execPath,
|
|
cliPath: path.join(runtimeRoot, 'dist', 'cli.js'),
|
|
cwd: projectPath,
|
|
configDir,
|
|
sessionDir,
|
|
conversationId,
|
|
workerGeneration: generation,
|
|
additionalArgs: [
|
|
'--provider', 'makelore-test',
|
|
'--model', 'test-model',
|
|
'--thinking', 'medium',
|
|
'--system-prompt', promptPath,
|
|
'--append-system-prompt', languagePromptPath,
|
|
'--extension', selectedExtensionPath,
|
|
'--session-id', `session-${conversationId}`,
|
|
],
|
|
env: {
|
|
MAKELORE_TEST_KEY: `secret-${conversationId}`,
|
|
MAKELORE_PI_BRIDGE_URL: 'http://127.0.0.1:9/v1/worker',
|
|
MAKELORE_PI_WORKER_TOKEN: `token-${conversationId}`,
|
|
MAKELORE_PI_CONTEXT_FILE: contextFile,
|
|
MAKELORE_PI_WORKER_ROLE: 'parent',
|
|
MAKELORE_PI_PROJECT_PATH: projectPath,
|
|
},
|
|
sensitiveValues: [`secret-${conversationId}`, `token-${conversationId}`],
|
|
};
|
|
};
|
|
|
|
await Promise.all(['left', 'right', 'broken'].map(async (conversationId) => {
|
|
await writeFile(path.join(root, `${conversationId}.json`), JSON.stringify({
|
|
conversationId,
|
|
workerGeneration: 1,
|
|
role: 'parent',
|
|
allowedToolNames: [],
|
|
tools: [],
|
|
}));
|
|
}));
|
|
|
|
const left = server.createWorker(workerOptions('left', 1));
|
|
const right = server.createWorker(workerOptions('right', 1));
|
|
try {
|
|
await Promise.all([left.start(), right.start()]);
|
|
const processId = server.processId;
|
|
expect(processId).toBeTypeOf('number');
|
|
expect(server.activeThreadCount).toBe(2);
|
|
|
|
await expect(left.request<{ sessionId: string }>({ type: 'get_state' }))
|
|
.resolves.toMatchObject({ data: { sessionId: 'session-left' } });
|
|
await expect(right.request<{ sessionId: string }>({ type: 'get_state' }))
|
|
.resolves.toMatchObject({ data: { sessionId: 'session-right' } });
|
|
await expect(left.request<{ models: Array<{ id: string }> }>({ type: 'get_available_models' }))
|
|
.resolves.toMatchObject({
|
|
data: { models: expect.arrayContaining([expect.objectContaining({ id: 'test-model' })]) },
|
|
});
|
|
|
|
const settled = (worker: typeof left) => new Promise<void>((resolve) => {
|
|
const unsubscribe = worker.subscribe((event) => {
|
|
if (event.type !== 'agent_settled') return;
|
|
unsubscribe();
|
|
resolve();
|
|
});
|
|
});
|
|
const leftSettled = settled(left);
|
|
const rightSettled = settled(right);
|
|
await Promise.all([
|
|
left.request({ type: 'prompt', message: 'LEFT_THREAD' }),
|
|
right.request({ type: 'prompt', message: 'RIGHT_THREAD' }),
|
|
]);
|
|
await expect.poll(() => provider.activeCount()).toBe(2);
|
|
expect(provider.maxActiveCount()).toBe(2);
|
|
expect(provider.requests.map(({ authorization }) => authorization).sort()).toEqual([
|
|
'Bearer secret-left',
|
|
'Bearer secret-right',
|
|
]);
|
|
expect(provider.requests.some(({ body }) => body.includes('LEFT_THREAD'))).toBe(true);
|
|
expect(provider.requests.some(({ body }) => body.includes('RIGHT_THREAD'))).toBe(true);
|
|
expect(provider.requests.every(({ body }) => body.includes('You are a test Agent.'))).toBe(true);
|
|
expect(provider.requests.every(({ body }) => body.includes('默认使用简体中文'))).toBe(true);
|
|
provider.release();
|
|
await Promise.all([leftSettled, rightSettled]);
|
|
|
|
const broken = server.createWorker(workerOptions('broken', 1, invalidJsonExtensionPath));
|
|
await broken.start();
|
|
const postAcceptanceFailure = new Promise<PiRpcEvent>((resolve) => {
|
|
const unsubscribe = broken.subscribe((event) => {
|
|
if (event.type !== 'makelore_thread_error') return;
|
|
unsubscribe();
|
|
resolve(event);
|
|
});
|
|
});
|
|
await expect(broken.request({ type: 'prompt', message: 'BREAK_AFTER_ACCEPTANCE' }))
|
|
.resolves.toMatchObject({ success: true });
|
|
await expect(postAcceptanceFailure).resolves.toMatchObject({
|
|
type: 'makelore_thread_error',
|
|
code: 'PROMPT_FAILED_AFTER_ACCEPTANCE',
|
|
});
|
|
await expect(broken.request<{ isStreaming: boolean }>({ type: 'get_state' }))
|
|
.resolves.toMatchObject({ data: { isStreaming: false } });
|
|
await broken.stop('test_injection');
|
|
|
|
await left.stop('test_injection');
|
|
expect(server.processId).toBe(processId);
|
|
expect(server.activeThreadCount).toBe(1);
|
|
await expect(right.request<{ sessionId: string }>({ type: 'get_state' }))
|
|
.resolves.toMatchObject({ data: { sessionId: 'session-right' } });
|
|
|
|
const invalidated = new Promise<string>((resolve) => {
|
|
right.subscribeInvalidation((error) => resolve(error.code));
|
|
});
|
|
process.kill(processId as number, 'SIGKILL');
|
|
await expect(invalidated).resolves.toBe('PI_RPC_EXITED');
|
|
await expect.poll(() => server.processId).toBeUndefined();
|
|
await expect(right.request({ type: 'get_state' })).rejects.toMatchObject({
|
|
code: 'PI_RPC_EXITED',
|
|
});
|
|
await right.stop('test_injection');
|
|
|
|
const replacementLeft = server.createWorker(workerOptions('left', 2));
|
|
const replacementRight = server.createWorker(workerOptions('right', 2));
|
|
await Promise.all([replacementLeft.start(), replacementRight.start()]);
|
|
expect(server.processId).toBeTypeOf('number');
|
|
expect(server.processId).not.toBe(processId);
|
|
expect(server.activeThreadCount).toBe(2);
|
|
await expect(replacementRight.request<{ sessionId: string }>({ type: 'get_state' }))
|
|
.resolves.toMatchObject({ data: { sessionId: 'session-right' } });
|
|
await Promise.all([
|
|
replacementLeft.stop('test_injection'),
|
|
replacementRight.stop('test_injection'),
|
|
]);
|
|
} finally {
|
|
await server.stop();
|
|
await provider.close();
|
|
}
|
|
expect(server.processId).toBeUndefined();
|
|
}, 20_000);
|
|
});
|