@dyyz1993/pi-coding-agent 0.70.6 → 0.74.5

package/dist/core/extensions/runner.js

@@ -1,9 +1,7 @@
 /**
  * Extension runner - executes extensions and manages their lifecycle.
  */
-import { randomUUID } from "node:crypto";
 import { theme } from "../../modes/interactive/theme/theme.js";
-import { getCwdDataDir, getGlobalDataDir, getProjectDataDir, getSessionDataDir, resolveProjectRoot, } from "../storage.js";
 // Extension shortcuts compete with canonical keybinding ids from keybindings.json.
 // Only editor-global shortcuts are reserved here. Picker-specific bindings are not.
 const RESERVED_KEYBINDINGS_FOR_EXTENSION_CONFLICTS = [
@@ -66,6 +64,7 @@ const noOpUIContext = {
     onTerminalInput: () => () => { },
     setStatus: () => { },
     setWorkingMessage: () => { },
+    setWorkingVisible: () => { },
     setWorkingIndicator: () => { },
     setHiddenThinkingLabel: () => { },
     setWidget: () => { },
@@ -79,6 +78,7 @@ const noOpUIContext = {
     editor: async () => undefined,
     addAutocompleteProvider: () => { },
     setEditorComponent: () => { },
+    getEditorComponent: () => undefined,
     get theme() {
         return theme;
     },
@@ -92,7 +92,6 @@ export class ExtensionRunner {
     extensions;
     runtime;
     uiContext;
-    uiContextOriginal;
     cwd;
     sessionManager;
     modelRegistry;
@@ -100,13 +99,20 @@ export class ExtensionRunner {
     getModel = () => undefined;
     isIdleFn = () => true;
     getSignalFn = () => undefined;
-    getSessionSignalFn = () => AbortSignal.abort();
     waitForIdleFn = async () => { };
     abortFn = () => { };
     hasPendingMessagesFn = () => false;
     getContextUsageFn = () => undefined;
     compactFn = () => { };
     getSystemPromptFn = () => "";
+    getSessionSignalFn = () => AbortSignal.abort();
+    getExtensionNameFn = () => "";
+    getProjectRootFn = () => this.cwd;
+    getSessionDataDirFn = () => "";
+    getProjectDataDirFn = () => "";
+    getCwdDataDirFn = () => "";
+    getGlobalDataDirFn = () => "";
+    respondUIFn = () => { };
     newSessionHandler = async () => ({ cancelled: false });
     forkHandler = async () => ({ cancelled: false });
     navigateTreeHandler = async () => ({ cancelled: false });
@@ -116,12 +122,30 @@ export class ExtensionRunner {
     shortcutDiagnostics = [];
     commandDiagnostics = [];
     staleMessage;
-    pendingUIResponses = new Map();
+    _fileSnapshotManager = null;
+    setFileSnapshotManager(manager) {
+        this._fileSnapshotManager = manager;
+    }
+    setContextDirFns(fns) {
+        if (fns.getExtensionName)
+            this.getExtensionNameFn = fns.getExtensionName;
+        if (fns.getProjectRoot)
+            this.getProjectRootFn = fns.getProjectRoot;
+        if (fns.getSessionDataDir)
+            this.getSessionDataDirFn = fns.getSessionDataDir;
+        if (fns.getProjectDataDir)
+            this.getProjectDataDirFn = fns.getProjectDataDir;
+        if (fns.getCwdDataDir)
+            this.getCwdDataDirFn = fns.getCwdDataDir;
+        if (fns.getGlobalDataDir)
+            this.getGlobalDataDirFn = fns.getGlobalDataDir;
+        if (fns.respondUI)
+            this.respondUIFn = fns.respondUI;
+    }
     constructor(extensions, runtime, cwd, sessionManager, modelRegistry) {
         this.extensions = extensions;
         this.runtime = runtime;
         this.uiContext = noOpUIContext;
-        this.uiContextOriginal = undefined;
         this.cwd = cwd;
         this.sessionManager = sessionManager;
         this.modelRegistry = modelRegistry;
@@ -131,7 +155,6 @@ export class ExtensionRunner {
         this.runtime.sendMessage = actions.sendMessage;
         this.runtime.sendUserMessage = actions.sendUserMessage;
         this.runtime.appendEntry = actions.appendEntry;
-        this.runtime.foldEntry = actions.foldEntry;
         this.runtime.setSessionName = actions.setSessionName;
         this.runtime.getSessionName = actions.getSessionName;
         this.runtime.setLabel = actions.setLabel;
@@ -143,21 +166,19 @@ export class ExtensionRunner {
         this.runtime.setModel = actions.setModel;
         this.runtime.getThinkingLevel = actions.getThinkingLevel;
         this.runtime.setThinkingLevel = actions.setThinkingLevel;
-        this.runtime.registerChannel = actions.registerChannel;
-        this.runtime.callLLM = actions.callLLM;
-        this.runtime.callLLMStructured = actions.callLLMStructured;
-        this.runtime.background = actions.background;
         // Context actions (required)
         this.getModel = contextActions.getModel;
         this.isIdleFn = contextActions.isIdle;
         this.getSignalFn = contextActions.getSignal;
-        this.getSessionSignalFn = contextActions.getSessionSignal;
         this.abortFn = contextActions.abort;
         this.hasPendingMessagesFn = contextActions.hasPendingMessages;
         this.shutdownHandler = contextActions.shutdown;
         this.getContextUsageFn = contextActions.getContextUsage;
         this.compactFn = contextActions.compact;
         this.getSystemPromptFn = contextActions.getSystemPrompt;
+        if (contextActions.getSessionSignal) {
+            this.getSessionSignalFn = contextActions.getSessionSignal;
+        }
         // Flush provider registrations queued during extension loading
         for (const { name, config, extensionPath } of this.runtime.pendingProviderRegistrations) {
             try {
@@ -178,10 +199,6 @@ export class ExtensionRunner {
             }
         }
         this.runtime.pendingProviderRegistrations = [];
-        // Set registerChannel on the runtime. If this is the throwing stub (non-RPC mode),
-        // pending channel registrations remain queued until bindExtensions() provides the
-        // real implementation via flushPendingChannels().
-        this.runtime.registerChannel = actions.registerChannel;
         // From this point on, provider registration/unregistration takes effect immediately
         // without requiring a /reload.
         this.runtime.registerProvider = (name, config) => {
@@ -216,51 +233,18 @@ export class ExtensionRunner {
         this.switchSessionHandler = async () => ({ cancelled: false });
         this.reloadHandler = async () => { };
     }
-    /**
-     * Flush pending channel registrations with the real registerChannel implementation.
-     * Called from bindCore() and again from bindExtensions() when the real registerChannel
-     * becomes available (e.g. in RPC mode).
-     */
-    flushPendingChannels(registerChannel) {
-        if (this.runtime.pendingChannelRegistrations.length === 0)
-            return;
-        for (const pending of this.runtime.pendingChannelRegistrations) {
-            try {
-                const channel = registerChannel(pending.name);
-                this.runtime.resolvedChannels.set(pending.name, channel);
-                pending.resolve(channel);
-            }
-            catch (err) {
-                pending.reject(err instanceof Error ? err : new Error(String(err)));
-            }
-        }
-        this.runtime.pendingChannelRegistrations = [];
-        this.runtime.registerChannel = registerChannel;
-    }
-    updateRegisterChannel(registerChannel) {
-        this.runtime.registerChannel = registerChannel;
-    }
     setUIContext(uiContext) {
-        this.uiContextOriginal = uiContext;
-        this.uiContext = this.wrapUIForInterception(uiContext ?? noOpUIContext);
+        this.uiContext = uiContext ?? noOpUIContext;
     }
     getUIContext() {
         return this.uiContext;
     }
     hasUI() {
-        return this.uiContextOriginal !== undefined && this.uiContextOriginal !== noOpUIContext;
+        return this.uiContext !== noOpUIContext;
     }
     getExtensionPaths() {
         return this.extensions.map((e) => e.path);
     }
-    getExtensionNameByPath(extensionPath) {
-        for (const ext of this.extensions) {
-            if (ext.path === extensionPath) {
-                return ext.name;
-            }
-        }
-        return undefined;
-    }
     /** Get all registered tools from all extensions (first registration per name wins). */
     getAllRegisteredTools() {
         const toolsByName = new Map();
@@ -333,7 +317,7 @@ export class ExtensionRunner {
     getShortcutDiagnostics() {
         return this.shortcutDiagnostics;
     }
-    invalidate(message = "This extension instance is stale after session replacement or reload.") {
+    invalidate(message = "This extension ctx is stale after session replacement or reload. Do not use a captured pi or command ctx after ctx.newSession(), ctx.fork(), ctx.switchSession(), or ctx.reload(). For newSession, fork, and switchSession, move post-replacement work into withSession and use the ctx passed to withSession. For reload, do not use the old ctx after await ctx.reload().") {
         if (!this.staleMessage) {
             this.staleMessage = message;
             this.runtime.invalidate(message);
@@ -421,7 +405,7 @@ export class ExtensionRunner {
      * Create an ExtensionContext for use in event handlers and tool execution.
      * Context values are resolved at call time, so changes via bindCore/bindUI are reflected.
      */
-    createContext(ext) {
+    createContext() {
         const runner = this;
         const getModel = this.getModel;
         return {
@@ -437,32 +421,6 @@ export class ExtensionRunner {
                 runner.assertActive();
                 return runner.cwd;
             },
-            get extensionName() {
-                runner.assertActive();
-                return ext?.name ?? "unknown";
-            },
-            get projectRoot() {
-                runner.assertActive();
-                return resolveProjectRoot(runner.cwd);
-            },
-            get sessionDataDir() {
-                runner.assertActive();
-                const sessionDir = runner.sessionManager.getSessionDir();
-                const sessionId = runner.sessionManager.getSessionId();
-                return getSessionDataDir(sessionDir, sessionId, ext?.name ?? "unknown");
-            },
-            get projectDataDir() {
-                runner.assertActive();
-                return getProjectDataDir(resolveProjectRoot(runner.cwd), ext?.name ?? "unknown");
-            },
-            get cwdDataDir() {
-                runner.assertActive();
-                return getCwdDataDir(runner.cwd, ext?.name ?? "unknown");
-            },
-            get globalDataDir() {
-                runner.assertActive();
-                return getGlobalDataDir(ext?.name ?? "unknown");
-            },
             get sessionManager() {
                 runner.assertActive();
                 return runner.sessionManager;
@@ -483,10 +441,6 @@ export class ExtensionRunner {
                 runner.assertActive();
                 return runner.getSignalFn();
             },
-            get sessionSignal() {
-                runner.assertActive();
-                return runner.getSessionSignalFn();
-            },
             abort: () => {
                 runner.assertActive();
                 runner.abortFn();
@@ -511,23 +465,49 @@ export class ExtensionRunner {
                 runner.assertActive();
                 return runner.getSystemPromptFn();
             },
-            respondUI: (id, result) => {
+            get extensionName() {
+                runner.assertActive();
+                return runner.getExtensionNameFn();
+            },
+            get projectRoot() {
+                runner.assertActive();
+                return runner.getProjectRootFn();
+            },
+            get sessionDataDir() {
+                runner.assertActive();
+                return runner.getSessionDataDirFn();
+            },
+            get projectDataDir() {
+                runner.assertActive();
+                return runner.getProjectDataDirFn();
+            },
+            get cwdDataDir() {
                 runner.assertActive();
-                runner.respondUI(id, result);
+                return runner.getCwdDataDirFn();
+            },
+            get globalDataDir() {
+                runner.assertActive();
+                return runner.getGlobalDataDirFn();
+            },
+            get sessionSignal() {
+                runner.assertActive();
+                return runner.getSessionSignalFn();
+            },
+            respondUI(id, result) {
+                runner.assertActive();
+                runner.respondUIFn(id, result);
+            },
+            get fileSnapshotManager() {
+                runner.assertActive();
+                return runner._fileSnapshotManager;
             },
         };
     }
-    createCommandContext(commandName) {
-        let cmdExt;
-        if (commandName) {
-            for (const ext of this.extensions) {
-                if (ext.commands.has(commandName)) {
-                    cmdExt = ext;
-                    break;
-                }
-            }
-        }
-        const context = Object.defineProperties({}, Object.getOwnPropertyDescriptors(this.createContext(cmdExt)));
+    createCommandContext() {
+        // Use property descriptors instead of object spread so the guarded getters from
+        // createContext() stay lazy. A spread would eagerly read them once and freeze the
+        // old values into the returned object, bypassing stale-instance checks.
+        const context = Object.defineProperties({}, Object.getOwnPropertyDescriptors(this.createContext()));
         context.waitForIdle = () => {
             this.assertActive();
             return this.waitForIdleFn();
@@ -561,12 +541,12 @@ export class ExtensionRunner {
             event.type === "session_before_tree");
     }
     async emit(event) {
+        const ctx = this.createContext();
         let result;
         for (const ext of this.extensions) {
             const handlers = ext.handlers.get(event.type);
             if (!handlers || handlers.length === 0)
                 continue;
-            const ctx = this.createContext(ext);
             for (const handler of handlers) {
                 try {
                     const handlerResult = await handler(event, ctx);
@@ -591,14 +571,53 @@ export class ExtensionRunner {
         }
         return result;
     }
+    async emitMessageEnd(event) {
+        const ctx = this.createContext();
+        let currentMessage = event.message;
+        let modified = false;
+        for (const ext of this.extensions) {
+            const handlers = ext.handlers.get("message_end");
+            if (!handlers || handlers.length === 0)
+                continue;
+            for (const handler of handlers) {
+                try {
+                    const currentEvent = { ...event, message: currentMessage };
+                    const handlerResult = (await handler(currentEvent, ctx));
+                    if (!handlerResult?.message)
+                        continue;
+                    if (handlerResult.message.role !== currentMessage.role) {
+                        this.emitError({
+                            extensionPath: ext.path,
+                            event: "message_end",
+                            error: "message_end handlers must return a message with the same role",
+                        });
+                        continue;
+                    }
+                    currentMessage = handlerResult.message;
+                    modified = true;
+                }
+                catch (err) {
+                    const message = err instanceof Error ? err.message : String(err);
+                    const stack = err instanceof Error ? err.stack : undefined;
+                    this.emitError({
+                        extensionPath: ext.path,
+                        event: "message_end",
+                        error: message,
+                        stack,
+                    });
+                }
+            }
+        }
+        return modified ? currentMessage : undefined;
+    }
     async emitToolResult(event) {
+        const ctx = this.createContext();
         const currentEvent = { ...event };
         let modified = false;
         for (const ext of this.extensions) {
             const handlers = ext.handlers.get("tool_result");
             if (!handlers || handlers.length === 0)
                 continue;
-            const ctx = this.createContext(ext);
             for (const handler of handlers) {
                 try {
                     const handlerResult = (await handler(currentEvent, ctx));
@@ -639,12 +658,12 @@ export class ExtensionRunner {
         };
     }
     async emitToolCall(event) {
+        const ctx = this.createContext();
         let result;
         for (const ext of this.extensions) {
             const handlers = ext.handlers.get("tool_call");
             if (!handlers || handlers.length === 0)
                 continue;
-            const ctx = this.createContext(ext);
             for (const handler of handlers) {
                 const handlerResult = await handler(event, ctx);
                 if (handlerResult) {
@@ -658,11 +677,11 @@ export class ExtensionRunner {
         return result;
     }
     async emitUserBash(event) {
+        const ctx = this.createContext();
         for (const ext of this.extensions) {
             const handlers = ext.handlers.get("user_bash");
             if (!handlers || handlers.length === 0)
                 continue;
-            const ctx = this.createContext(ext);
             for (const handler of handlers) {
                 try {
                     const handlerResult = await handler(event, ctx);
@@ -684,157 +703,13 @@ export class ExtensionRunner {
         }
         return undefined;
     }
-    createAsyncUIPromise(id, extract) {
-        if (!this.hasHandlers("ui"))
-            return undefined;
-        return new Promise((resolve) => {
-            this.pendingUIResponses.set(id, (result) => {
-                this.pendingUIResponses.delete(id);
-                resolve(extract(result));
-            });
-        });
-    }
-    wrapUIForInterception(original) {
-        return {
-            ...original,
-            confirm: async (title, message, opts) => {
-                if (!this.hasHandlers("ui"))
-                    return original.confirm(title, message, opts);
-                const id = randomUUID();
-                const asyncPromise = this.createAsyncUIPromise(id, (r) => r?.action === "responded" && r.confirmed !== undefined ? r.confirmed : false);
-                const result = await this.emitUIEvent({
-                    type: "ui",
-                    id,
-                    method: "confirm",
-                    title,
-                    message,
-                    signal: opts?.signal,
-                    timeout: opts?.timeout,
-                });
-                if (result?.action === "responded" && result.confirmed !== undefined) {
-                    this.pendingUIResponses.delete(id);
-                    return result.confirmed;
-                }
-                return Promise.race([original.confirm(title, message, opts), asyncPromise]);
-            },
-            select: async (title, options, opts) => {
-                if (!this.hasHandlers("ui"))
-                    return original.select(title, options, opts);
-                const id = randomUUID();
-                const multiple = opts?.multiple;
-                const asyncPromise = this.createAsyncUIPromise(id, (r) => r?.action === "responded" ? r.value : undefined);
-                const result = await this.emitUIEvent({
-                    type: "ui",
-                    id,
-                    method: "select",
-                    title,
-                    options,
-                    multiple,
-                    signal: opts?.signal,
-                    timeout: opts?.timeout,
-                });
-                if (result?.action === "responded") {
-                    this.pendingUIResponses.delete(id);
-                    if (multiple && result.value !== undefined) {
-                        return Array.isArray(result.value) ? result.value : [result.value];
-                    }
-                    return result.value;
-                }
-                return Promise.race([original.select(title, options, opts), asyncPromise]);
-            },
-            input: async (title, placeholder, opts) => {
-                if (!this.hasHandlers("ui"))
-                    return original.input(title, placeholder, opts);
-                const id = randomUUID();
-                const asyncPromise = this.createAsyncUIPromise(id, (r) => r?.action === "responded" ? r.value : undefined);
-                const result = await this.emitUIEvent({
-                    type: "ui",
-                    id,
-                    method: "input",
-                    title,
-                    placeholder,
-                    signal: opts?.signal,
-                    timeout: opts?.timeout,
-                });
-                if (result?.action === "responded") {
-                    this.pendingUIResponses.delete(id);
-                    return result.value;
-                }
-                return Promise.race([original.input(title, placeholder, opts), asyncPromise]);
-            },
-            editor: async (title, prefill) => {
-                if (!this.hasHandlers("ui"))
-                    return original.editor(title, prefill);
-                const id = randomUUID();
-                const asyncPromise = this.createAsyncUIPromise(id, (r) => r?.action === "responded" ? r.value : undefined);
-                const result = await this.emitUIEvent({
-                    type: "ui",
-                    id,
-                    method: "editor",
-                    title,
-                    prefill,
-                });
-                if (result?.action === "responded") {
-                    this.pendingUIResponses.delete(id);
-                    return result.value;
-                }
-                return Promise.race([original.editor(title, prefill), asyncPromise]);
-            },
-            notify: (message, notifyType) => {
-                if (this.hasHandlers("ui")) {
-                    this.emitUIEvent({
-                        type: "ui",
-                        id: randomUUID(),
-                        method: "notify",
-                        title: message,
-                        message,
-                        notifyType,
-                    }).catch(() => { });
-                }
-                return original.notify(message, notifyType);
-            },
-        };
-    }
-    async emitUIEvent(event) {
-        for (const ext of this.extensions) {
-            const ctx = this.createContext(ext);
-            for (const handler of ext.handlers.get("ui") ?? []) {
-                try {
-                    const result = (await handler(event, ctx));
-                    if (result?.action === "responded")
-                        return result;
-                }
-                catch (err) {
-                    const message = err instanceof Error ? err.message : String(err);
-                    const stack = err instanceof Error ? err.stack : undefined;
-                    this.emitError({
-                        extensionPath: ext.path,
-                        event: "ui",
-                        error: message,
-                        stack,
-                    });
-                }
-            }
-        }
-        return undefined;
-    }
-    async emitUI(event) {
-        return this.emitUIEvent(event);
-    }
-    respondUI(id, result) {
-        const resolve = this.pendingUIResponses.get(id);
-        if (resolve) {
-            this.pendingUIResponses.delete(id);
-            resolve(result);
-        }
-    }
     async emitContext(messages) {
+        const ctx = this.createContext();
         let currentMessages = structuredClone(messages);
         for (const ext of this.extensions) {
             const handlers = ext.handlers.get("context");
             if (!handlers || handlers.length === 0)
                 continue;
-            const ctx = this.createContext(ext);
             for (const handler of handlers) {
                 try {
                     const event = { type: "context", messages: currentMessages };
@@ -858,12 +733,12 @@ export class ExtensionRunner {
         return currentMessages;
     }
     async emitBeforeProviderRequest(payload) {
+        const ctx = this.createContext();
         let currentPayload = payload;
         for (const ext of this.extensions) {
             const handlers = ext.handlers.get("before_provider_request");
             if (!handlers || handlers.length === 0)
                 continue;
-            const ctx = this.createContext(ext);
             for (const handler of handlers) {
                 try {
                     const event = {
@@ -891,17 +766,17 @@ export class ExtensionRunner {
     }
     async emitBeforeAgentStart(prompt, images, systemPrompt, systemPromptOptions) {
         let currentSystemPrompt = systemPrompt;
+        const ctx = Object.defineProperties({}, Object.getOwnPropertyDescriptors(this.createContext()));
+        ctx.getSystemPrompt = () => {
+            this.assertActive();
+            return currentSystemPrompt;
+        };
         const messages = [];
         let systemPromptModified = false;
         for (const ext of this.extensions) {
             const handlers = ext.handlers.get("before_agent_start");
             if (!handlers || handlers.length === 0)
                 continue;
-            const ctx = Object.defineProperties({}, Object.getOwnPropertyDescriptors(this.createContext(ext)));
-            ctx.getSystemPrompt = () => {
-                this.assertActive();
-                return currentSystemPrompt;
-            };
             for (const handler of handlers) {
                 try {
                     const event = {
@@ -944,6 +819,7 @@ export class ExtensionRunner {
         return undefined;
     }
     async emitResourcesDiscover(cwd, reason) {
+        const ctx = this.createContext();
         const skillPaths = [];
         const promptPaths = [];
         const themePaths = [];
@@ -951,7 +827,6 @@ export class ExtensionRunner {
             const handlers = ext.handlers.get("resources_discover");
             if (!handlers || handlers.length === 0)
                 continue;
-            const ctx = this.createContext(ext);
             for (const handler of handlers) {
                 try {
                     const event = { type: "resources_discover", cwd, reason };
@@ -983,10 +858,10 @@ export class ExtensionRunner {
     }
     /** Emit input event. Transforms chain, "handled" short-circuits. */
     async emitInput(text, images, source) {
+        const ctx = this.createContext();
         let currentText = text;
         let currentImages = images;
         for (const ext of this.extensions) {
-            const ctx = this.createContext(ext);
             for (const handler of ext.handlers.get("input") ?? []) {
                 try {
                     const event = { type: "input", text: currentText, images: currentImages, source };