@ashx-j/lunr 0.2.17 → 0.2.18

package/dist/builtin-extensions/pi-mcp-adapter/index.js

@@ -1,46 +1,238 @@
+import { existsSync } from "node:fs";
 import { Type } from "typebox";
-import { showStatus, showTools, reconnectServers, authenticateServer, logoutServer, openMcpAuthPanel, openMcpPanel, openMcpSetup } from "./commands.js";
 import { loadMcpConfig } from "./config.js";
-import { buildProxyDescription, createDirectToolExecutor, getMissingConfiguredDirectToolServers, resolveDirectTools } from "./direct-tools.js";
-import { flushMetadataCache, initializeMcp, updateStatusBar } from "./init.js";
-import { loadMetadataCache } from "./metadata-cache.js";
-import { executeAuthComplete, executeAuthStart, executeCall, executeConnect, executeDescribe, executeList, executeSearch, executeStatus, executeUiMessages } from "./proxy-modes.js";
+import { buildProxyDescription, getMissingConfiguredDirectToolServers, resolveDirectTools } from "./direct-tools.js";
+import { getMetadataCachePath, loadMetadataCache } from "./metadata-cache.js";
 import { getConfigPathFromArgv, normalizeDirectToolInputSchema, truncateAtWord } from "./utils.js";
-import { initializeOAuth, shutdownOAuth } from "./mcp-auth-flow.js";
 import { createMcpDirectToolCallRenderer, renderMcpProxyToolCall, renderMcpToolResult } from "./tool-result-renderer.js";
 import { toolErrorOverride } from "./error-signal.js";
