Files
makelore/resources/pi-agent-server.mjs

715 lines
27 KiB
JavaScript

import { randomUUID } from 'node:crypto';
import { readFile } from 'node:fs/promises';
import path from 'node:path';
import { pathToFileURL } from 'node:url';
function runtimeRootFromArgs(argv) {
const index = argv.indexOf('--runtime-root');
const value = index >= 0 ? argv[index + 1] : undefined;
if (!value || value.startsWith('--')) throw new Error('--runtime-root is required');
return path.resolve(value);
}
const runtimeRoot = runtimeRootFromArgs(process.argv.slice(2));
const runtimeModule = (...segments) => pathToFileURL(path.join(runtimeRoot, ...segments)).href;
const runtimePackageUrl = pathToFileURL(path.join(runtimeRoot, 'package.json')).href;
const outputGuard = await import(runtimeModule('dist', 'core', 'output-guard.js'));
outputGuard.takeOverStdout();
const [pi, piAi, httpDispatcher, jsonEvents, jsonl, themeModule, shellModule] = await Promise.all([
import(runtimeModule('dist', 'index.js')),
import(import.meta.resolve('@earendil-works/pi-ai', runtimePackageUrl)),
import(runtimeModule('dist', 'core', 'http-dispatcher.js')),
import(runtimeModule('dist', 'modes', 'json-event.js')),
import(runtimeModule('dist', 'modes', 'rpc', 'jsonl.js')),
import(runtimeModule('dist', 'modes', 'interactive', 'theme', 'theme.js')),
import(runtimeModule('dist', 'utils', 'shell.js')),
]);
process.title = 'makelore-pi-agent-server';
process.env.PI_CODING_AGENT = 'true';
process.env.AI_AGENT = 'pi';
httpDispatcher.configureHttpDispatcher();
let networkSettingsKey;
function configureNetwork(settingsManager) {
const httpProxy = settingsManager.getGlobalSettings().httpProxy;
const idleTimeoutMs = settingsManager.getHttpIdleTimeoutMs();
const key = JSON.stringify([httpProxy ?? null, idleTimeoutMs]);
if (networkSettingsKey === key) return;
httpDispatcher.applyHttpProxySettings(httpProxy);
httpDispatcher.configureHttpDispatcher(idleTimeoutMs);
networkSettingsKey = key;
}
const SERVER_CHANNEL = '@makelore/server';
const PROTOCOL_VERSION = 1;
const threads = new Map();
const openingThreads = new Set();
let shuttingDown = false;
let detachInput = () => undefined;
function output(channel, payload) {
outputGuard.writeRawStdout(jsonl.serializeJsonLine({ channel, payload }));
}
function success(id, command, data) {
return data === undefined
? { id, type: 'response', command, success: true }
: { id, type: 'response', command, success: true, data };
}
function failure(id, command, error) {
const message = error instanceof Error ? error.message : String(error);
return { id, type: 'response', command, success: false, error: message };
}
function record(value) {
return value !== null && typeof value === 'object' && !Array.isArray(value)
? value
: null;
}
function requiredString(value, label) {
if (typeof value !== 'string' || !value.trim()) throw new Error(label + ' is required');
return value;
}
function stringArray(value, label) {
if (!Array.isArray(value) || value.some((item) => typeof item !== 'string')) {
throw new Error(label + ' must be a string array');
}
return value;
}
function optionValue(args, name) {
const index = args.lastIndexOf(name);
return index >= 0 ? args[index + 1] : undefined;
}
function optionValues(args, name) {
const values = [];
for (let index = 0; index < args.length; index += 1) {
if (args[index] === name && typeof args[index + 1] === 'string') values.push(args[index + 1]);
}
return values;
}
function referencedEnvironmentNames(value) {
if (typeof value !== 'string') return [];
const names = [];
const pattern = /\$([A-Za-z_][A-Za-z0-9_]*)/g;
let match;
while ((match = pattern.exec(value)) !== null) names.push(match[1]);
return names;
}
async function credentialStoreFor(options, providerId) {
const models = JSON.parse(await readFile(path.join(options.configDir, 'models.json'), 'utf8'));
const provider = record(record(models)?.providers)?.[providerId];
if (!record(provider)) throw new Error('Configured Pi provider is unavailable');
const environment = record(options.env) ?? {};
const references = new Set([
...referencedEnvironmentNames(provider.apiKey),
...Object.values(record(provider.headers) ?? {}).flatMap(referencedEnvironmentNames),
]);
const credentialEnvironment = {};
for (const name of references) {
const value = environment[name];
if (typeof value === 'string') credentialEnvironment[name] = value;
}
const apiKeyName = referencedEnvironmentNames(provider.apiKey)[0];
const apiKey = apiKeyName ? credentialEnvironment[apiKeyName] : undefined;
if (!apiKey) throw new Error('Configured Pi provider credential is unavailable');
const credentials = new piAi.InMemoryCredentialStore();
await credentials.modify(providerId, async () => ({
type: 'api_key',
key: apiKey,
env: credentialEnvironment,
}));
return credentials;
}
async function sessionManagerFor(options, cwd, sessionId, forkSessionId) {
const sessions = await pi.SessionManager.list(cwd, options.sessionDir);
if (forkSessionId) {
const source = sessions.find((item) => item.id === forkSessionId);
if (!source) throw new Error('Fork source Pi session is unavailable');
const target = sessions.find((item) => item.id === sessionId);
if (target) throw new Error('Fork target Pi session already exists');
return pi.SessionManager.forkFrom(source.path, cwd, options.sessionDir, { id: sessionId });
}
const existing = sessions.find((item) => item.id === sessionId);
return existing
? pi.SessionManager.open(existing.path, options.sessionDir)
: pi.SessionManager.create(cwd, options.sessionDir, { id: sessionId });
}
function runtimeExtensionDefaults(options) {
const environment = record(options.env) ?? {};
const value = (key, fallback = '') => (
typeof environment[key] === 'string' && environment[key]
? environment[key]
: fallback
);
return {
bridgeUrl: value('MAKELORE_PI_BRIDGE_URL'),
workerToken: value('MAKELORE_PI_WORKER_TOKEN'),
contextFile: value('MAKELORE_PI_CONTEXT_FILE'),
workerRole: value('MAKELORE_PI_WORKER_ROLE', 'parent'),
projectPath: value('MAKELORE_PI_PROJECT_PATH', options.cwd),
};
}
async function createThreadRuntime(channel, rawOptions) {
const options = record(rawOptions);
if (!options) throw new Error('Thread options are invalid');
const args = stringArray(options.additionalArgs ?? [], 'Thread arguments');
const cwd = path.resolve(requiredString(options.cwd, 'Thread cwd'));
const configDir = path.resolve(requiredString(options.configDir, 'Thread config directory'));
const sessionDir = path.resolve(requiredString(options.sessionDir, 'Thread session directory'));
const providerId = requiredString(optionValue(args, '--provider'), 'Thread provider');
const modelId = requiredString(optionValue(args, '--model'), 'Thread model');
const thinkingLevel = requiredString(optionValue(args, '--thinking'), 'Thread thinking level');
const systemPrompt = requiredString(optionValue(args, '--system-prompt'), 'Thread system prompt');
const appendSystemPrompts = optionValues(args, '--append-system-prompt');
const sessionId = requiredString(optionValue(args, '--session-id'), 'Thread session id');
const forkSessionId = optionValue(args, '--fork');
const [makeloreExtensionPath, ...additionalExtensionPaths] = optionValues(args, '--extension')
.map((value) => path.resolve(value));
const skillPaths = optionValues(args, '--skill').map((value) => path.resolve(value));
if (!makeloreExtensionPath) throw new Error('Thread requires a Makelore extension');
if (appendSystemPrompts.length !== 1) throw new Error('Thread requires exactly one Makelore language prompt');
const extensionModule = await import(pathToFileURL(makeloreExtensionPath).href);
if (typeof extensionModule.createMakeloreRuntime !== 'function') {
throw new Error('Makelore runtime extension factory is unavailable');
}
const initialSessionManager = await sessionManagerFor(
{ ...options, configDir, sessionDir },
cwd,
sessionId,
forkSessionId,
);
const createRuntime = async ({ cwd: runtimeCwd, agentDir, sessionManager, sessionStartEvent }) => {
const credentials = await credentialStoreFor({ ...options, configDir }, providerId);
const modelRuntime = await pi.ModelRuntime.create({
credentials,
modelsPath: path.join(configDir, 'models.json'),
allowModelNetwork: false,
refreshOnCreate: false,
});
const settingsManager = pi.SettingsManager.create(runtimeCwd, agentDir, { projectTrusted: true });
configureNetwork(settingsManager);
const services = await pi.createAgentSessionServices({
cwd: runtimeCwd,
agentDir,
settingsManager,
modelRuntime,
resourceLoaderOptions: {
extensionFactories: [{
name: 'makelore-runtime',
factory: extensionModule.createMakeloreRuntime(runtimeExtensionDefaults({
...options,
cwd: runtimeCwd,
})),
}],
additionalExtensionPaths,
additionalSkillPaths: skillPaths,
noExtensions: true,
noSkills: true,
noPromptTemplates: true,
noThemes: true,
noContextFiles: true,
systemPrompt,
appendSystemPrompt: appendSystemPrompts,
},
});
const extensionErrors = services.resourceLoader.getExtensions().errors;
const errors = [
...services.diagnostics.filter((item) => item.type === 'error').map((item) => item.message),
...extensionErrors.map((item) => item.error),
];
if (errors.length > 0) throw new Error(errors.join('; '));
const model = modelRuntime.getModel(providerId, modelId);
if (!model) throw new Error('Configured Pi model is unavailable');
const created = await pi.createAgentSessionFromServices({
services,
sessionManager,
sessionStartEvent,
model,
thinkingLevel,
tools: Array.isArray(options.tools) ? options.tools : undefined,
});
return { ...created, services, diagnostics: services.diagnostics };
};
const runtime = await pi.createAgentSessionRuntime(createRuntime, {
cwd,
agentDir: configDir,
sessionManager: initialSessionManager,
});
return new AgentThread(channel, runtime);
}
class AgentThread {
constructor(channel, runtime) {
this.channel = channel;
this.runtime = runtime;
this.pendingExtensionRequests = new Map();
this.unsubscribe = () => undefined;
this.unsubscribeBackpressure = () => undefined;
this.closed = false;
}
async start() {
this.runtime.setRebindSession(async () => await this.rebind());
await this.rebind();
}
emit(payload) {
if (!this.closed) output(this.channel, payload);
}
createDialogPromise(options, defaultValue, request, parseResponse) {
if (options?.signal?.aborted) return Promise.resolve(defaultValue);
const id = randomUUID();
return new Promise((resolve) => {
let timeout;
const cleanup = () => {
if (timeout) clearTimeout(timeout);
options?.signal?.removeEventListener('abort', onAbort);
this.pendingExtensionRequests.delete(id);
};
const finish = (value) => {
cleanup();
resolve(value);
};
const onAbort = () => finish(defaultValue);
options?.signal?.addEventListener('abort', onAbort, { once: true });
if (options?.timeout) timeout = setTimeout(onAbort, options.timeout);
this.pendingExtensionRequests.set(id, {
respond: (response) => finish(parseResponse(response)),
cancel: () => finish(defaultValue),
});
this.emit({ type: 'extension_ui_request', id, ...request });
});
}
createExtensionUiContext() {
return {
select: (title, options, dialogOptions) => this.createDialogPromise(
dialogOptions,
undefined,
{ method: 'select', title, options, timeout: dialogOptions?.timeout },
(response) => response.cancelled ? undefined : response.value,
),
confirm: (title, message, dialogOptions) => this.createDialogPromise(
dialogOptions,
false,
{ method: 'confirm', title, message, timeout: dialogOptions?.timeout },
(response) => response.cancelled ? false : response.confirmed === true,
),
input: (title, placeholder, dialogOptions) => this.createDialogPromise(
dialogOptions,
undefined,
{ method: 'input', title, placeholder, timeout: dialogOptions?.timeout },
(response) => response.cancelled ? undefined : response.value,
),
editor: (title, prefill, dialogOptions) => this.createDialogPromise(
dialogOptions,
undefined,
{ method: 'editor', title, prefill, timeout: dialogOptions?.timeout },
(response) => response.cancelled ? undefined : response.value,
),
notify: (message, type) => this.emit({
type: 'extension_ui_request', id: randomUUID(), method: 'notify', message, notifyType: type,
}),
onTerminalInput: () => () => undefined,
setStatus: (key, text) => this.emit({
type: 'extension_ui_request', id: randomUUID(), method: 'setStatus', statusKey: key, statusText: text,
}),
setWorkingMessage: () => undefined,
setWorkingVisible: () => undefined,
setWorkingIndicator: () => undefined,
setHiddenThinkingLabel: () => undefined,
setWidget: (key, content, options) => {
if (content === undefined || Array.isArray(content)) {
this.emit({
type: 'extension_ui_request', id: randomUUID(), method: 'setWidget',
widgetKey: key, widgetLines: content, widgetPlacement: options?.placement,
});
}
},
setFooter: () => undefined,
setHeader: () => undefined,
setTitle: (title) => this.emit({
type: 'extension_ui_request', id: randomUUID(), method: 'setTitle', title,
}),
custom: async () => undefined,
pasteToEditor: (text) => this.emit({
type: 'extension_ui_request', id: randomUUID(), method: 'set_editor_text', text,
}),
setEditorText: (text) => this.emit({
type: 'extension_ui_request', id: randomUUID(), method: 'set_editor_text', text,
}),
getEditorText: () => '',
addAutocompleteProvider: () => undefined,
setEditorComponent: () => undefined,
getEditorComponent: () => undefined,
get theme() { return themeModule.theme; },
getAllThemes: () => [],
getTheme: () => undefined,
setTheme: () => ({ success: false, error: 'Theme switching not supported in RPC mode' }),
getToolsExpanded: () => false,
setToolsExpanded: () => undefined,
};
}
async rebind() {
const session = this.runtime.session;
await session.bindExtensions({
uiContext: this.createExtensionUiContext(),
mode: 'rpc',
commandContextActions: {
waitForIdle: () => session.waitForIdle(),
newSession: async (options) => this.runtime.newSession(options),
fork: async (entryId, options) => {
const result = await this.runtime.fork(entryId, options);
return { cancelled: result.cancelled };
},
navigateTree: async (targetId, options) => {
const result = await session.navigateTree(targetId, options);
return { cancelled: result.cancelled };
},
switchSession: async (sessionPath, options) => this.runtime.switchSession(sessionPath, options),
reload: async () => await session.reload(),
},
shutdownHandler: () => undefined,
onError: (error) => this.emit({
type: 'extension_error',
extensionPath: error.extensionPath,
event: error.event,
error: error.error,
}),
});
this.unsubscribe();
this.unsubscribeBackpressure();
this.unsubscribe = session.subscribe((event) => this.emit(jsonEvents.toJsonEvent(event)));
this.unsubscribeBackpressure = session.agent.subscribe(async () => {
await outputGuard.waitForRawStdoutBackpressure();
});
}
extensionResponse(command) {
const pending = this.pendingExtensionRequests.get(command.id);
if (pending) pending.respond(command);
}
async handle(command) {
if (this.closed) throw new Error('Agent thread is closed');
const session = this.runtime.session;
const id = command.id;
switch (command.type) {
case 'extension_ui_response':
this.extensionResponse(command);
return undefined;
case 'prompt': {
let accepted = false;
void session.prompt(command.message, {
images: command.images,
streamingBehavior: command.streamingBehavior,
source: 'rpc',
preflightResult: (ok) => {
if (!ok || accepted) return;
accepted = true;
this.emit(success(id, 'prompt'));
},
}).catch((error) => {
if (!accepted) this.emit(failure(id, 'prompt', error));
});
return undefined;
}
case 'steer':
await session.steer(command.message, command.images);
return success(id, 'steer');
case 'follow_up':
await session.followUp(command.message, command.images);
return success(id, 'follow_up');
case 'abort':
await session.abort();
return success(id, 'abort');
case 'new_session': {
const result = await this.runtime.newSession(
command.parentSession ? { parentSession: command.parentSession } : undefined,
);
if (!result.cancelled) await this.rebind();
return success(id, 'new_session', result);
}
case 'get_state':
return success(id, 'get_state', {
model: session.model,
thinkingLevel: session.thinkingLevel,
isStreaming: session.isStreaming,
isCompacting: session.isCompacting,
steeringMode: session.steeringMode,
followUpMode: session.followUpMode,
sessionFile: session.sessionFile,
sessionId: session.sessionId,
sessionName: session.sessionName,
autoCompactionEnabled: session.autoCompactionEnabled,
messageCount: session.messages.length,
pendingMessageCount: session.pendingMessageCount,
});
case 'set_model': {
const model = session.modelRuntime.getAvailableSnapshot().find((candidate) => (
candidate.provider === command.provider && candidate.id === command.modelId
));
if (!model) throw new Error('Model not found: ' + command.provider + '/' + command.modelId);
await session.setModel(model);
return success(id, 'set_model', model);
}
case 'cycle_model':
return success(id, 'cycle_model', await session.cycleModel() ?? null);
case 'get_available_models':
return success(id, 'get_available_models', {
models: session.modelRuntime.getAvailableSnapshot(),
});
case 'set_thinking_level':
session.setThinkingLevel(command.level);
return success(id, 'set_thinking_level');
case 'cycle_thinking_level':
{
const level = session.cycleThinkingLevel();
return success(id, 'cycle_thinking_level', level ? { level } : null);
}
case 'get_available_thinking_levels':
return success(id, 'get_available_thinking_levels', {
levels: session.getAvailableThinkingLevels(),
});
case 'set_steering_mode':
session.setSteeringMode(command.mode);
return success(id, 'set_steering_mode');
case 'set_follow_up_mode':
session.setFollowUpMode(command.mode);
return success(id, 'set_follow_up_mode');
case 'compact':
return success(id, 'compact', await session.compact(command.customInstructions));
case 'set_auto_compaction':
session.setAutoCompactionEnabled(command.enabled);
return success(id, 'set_auto_compaction');
case 'set_auto_retry':
session.setAutoRetryEnabled(command.enabled);
return success(id, 'set_auto_retry');
case 'abort_retry':
session.abortRetry();
return success(id, 'abort_retry');
case 'bash': {
const eventResult = await session.extensionRunner.emitUserBash({
type: 'user_bash',
command: command.command,
excludeFromContext: command.excludeFromContext ?? false,
cwd: session.sessionManager.getCwd(),
});
if (eventResult?.result) {
session.recordBashResult(command.command, eventResult.result, {
excludeFromContext: command.excludeFromContext,
});
return success(id, 'bash', eventResult.result);
}
return success(id, 'bash', await session.executeBash(command.command, undefined, {
excludeFromContext: command.excludeFromContext,
id,
operations: eventResult?.operations,
}));
}
case 'abort_bash':
session.abortBash();
return success(id, 'abort_bash');
case 'get_session_stats':
return success(id, 'get_session_stats', session.getSessionStats());
case 'export_html':
return success(id, 'export_html', { path: await session.exportToHtml(command.outputPath) });
case 'switch_session': {
const result = await this.runtime.switchSession(command.sessionPath);
if (!result.cancelled) await this.rebind();
return success(id, 'switch_session', result);
}
case 'fork': {
const result = await this.runtime.fork(command.entryId);
if (!result.cancelled) await this.rebind();
return success(id, 'fork', { text: result.selectedText, cancelled: result.cancelled });
}
case 'clone': {
const leafId = session.sessionManager.getLeafId();
if (!leafId) throw new Error('Cannot clone session without a current entry');
const result = await this.runtime.fork(leafId, { position: 'at' });
if (!result.cancelled) await this.rebind();
return success(id, 'clone', { cancelled: result.cancelled });
}
case 'get_fork_messages':
return success(id, 'get_fork_messages', { messages: session.getUserMessagesForForking() });
case 'get_entries': {
let entries = session.sessionManager.getEntries();
if (command.since !== undefined) {
const index = entries.findIndex((entry) => entry.id === command.since);
if (index < 0) throw new Error('Entry not found: ' + command.since);
entries = entries.slice(index + 1);
}
return success(id, 'get_entries', { entries, leafId: session.sessionManager.getLeafId() });
}
case 'get_tree':
return success(id, 'get_tree', {
tree: session.sessionManager.getTree(), leafId: session.sessionManager.getLeafId(),
});
case 'get_last_assistant_text':
return success(id, 'get_last_assistant_text', { text: session.getLastAssistantText() });
case 'set_session_name':
session.setSessionName(requiredString(command.name, 'Session name'));
return success(id, 'set_session_name');
case 'get_messages':
return success(id, 'get_messages', { messages: session.messages });
case 'get_commands': {
const commands = [];
for (const item of session.extensionRunner.getRegisteredCommands()) {
commands.push({
name: item.invocationName,
description: item.description,
source: 'extension',
sourceInfo: item.sourceInfo,
});
}
for (const item of session.promptTemplates) {
commands.push({
name: item.name, description: item.description, source: 'prompt', sourceInfo: item.sourceInfo,
});
}
for (const item of session.resourceLoader.getSkills().skills) {
commands.push({
name: 'skill:' + item.name,
description: item.description,
source: 'skill',
sourceInfo: item.sourceInfo,
});
}
return success(id, 'get_commands', { commands });
}
default:
throw new Error('Unknown command: ' + command.type);
}
}
async dispose() {
if (this.closed) return;
this.closed = true;
this.unsubscribe();
this.unsubscribeBackpressure();
for (const pending of this.pendingExtensionRequests.values()) pending.cancel();
this.pendingExtensionRequests.clear();
await this.runtime.dispose();
}
}
async function handleControl(command) {
switch (command.type) {
case 'server_initialize':
return success(command.id, command.type, {
protocolVersion: PROTOCOL_VERSION,
threadCount: threads.size,
});
case 'thread_open': {
const threadId = requiredString(command.threadId, 'Thread id');
if (threads.has(threadId) || openingThreads.has(threadId)) {
throw new Error('Agent thread already exists');
}
openingThreads.add(threadId);
let thread;
try {
thread = await createThreadRuntime(threadId, command.options);
await thread.start();
threads.set(threadId, thread);
return success(command.id, command.type, { threadId });
} catch (error) {
await thread?.dispose().catch(() => undefined);
throw error;
} finally {
openingThreads.delete(threadId);
}
}
case 'thread_close': {
const threadId = requiredString(command.threadId, 'Thread id');
const thread = threads.get(threadId);
threads.delete(threadId);
if (thread) await thread.dispose();
return success(command.id, command.type, { threadId, closed: Boolean(thread) });
}
case 'server_shutdown':
setTimeout(() => { void shutdown(0); }, 0);
return success(command.id, command.type, { threadCount: threads.size });
default:
throw new Error('Unknown server command: ' + command.type);
}
}
async function handleEnvelope(line) {
let envelope;
try {
envelope = JSON.parse(line);
} catch (error) {
process.stderr.write('Invalid Makelore Agent Server JSON: ' + String(error) + '\n');
await shutdown(1);
return;
}
const value = record(envelope);
const payload = record(value?.payload);
if (!value || typeof value.channel !== 'string' || !payload || typeof payload.type !== 'string') {
process.stderr.write('Invalid Makelore Agent Server envelope\n');
await shutdown(1);
return;
}
try {
const response = value.channel === SERVER_CHANNEL
? await handleControl(payload)
: await threads.get(value.channel)?.handle(payload);
if (response) output(value.channel, response);
else if (value.channel !== SERVER_CHANNEL && !threads.has(value.channel)) {
output(value.channel, failure(payload.id, payload.type, new Error('Agent thread is unavailable')));
}
await outputGuard.waitForRawStdoutBackpressure();
} catch (error) {
output(value.channel, failure(payload.id, payload.type, error));
await outputGuard.waitForRawStdoutBackpressure();
}
}
async function shutdown(exitCode) {
if (shuttingDown) return;
shuttingDown = true;
detachInput();
const active = [...threads.values()];
threads.clear();
await Promise.allSettled(active.map((thread) => thread.dispose()));
shellModule.killTrackedDetachedChildren();
process.stdin.pause();
await outputGuard.flushRawStdout();
process.exit(exitCode);
}
const onInputEnd = () => { void shutdown(0); };
process.stdin.on('end', onInputEnd);
const detachJsonl = jsonl.attachJsonlLineReader(process.stdin, (line) => {
void handleEnvelope(line);
});
detachInput = () => {
detachJsonl();
process.stdin.off('end', onInputEnd);
};
for (const signal of process.platform === 'win32' ? ['SIGTERM'] : ['SIGTERM', 'SIGHUP']) {
process.on(signal, () => { void shutdown(signal === 'SIGHUP' ? 129 : 143); });
}
await new Promise(() => undefined);