@larose/pi-web 0.3.0

package/dist/server/runtime-registry.js

@@ -0,0 +1,703 @@
+import { createAgentSessionFromServices, createAgentSessionRuntime, createAgentSessionServices, getAgentDir, ModelRuntime, } from "@earendil-works/pi-coding-agent";
+import { AppError } from "./errors.js";
+import { EventBuffer } from "./event-buffer.js";
+import { WebExtensionUI, } from "./extension-ui.js";
+function errorMessage(error) {
+    return error instanceof Error ? error.message : String(error);
+}
+function sessionCommands(session) {
+    const commands = [];
+    for (const command of session.extensionRunner.getRegisteredCommands()) {
+        commands.push({
+            name: command.invocationName,
+            ...(command.description === undefined ? {} : { description: command.description }),
+            source: "extension",
+            sourceInfo: command.sourceInfo,
+        });
+    }
+    for (const template of session.promptTemplates) {
+        commands.push({
+            name: template.name,
+            ...(template.description === undefined ? {} : { description: template.description }),
+            source: "prompt",
+            sourceInfo: template.sourceInfo,
+        });
+    }
+    for (const skill of session.resourceLoader.getSkills().skills) {
+        commands.push({
+            name: `skill:${skill.name}`,
+            ...(skill.description === undefined ? {} : { description: skill.description }),
+            source: "skill",
+            sourceInfo: skill.sourceInfo,
+        });
+    }
+    return commands;
+}
+class SdkRuntimeSession {
+    extensionUI = new WebExtensionUI();
+    runtime;
+    listeners = new Set();
+    replacementListeners = new Set();
+    unsubscribeSession;
+    binding;
+    commandCatalogue;
+    currentSessionId;
+    currentSessionManager;
+    disposed = false;
+    constructor(runtime) {
+        this.runtime = runtime;
+        this.commandCatalogue = sessionCommands(runtime.session);
+        this.currentSessionId = runtime.session.sessionId;
+        this.currentSessionManager = runtime.session.sessionManager;
+        this.attachSession(runtime.session);
+        this.binding = this.bindSession(runtime.session);
+        void this.binding.catch(() => undefined);
+        runtime.setBeforeSessionInvalidate(() => this.extensionUI.reset());
+        runtime.setRebindSession(async (session) => {
+            const previousId = this.currentSessionId;
+            this.commandCatalogue = sessionCommands(session);
+            this.currentSessionId = session.sessionId;
+            this.currentSessionManager = session.sessionManager;
+            this.attachSession(session);
+            this.binding = this.bindSession(session);
+            void this.binding.catch(() => undefined);
+            for (const listener of this.replacementListeners) {
+                listener(previousId, session.sessionId);
+            }
+            await this.binding;
+        });
+    }
+    get sessionId() {
+        return this.currentSessionId;
+    }
+    get sessionManager() {
+        return this.currentSessionManager;
+    }
+    get sessionName() {
+        return this.runtime.session.sessionName;
+    }
+    get cwd() {
+        return this.runtime.session.sessionManager.getCwd();
+    }
+    get isStreaming() {
+        return this.runtime.session.isStreaming;
+    }
+    get isCompacting() {
+        return this.runtime.session.isCompacting;
+    }
+    get pendingMessageCount() {
+        return this.runtime.session.pendingMessageCount;
+    }
+    getSteeringMessages() {
+        return [...this.runtime.session.getSteeringMessages()];
+    }
+    getFollowUpMessages() {
+        return [...this.runtime.session.getFollowUpMessages()];
+    }
+    get messages() {
+        return this.runtime.session.messages;
+    }
+    get commands() {
+        return this.commandCatalogue;
+    }
+    get model() {
+        return this.runtime.session.model;
+    }
+    get thinkingLevel() {
+        return this.runtime.session.thinkingLevel;
+    }
+    get autoCompactionEnabled() {
+        return this.runtime.session.autoCompactionEnabled;
+    }
+    get agent() {
+        return this.runtime.session.agent;
+    }
+    get ready() {
+        return this.binding;
+    }
+    getSessionStats() {
+        return this.runtime.session.getSessionStats();
+    }
+    subscribe(listener) {
+        this.listeners.add(listener);
+        return () => this.listeners.delete(listener);
+    }
+    onSessionReplaced(listener) {
+        this.replacementListeners.add(listener);
+        return () => this.replacementListeners.delete(listener);
+    }
+    prompt(message, options) {
+        return this.runtime.session.prompt(message, options);
+    }
+    steer(message) {
+        return this.runtime.session.steer(message);
+    }
+    clearQueue() {
+        return this.runtime.session.clearQueue();
+    }
+    async restoreQueue(queue) {
+        for (const message of queue.steering) {
+            await this.runtime.session.sendUserMessage(message, {
+                deliverAs: "steer",
+                expandPromptTemplates: false,
+            });
+        }
+        for (const message of queue.followUp) {
+            await this.runtime.session.sendUserMessage(message, {
+                deliverAs: "followUp",
+                expandPromptTemplates: false,
+            });
+        }
+    }
+    abort() {
+        return this.runtime.session.abort();
+    }
+    setSessionName(name) {
+        this.runtime.session.setSessionName(name);
+    }
+    async dispose() {
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
+        this.extensionUI.dispose();
+        await this.runtime.dispose();
+        this.unsubscribeSession?.();
+        this.listeners.clear();
+        this.replacementListeners.clear();
+    }
+    attachSession(session) {
+        this.unsubscribeSession?.();
+        this.unsubscribeSession = session.subscribe((event) => {
+            for (const listener of this.listeners) {
+                listener(event);
+            }
+        });
+    }
+    bindSession(session) {
+        return session.bindExtensions({
+            mode: "rpc",
+            uiContext: this.extensionUI.context,
+            commandContextActions: {
+                waitForIdle: () => session.waitForIdle(),
+                newSession: (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: (sessionPath, options) => this.runtime.switchSession(sessionPath, options),
+                reload: async () => {
+                    await session.reload();
+                    this.commandCatalogue = sessionCommands(session);
+                },
+            },
+            shutdownHandler: () => this.extensionUI.context.notify("An extension requested shutdown. Stop Pi Web from its server process.", "warning"),
+            onError: (error) => this.extensionUI.context.notify(`Extension error in ${error.extensionPath} (${error.event}): ${error.error}`, "error"),
+        });
+    }
+}
+export class RuntimeHandle {
+    events;
+    session;
+    initialCwd;
+    idleTimeoutMs;
+    onIdle;
+    onSessionIdChanged;
+    unsubscribeSession;
+    unsubscribeExtensionUI;
+    unsubscribeReplacement;
+    sessionId;
+    activeTools = new Set();
+    activePrompt;
+    connectedClients = 0;
+    disposed = false;
+    idleTimer;
+    lastSettledAt = Date.now();
+    operationTail = Promise.resolve();
+    queuedMessages = 0;
+    retrying = false;
+    suppressQueueUpdates = false;
+    suppressedQueueUpdate;
+    constructor(session, cwd, idleTimeoutMs, eventCapacity, onIdle, onSessionIdChanged = () => undefined) {
+        this.session = session;
+        this.sessionId = session.sessionId;
+        this.initialCwd = cwd;
+        this.idleTimeoutMs = idleTimeoutMs;
+        this.events = new EventBuffer(eventCapacity);
+        this.onIdle = onIdle;
+        this.onSessionIdChanged = onSessionIdChanged;
+        this.unsubscribeSession = session.subscribe((event) => this.onSessionEvent(event));
+        this.unsubscribeExtensionUI = session.extensionUI?.setPublisher((event) => {
+            this.events.publish(event);
+            if (event.type === "extension_ui_request" &&
+                (event.method === "select" ||
+                    event.method === "confirm" ||
+                    event.method === "input" ||
+                    event.method === "editor")) {
+                this.cancelEviction();
+            }
+            if (event.type === "extension_ui_closed") {
+                this.scheduleEviction();
+            }
+        });
+        this.unsubscribeReplacement = session.onSessionReplaced?.((previousId, sessionId) => {
+            this.sessionId = sessionId;
+            this.onSessionIdChanged(this, previousId, sessionId, session.sessionManager);
+            this.events.publish({ type: "session_replaced", previousId, sessionId });
+        });
+        this.scheduleEviction();
+    }
+    get id() {
+        return this.sessionId;
+    }
+    get cwd() {
+        return this.session.cwd ?? this.initialCwd;
+    }
+    snapshot() {
+        const streamingMessage = this.session.agent.state.streamingMessage;
+        const stats = this.session.getSessionStats();
+        return {
+            id: this.id,
+            cwd: this.cwd,
+            ...(this.session.sessionName ? { sessionName: this.session.sessionName } : {}),
+            isStreaming: this.session.isStreaming,
+            isCompacting: this.session.isCompacting,
+            isWorking: this.isWorking(),
+            pendingMessageCount: this.session.pendingMessageCount,
+            queue: {
+                steering: [...this.session.getSteeringMessages()],
+                followUp: [...this.session.getFollowUpMessages()],
+            },
+            connectedClients: this.connectedClients,
+            model: this.session.model
+                ? {
+                    provider: this.session.model.provider,
+                    id: this.session.model.id,
+                    name: this.session.model.name,
+                }
+                : null,
+            thinkingLevel: this.session.thinkingLevel,
+            autoCompactionEnabled: this.session.autoCompactionEnabled,
+            usage: {
+                tokens: stats.tokens,
+                context: stats.contextUsage ?? null,
+            },
+            messages: this.session.messages,
+            commands: this.session.commands,
+            extensionUI: this.session.extensionUI?.snapshot() ?? { pending: [], statuses: [], widgets: [] },
+            ...(streamingMessage ? { streamingMessage } : {}),
+            latestEventId: this.events.latestId,
+        };
+    }
+    attachClient(listener) {
+        this.assertOpen();
+        this.connectedClients += 1;
+        this.cancelEviction();
+        const unsubscribe = this.events.subscribe(listener);
+        let attached = true;
+        return () => {
+            if (!attached) {
+                return;
+            }
+            attached = false;
+            unsubscribe();
+            this.connectedClients -= 1;
+            this.scheduleEviction();
+        };
+    }
+    send(message) {
+        const trimmed = message.trim();
+        if (!trimmed) {
+            return Promise.reject(new AppError("invalid_message", "A message is required"));
+        }
+        return this.enqueue(async () => {
+            this.assertOpen();
+            this.cancelEviction();
+            if (this.session.isStreaming || this.activePrompt) {
+                await this.session.prompt(message, {
+                    streamingBehavior: "steer",
+                    preflightResult: () => undefined,
+                });
+                return { queued: true };
+            }
+            await this.startPrompt(message);
+            return { queued: false };
+        });
+    }
+    removePendingSteering(index, expectedMessage, expectedQueue) {
+        return this.enqueue(async () => {
+            this.assertOpen();
+            this.cancelEviction();
+            const steering = [...this.session.getSteeringMessages()];
+            const unchanged = steering.length === expectedQueue.length &&
+                steering.every((message, queueIndex) => message === expectedQueue[queueIndex]);
+            if (!unchanged || steering[index] !== expectedMessage) {
+                throw new AppError("stale_pending_message", "The pending message queue has changed", 409);
+            }
+            this.suppressQueueUpdates = true;
+            try {
+                const queue = this.session.clearQueue();
+                queue.steering.splice(index, 1);
+                await this.session.restoreQueue(queue);
+            }
+            finally {
+                this.suppressQueueUpdates = false;
+                const update = this.suppressedQueueUpdate;
+                this.suppressedQueueUpdate = undefined;
+                if (update) {
+                    this.events.publish({ type: "agent_event", event: update });
+                }
+                this.scheduleEviction();
+            }
+        });
+    }
+    respondToExtensionUI(response) {
+        this.assertOpen();
+        if (!this.session.extensionUI?.respond(response)) {
+            throw new AppError("unknown_ui_request", `Unknown extension UI request: ${response.id}`, 409);
+        }
+    }
+    setSessionName(name) {
+        this.assertOpen();
+        this.session.setSessionName(name ?? "");
+        return this.session.sessionName;
+    }
+    abort() {
+        return this.enqueue(async () => {
+            this.assertOpen();
+            this.cancelEviction();
+            const queued = this.session.clearQueue();
+            this.queuedMessages = 0;
+            this.session.extensionUI?.cancelPending();
+            await this.session.abort();
+            this.retrying = false;
+            this.activeTools.clear();
+            this.lastSettledAt = Date.now();
+            this.scheduleEviction();
+            return { restoredMessages: [...queued.steering, ...queued.followUp] };
+        });
+    }
+    isEvictable() {
+        return !this.disposed && this.connectedClients === 0 && !this.isWorking();
+    }
+    async dispose() {
+        if (this.disposed) {
+            return;
+        }
+        this.disposed = true;
+        this.cancelEviction();
+        this.session.clearQueue();
+        try {
+            await this.session.abort();
+        }
+        finally {
+            this.events.publish({ type: "runtime_disposed" });
+            this.unsubscribeSession();
+            this.unsubscribeExtensionUI?.();
+            this.unsubscribeReplacement?.();
+            await this.session.dispose();
+        }
+    }
+    enqueue(operation) {
+        const result = this.operationTail.then(operation, operation);
+        this.operationTail = result.then(() => undefined, () => undefined);
+        return result;
+    }
+    async startPrompt(message) {
+        let preflightSettled = false;
+        let acceptPrompt;
+        let rejectPrompt;
+        const accepted = new Promise((resolve, reject) => {
+            acceptPrompt = resolve;
+            rejectPrompt = reject;
+        });
+        const completion = (async () => {
+            await this.session.ready;
+            await this.session.prompt(message, {
+                preflightResult: (success) => {
+                    if (preflightSettled) {
+                        return;
+                    }
+                    preflightSettled = true;
+                    if (success) {
+                        acceptPrompt();
+                    }
+                    else {
+                        rejectPrompt(new AppError("prompt_rejected", "Pi rejected the message", 409));
+                    }
+                },
+            });
+        })();
+        this.activePrompt = completion;
+        void completion
+            .then(() => {
+            if (preflightSettled) {
+                return;
+            }
+            preflightSettled = true;
+            acceptPrompt();
+        })
+            .catch((error) => {
+            if (!preflightSettled) {
+                preflightSettled = true;
+                rejectPrompt(error);
+            }
+            else {
+                this.events.publish({ type: "runtime_error", message: errorMessage(error) });
+            }
+        })
+            .finally(() => {
+            if (this.activePrompt === completion) {
+                this.activePrompt = undefined;
+            }
+            this.scheduleEviction();
+        });
+        await accepted;
+    }
+    onSessionEvent(event) {
+        if (this.disposed) {
+            return;
+        }
+        switch (event.type) {
+            case "agent_start":
+                this.cancelEviction();
+                break;
+            case "tool_execution_start":
+                this.activeTools.add(event.toolCallId);
+                this.cancelEviction();
+                break;
+            case "tool_execution_end":
+                this.activeTools.delete(event.toolCallId);
+                break;
+            case "queue_update":
+                this.queuedMessages = event.steering.length + event.followUp.length;
+                if (this.queuedMessages > 0) {
+                    this.cancelEviction();
+                }
+                break;
+            case "compaction_start":
+            case "auto_retry_start":
+            case "summarization_retry_scheduled":
+            case "summarization_retry_attempt_start":
+                this.retrying = true;
+                this.cancelEviction();
+                break;
+            case "compaction_end":
+                this.retrying = event.willRetry;
+                break;
+            case "auto_retry_end":
+            case "summarization_retry_finished":
+                this.retrying = false;
+                break;
+            case "agent_settled":
+                this.retrying = false;
+                this.activeTools.clear();
+                this.lastSettledAt = Date.now();
+                this.scheduleEviction();
+                break;
+        }
+        if (event.type === "queue_update" && this.suppressQueueUpdates) {
+            this.suppressedQueueUpdate = event;
+            return;
+        }
+        this.events.publish({ type: "agent_event", event });
+    }
+    isWorking() {
+        return (this.session.isStreaming ||
+            this.session.isCompacting ||
+            this.activePrompt !== undefined ||
+            this.activeTools.size > 0 ||
+            this.queuedMessages > 0 ||
+            (this.session.extensionUI?.snapshot().pending.length ?? 0) > 0 ||
+            this.retrying);
+    }
+    scheduleEviction() {
+        this.cancelEviction();
+        if (!this.isEvictable()) {
+            return;
+        }
+        const elapsed = Date.now() - this.lastSettledAt;
+        const delay = Math.max(0, this.idleTimeoutMs - elapsed);
+        this.idleTimer = setTimeout(() => {
+            this.idleTimer = undefined;
+            if (this.isEvictable()) {
+                this.onIdle(this);
+            }
+        }, delay);
+        this.idleTimer.unref();
+    }
+    cancelEviction() {
+        if (!this.idleTimer) {
+            return;
+        }
+        clearTimeout(this.idleTimer);
+        this.idleTimer = undefined;
+    }
+    assertOpen() {
+        if (this.disposed) {
+            throw new AppError("runtime_closed", "The session runtime is closed", 409);
+        }
+    }
+}
+export class RuntimeRegistry {
+    repository;
+    options;
+    handles = new Map();
+    starts = new Map();
+    deletions = new Map();
+    closing = false;
+    constructor(repository, options) {
+        this.repository = repository;
+        this.options = {
+            idleTimeoutMs: options.idleTimeoutMs ?? 10 * 60_000,
+            eventCapacity: options.eventCapacity ?? 512,
+            createSession: options.createSession,
+        };
+    }
+    async get(id) {
+        if (this.closing) {
+            throw new AppError("server_stopping", "The server is stopping", 503);
+        }
+        const deleting = this.deletions.get(id);
+        if (deleting) {
+            await deleting;
+        }
+        if (this.closing) {
+            throw new AppError("server_stopping", "The server is stopping", 503);
+        }
+        const existing = this.handles.get(id);
+        if (existing) {
+            return existing;
+        }
+        const starting = this.starts.get(id);
+        if (starting) {
+            return starting;
+        }
+        const start = this.start(id);
+        this.starts.set(id, start);
+        try {
+            return await start;
+        }
+        finally {
+            if (this.starts.get(id) === start) {
+                this.starts.delete(id);
+            }
+        }
+    }
+    getActive(id) {
+        return this.handles.get(id);
+    }
+    activeIds() {
+        return [...this.handles.keys()];
+    }
+    async delete(id) {
+        if (this.closing) {
+            throw new AppError("server_stopping", "The server is stopping", 503);
+        }
+        const existing = this.deletions.get(id);
+        if (existing) {
+            return existing;
+        }
+        const deletion = this.deleteSession(id);
+        this.deletions.set(id, deletion);
+        try {
+            await deletion;
+        }
+        finally {
+            if (this.deletions.get(id) === deletion) {
+                this.deletions.delete(id);
+            }
+        }
+    }
+    async close() {
+        if (this.closing) {
+            return;
+        }
+        this.closing = true;
+        await Promise.allSettled([...this.starts.values()]);
+        await Promise.allSettled([...this.deletions.values()]);
+        const handles = [...this.handles.values()];
+        this.handles.clear();
+        await Promise.allSettled(handles.map((handle) => handle.dispose()));
+    }
+    async deleteSession(id) {
+        const starting = this.starts.get(id);
+        if (starting) {
+            await starting.catch(() => undefined);
+        }
+        const handle = this.handles.get(id);
+        if (handle?.isWorking()) {
+            throw new AppError("session_busy", "The session is busy and cannot be deleted", 409);
+        }
+        if (handle) {
+            for (const [key, candidate] of this.handles) {
+                if (candidate === handle) {
+                    this.handles.delete(key);
+                }
+            }
+            await handle.dispose();
+        }
+        await this.repository.delete(id);
+    }
+    async start(id) {
+        const manager = await this.repository.openManager(id);
+        const session = await this.options.createSession(manager);
+        if (this.closing) {
+            await session.dispose();
+            throw new AppError("server_stopping", "The server is stopping", 503);
+        }
+        const handle = new RuntimeHandle(session, manager.getCwd(), this.options.idleTimeoutMs, this.options.eventCapacity, (idleHandle) => void this.evict(idleHandle), (changedHandle, previousId, sessionId, replacementManager) => {
+            this.rekey(changedHandle, previousId, sessionId, replacementManager);
+        });
+        this.handles.set(id, handle);
+        return handle;
+    }
+    rekey(handle, previousId, sessionId, manager) {
+        this.repository.adoptManager(manager);
+        if (this.handles.get(previousId) === handle) {
+            this.handles.delete(previousId);
+        }
+        const existing = this.handles.get(sessionId);
+        if (existing && existing !== handle) {
+            void existing.dispose();
+        }
+        this.handles.set(sessionId, handle);
+    }
+    async evict(handle) {
+        if (!handle.isEvictable()) {
+            return;
+        }
+        for (const [id, candidate] of this.handles) {
+            if (candidate === handle) {
+                this.handles.delete(id);
+            }
+        }
+        await handle.dispose();
+    }
+}
+export async function createRuntimeRegistry(repository) {
+    const modelRuntime = await ModelRuntime.create();
+    const agentDir = getAgentDir();
+    const createSession = async (manager) => {
+        const createRuntime = async ({ cwd, sessionManager, sessionStartEvent, }) => {
+            const services = await createAgentSessionServices({ cwd, agentDir, modelRuntime });
+            const created = await createAgentSessionFromServices({
+                services,
+                sessionManager,
+                ...(sessionStartEvent === undefined ? {} : { sessionStartEvent }),
+            });
+            return { ...created, services, diagnostics: services.diagnostics };
+        };
+        const runtime = await createAgentSessionRuntime(createRuntime, {
+            cwd: manager.getCwd(),
+            agentDir,
+            sessionManager: manager,
+        });
+        return new SdkRuntimeSession(runtime);
+    };
+    return new RuntimeRegistry(repository, { createSession });
+}