+import { abortable, throwIfAborted } from "./abort.js";
+function shouldBackgroundAutostart(config, cache) {
+    const serverEntries = Object.entries(config.mcpServers);
+    if (serverEntries.length === 0)
+        return false;
+    if (!existsSync(getMetadataCachePath()))
+        return true;
+    for (const [, definition] of serverEntries) {
+        const mode = definition.lifecycle ?? "lazy";
+        if (mode === "keep-alive" || mode === "eager")
+            return true;
+    }
+    if (process.env.MCP_DIRECT_TOOLS === "__none__")
+        return false;
+    return getMissingConfiguredDirectToolServers(config, cache).length > 0;
+}
 export default function mcpAdapter(pi) {
     let state = null;
     let initPromise = null;
     let lifecycleGeneration = 0;
-    async function shutdownState(currentState, reason) {
+    let sessionAbort = new AbortController();
+    const initializingStates = new Set();
+    const shutdownPromises = new WeakMap();
+    let heavyModules = null;
+    let heavyModulesPromise = null;
+    const directExecutors = new Map();
+    function loadHeavyModules() {
+        if (heavyModules)
+            return Promise.resolve(heavyModules);
+        if (!heavyModulesPromise) {
+            heavyModulesPromise = (async () => {
+                const [init, authFlow, commands, proxyModes, directToolExecutor] = await Promise.all([
+                    import("./init.js"),
+                    import("./mcp-auth-flow.js"),
+                    import("./commands.js"),
+                    import("./proxy-modes.js"),
+                    import("./direct-tool-executor.js"),
+                ]);
+                const loaded = { init, authFlow, commands, proxyModes, directToolExecutor };
+                heavyModules = loaded;
+                return loaded;
+            })().catch((error) => {
+                heavyModulesPromise = null;
+                throw error;
+            });
+        }
+        return heavyModulesPromise;
+    }
+    function shutdownState(currentState, reason) {
         if (!currentState)
+            return Promise.resolve();
+        const existing = shutdownPromises.get(currentState);
+        if (existing)
+            return existing;
+        const shutdown = (async () => {
+            if (currentState.uiServer) {
+                currentState.uiServer.close(reason);
+                currentState.uiServer = null;
+            }
+            let flushError;
+            try {
+                if (heavyModules) {
+                    heavyModules.init.flushMetadataCache(currentState);
+                }
+            }
+            catch (error) {
+                flushError = error;
+            }
+            try {
+                await currentState.lifecycle.gracefulShutdown();
+            }
+            catch (error) {
+                if (flushError) {
+                    console.error("MCP: graceful shutdown failed after metadata flush error", error);
+                }
+                else {
+                    throw error;
+                }
+            }
+            if (flushError) {
+                throw flushError;
+            }
+        })();
+        shutdownPromises.set(currentState, shutdown);
+        return shutdown;
+    }
+    async function shutdownInitializingStates() {
+        const pending = [...initializingStates];
+        initializingStates.clear();
+        await Promise.all(pending.map((entry) => shutdownState(entry, "initialization_cancelled")));
+    }
+    async function shutdownLoadedRuntime() {
+        if (!heavyModules)
             return;
-        if (currentState.uiServer) {
-            currentState.uiServer.close(reason);
-            currentState.uiServer = null;
+        await heavyModules.authFlow.shutdownOAuth();
+    }
+    async function runInitialize(generation, sessionPi, ctx) {
+        const signal = sessionAbort.signal;
+        const heavy = await loadHeavyModules();
+        throwIfAborted(signal);
+        if (generation !== lifecycleGeneration) {
+            throw new Error("MCP initialization aborted (session changed)");
+        }
+        await heavy.authFlow.initializeOAuth().catch(err => {
+            console.error("MCP OAuth initialization failed:", err);
+        });
+        if (generation !== lifecycleGeneration) {
+            throw new Error("MCP initialization aborted (session changed)");
         }
-        let flushError;
+        let partial;
         try {
-            flushMetadataCache(currentState);
+            const initialized = await heavy.init.initializeMcp(sessionPi, ctx, {
+                signal,
+                onState: (entry) => {
+                    partial = entry;
+                    initializingStates.add(entry);
+                },
+            });
+            if (!partial) {
+                partial = initialized;
+                initializingStates.add(initialized);
+            }
+            throwIfAborted(signal);
+            return initialized;
         }
         catch (error) {
-            flushError = error;
+            if (partial) {
+                initializingStates.delete(partial);
+                await shutdownState(partial, "initialization_failed");
+            }
+            throw error;
         }
-        try {
-            await currentState.lifecycle.gracefulShutdown();
+        finally {
+            if (partial)
+                initializingStates.delete(partial);
         }
-        catch (error) {
-            if (flushError) {
-                console.error("MCP: graceful shutdown failed after metadata flush error", error);
+    }
+    function startSessionInit(generation, sessionPi, ctx) {
+        const promise = runInitialize(generation, sessionPi, ctx);
+        initPromise = promise;
+        promise.then(async (nextState) => {
+            if (generation !== lifecycleGeneration || initPromise !== promise) {
+                try {
+                    await shutdownState(nextState, "stale_session_start");
+                }
+                catch (error) {
+                    console.error("MCP: failed to clean stale session state", error);
+                }
+                return;
             }
-            else {
-                throw error;
+            state = nextState;
+            if (heavyModules) {
+                heavyModules.init.updateStatusBar(nextState);
+            }
+            initPromise = null;
+        }).catch(err => {
+            if (generation !== lifecycleGeneration) {
+                return;
             }
+            if (initPromise !== promise && initPromise !== null) {
+                return;
+            }
+            console.error("MCP initialization failed:", err);
+            initPromise = null;
+        });
+        return promise;
+    }
+    async function ensureState(sessionPi, ctx, signal) {
+        throwIfAborted(signal);
+        if (state)
+            return state;
+        const generation = lifecycleGeneration;
+        if (!initPromise) {
+            startSessionInit(generation, sessionPi, ctx);
         }
-        if (flushError) {
-            throw flushError;
+        const promise = initPromise;
+        if (!promise) {
+            throw new Error("MCP not initialized");
         }
+        try {
+            const nextState = await abortable(promise, signal);
+            if (generation !== lifecycleGeneration) {
+                throw new Error("MCP initialization aborted (session changed)");
+            }
+            if (state)
+                return state;
+            state = nextState;
+            if (initPromise === promise) {
+                initPromise = null;
+            }
+            return state;
+        }
+        catch (error) {
+            if (initPromise === promise && !signal?.aborted) {
+                initPromise = null;
+            }
+            throw error;
+        }
+    }
+    function assertCurrentState(currentState) {
+        if (currentState !== state)
+            throw new Error("MCP session changed before execution");
+    }
+    function createLazyDirectToolExecute(spec) {
+        return async function execute(toolCallId, params, incomingSignal, onUpdate, ctx) {
+            const signal = incomingSignal ? AbortSignal.any([incomingSignal, sessionAbort.signal]) : sessionAbort.signal;
+            let currentState;
+            try {
+                currentState = await ensureState(pi, ctx, signal);
+            }
+            catch (error) {
+                const message = error instanceof Error ? error.message : String(error);
+                return {
+                    content: [{ type: "text", text: `MCP initialization failed: ${message}` }],
+                    details: { error: "init_failed", message },
+                };
+            }
+            const heavy = await loadHeavyModules();
+            throwIfAborted(signal);
+            assertCurrentState(currentState);
+            let executor = directExecutors.get(spec.prefixedName);
+            if (!executor) {
+                executor = heavy.directToolExecutor.createDirectToolExecutor(() => currentState, () => null, spec);
+                directExecutors.set(spec.prefixedName, executor);
+            }
+            return executor(toolCallId, params, signal, onUpdate, ctx);
+        };
     }
     const earlyConfigPath = getConfigPathFromArgv();
     const earlyConfig = loadMcpConfig(earlyConfigPath);
@@ -61,7 +253,7 @@ export default function mcpAdapter(pi) {
             description: spec.description || "(no description)",
             promptSnippet: truncateAtWord(spec.description, 100) || `MCP tool from ${spec.serverName}`,
             parameters: Type.Unsafe(normalizeDirectToolInputSchema(spec.inputSchema)),
-            execute: createDirectToolExecutor(() => state, () => initPromise, spec),
+            execute: createLazyDirectToolExecute(spec),
             renderCall: createMcpDirectToolCallRenderer(spec.prefixedName),
             renderResult: renderMcpToolResult,
         });
@@ -73,13 +265,17 @@ export default function mcpAdapter(pi) {
     });
     pi.on("session_start", async (_event, ctx) => {
         const generation = ++lifecycleGeneration;
+        sessionAbort.abort(new Error("MCP session changed"));
+        sessionAbort = new AbortController();
         const previousState = state;
         state = null;
         initPromise = null;
+        directExecutors.clear();
         try {
             await Promise.all([
                 shutdownState(previousState, "session_restart"),
-                shutdownOAuth(),
+                shutdownInitializingStates(),
+                shutdownLoadedRuntime(),
             ]);
         }
         catch (error) {
@@ -88,44 +284,26 @@ export default function mcpAdapter(pi) {
         if (generation !== lifecycleGeneration) {
             return;
         }
-        await initializeOAuth().catch(err => {
-            console.error("MCP OAuth initialization failed:", err);
-        });
-        const promise = initializeMcp(pi, ctx);
-        initPromise = promise;
-        promise.then(async (nextState) => {
-            if (generation !== lifecycleGeneration || initPromise !== promise) {
-                try {
-                    await shutdownState(nextState, "stale_session_start");
-                }
-                catch (error) {
-                    console.error("MCP: failed to clean stale session state", error);
-                }
-                return;
-            }
-            state = nextState;
-            updateStatusBar(nextState);
-            initPromise = null;
-        }).catch(err => {
-            if (generation !== lifecycleGeneration) {
-                return;
-            }
-            if (initPromise !== promise && initPromise !== null) {
-                return;
-            }
-            console.error("MCP initialization failed:", err);
-            initPromise = null;
-        });
+        const configPath = pi.getFlag("mcp-config") ?? earlyConfigPath;
+        const sessionConfig = loadMcpConfig(configPath, ctx.cwd);
+        const sessionCache = loadMetadataCache();
+        if (!shouldBackgroundAutostart(sessionConfig, sessionCache)) {
+            return;
+        }
+        startSessionInit(generation, pi, ctx);
     });
     pi.on("session_shutdown", async () => {
         ++lifecycleGeneration;
+        sessionAbort.abort(new Error("MCP session ended"));
         const currentState = state;
         state = null;
         initPromise = null;
+        directExecutors.clear();
         try {
             await Promise.all([
                 shutdownState(currentState, "session_shutdown"),
-                shutdownOAuth(),
+                shutdownInitializingStates(),
+                shutdownLoadedRuntime(),
             ]);
         }
         catch (error) {
@@ -137,35 +315,31 @@ export default function mcpAdapter(pi) {
     pi.registerCommand("mcp", {
         description: "Show MCP server status",
         handler: async (args, ctx) => {
-            if (!state && initPromise) {
-                try {
-                    state = await initPromise;
-                }
-                catch (error) {
-                    const message = error instanceof Error ? error.message : String(error);
-                    if (ctx.hasUI)
-                        ctx.ui.notify(`MCP initialization failed: ${message}`, "error");
-                    return;
-                }
+            let currentState;
+            try {
+                currentState = await ensureState(pi, ctx);
             }
-            if (!state) {
+            catch (error) {
+                const message = error instanceof Error ? error.message : String(error);
                 if (ctx.hasUI)
-                    ctx.ui.notify("MCP not initialized", "error");
+                    ctx.ui.notify(`MCP initialization failed: ${message}`, "error");
                 return;
             }
+            const heavy = await loadHeavyModules();
+            assertCurrentState(currentState);
             const parts = args?.trim()?.split(/\s+/) ?? [];
             const subcommand = parts[0] ?? "";
             const targetServer = parts[1];
             const rest = parts.slice(1).join(" ");
             switch (subcommand) {
                 case "reconnect":
-                    await reconnectServers(state, ctx, targetServer);
+                    await heavy.commands.reconnectServers(currentState, ctx, targetServer);
                     break;
                 case "tools":
-                    await showTools(state, ctx);
+                    await heavy.commands.showTools(currentState, ctx);
                     break;
                 case "setup": {
-                    const result = await openMcpSetup(state, pi, ctx, earlyConfigPath, "setup");
+                    const result = await heavy.commands.openMcpSetup(currentState, pi, ctx, earlyConfigPath, "setup");
                     if (result?.configChanged) {
                         await ctx.reload();
                         return;
@@ -179,21 +353,21 @@ export default function mcpAdapter(pi) {
                             ctx.ui.notify("Usage: /mcp logout <server>", "error");
                         return;
                     }
-                    await logoutServer(serverName, state, ctx);
+                    await heavy.commands.logoutServer(serverName, currentState, ctx);
                     break;
                 }
                 case "status":
                 case "":
                 default:
                     if (ctx.hasUI) {
-                        const result = await openMcpPanel(state, pi, ctx, earlyConfigPath);
+                        const result = await heavy.commands.openMcpPanel(currentState, pi, ctx, earlyConfigPath);
                         if (result?.configChanged) {
                             await ctx.reload();
                             return;
                         }
                     }
                     else {
-                        await showStatus(state, ctx);
+                        await heavy.commands.showStatus(currentState, ctx);
                     }
                     break;
             }
@@ -206,27 +380,23 @@ export default function mcpAdapter(pi) {
             if (!serverName && !ctx.hasUI) {
                 return;
             }
-            if (!state && initPromise) {
-                try {
-                    state = await initPromise;
-                }
-                catch (error) {
-                    const message = error instanceof Error ? error.message : String(error);
-                    if (ctx.hasUI)
-                        ctx.ui.notify(`MCP initialization failed: ${message}`, "error");
-                    return;
-                }
+            let currentState;
+            try {
+                currentState = await ensureState(pi, ctx);
             }
-            if (!state) {
+            catch (error) {
+                const message = error instanceof Error ? error.message : String(error);
                 if (ctx.hasUI)
-                    ctx.ui.notify("MCP not initialized", "error");
+                    ctx.ui.notify(`MCP initialization failed: ${message}`, "error");
                 return;
             }
+            const heavy = await loadHeavyModules();
+            assertCurrentState(currentState);
             if (!serverName) {
-                await openMcpAuthPanel(state, pi, ctx, earlyConfigPath);
+                await heavy.commands.openMcpAuthPanel(currentState, pi, ctx, earlyConfigPath);
                 return;
             }
-            await authenticateServer(serverName, state.config, ctx);
+            await heavy.commands.authenticateServer(serverName, currentState.config, ctx);
         },
     });
     if (shouldRegisterProxyTool) {
@@ -248,7 +418,8 @@ export default function mcpAdapter(pi) {
                 action: Type.Optional(Type.String({ description: "Action: 'ui-messages', 'auth-start', or 'auth-complete'" })),
             }),
             renderResult: renderMcpToolResult,
-            async execute(_toolCallId, params, signal, _onUpdate, _ctx) {
+            async execute(_toolCallId, params, incomingSignal, _onUpdate, ctx) {
+                const signal = incomingSignal ? AbortSignal.any([incomingSignal, sessionAbort.signal]) : sessionAbort.signal;
                 let parsedArgs;
                 if (params.args) {
                     try {
@@ -265,26 +436,22 @@ export default function mcpAdapter(pi) {
                         throw error;
                     }
                 }
-                if (!state && initPromise) {
-                    try {
-                        state = await initPromise;
-                    }
-                    catch (error) {
-                        const message = error instanceof Error ? error.message : String(error);
-                        return {
-                            content: [{ type: "text", text: `MCP initialization failed: ${message}` }],
-                            details: { error: "init_failed", message },
-                        };
-                    }
+                let currentState;
+                try {
+                    currentState = await ensureState(pi, ctx, signal);
                 }
-                if (!state) {
+                catch (error) {
+                    const message = error instanceof Error ? error.message : String(error);
                     return {
-                        content: [{ type: "text", text: "MCP not initialized" }],
-                        details: { error: "not_initialized" },
+                        content: [{ type: "text", text: `MCP initialization failed: ${message}` }],
+                        details: { error: "init_failed", message },
                     };
                 }
+                const heavy = await loadHeavyModules();
+                throwIfAborted(signal);
+                assertCurrentState(currentState);
                 if (params.action === "ui-messages") {
-                    return executeUiMessages(state);
+                    return heavy.proxyModes.executeUiMessages(currentState);
                 }
                 if (params.action === "auth-start") {
                     if (!params.server) {
@@ -293,7 +460,7 @@ export default function mcpAdapter(pi) {
                             details: { mode: "auth-start", error: "missing_server" },
                         };
                     }
-                    return executeAuthStart(state, params.server);
+                    return heavy.proxyModes.executeAuthStart(currentState, params.server);
                 }
                 if (params.action === "auth-complete") {
                     if (!params.server) {
@@ -309,24 +476,24 @@ export default function mcpAdapter(pi) {
                             details: { mode: "auth-complete", error: "missing_input" },
                         };
                     }
-                    return executeAuthComplete(state, params.server, input);
+                    return heavy.proxyModes.executeAuthComplete(currentState, params.server, input);
                 }
                 if (params.tool) {
-                    return executeCall(state, params.tool, parsedArgs, params.server, getPiTools, signal);
+                    return heavy.proxyModes.executeCall(currentState, params.tool, parsedArgs, params.server, getPiTools, signal);
                 }
                 if (params.connect) {
-                    return executeConnect(state, params.connect, signal);
+                    return heavy.proxyModes.executeConnect(currentState, params.connect, signal);
                 }
                 if (params.describe) {
-                    return executeDescribe(state, params.describe);
+                    return heavy.proxyModes.executeDescribe(currentState, params.describe);
                 }
                 if (params.search) {
-                    return executeSearch(state, params.search, params.regex, params.server, params.includeSchemas);
+                    return heavy.proxyModes.executeSearch(currentState, params.search, params.regex, params.server, params.includeSchemas);
                 }
                 if (params.server) {
-                    return executeList(state, params.server);
+                    return heavy.proxyModes.executeList(currentState, params.server);
                 }
-                return executeStatus(state);
+                return heavy.proxyModes.executeStatus(currentState);
             },
         });
     }