diff --git a/packages/opencode/src/config/config.ts b/packages/opencode/src/config/config.ts index bf4a6035bd8..640f29225fa 100644 --- a/packages/opencode/src/config/config.ts +++ b/packages/opencode/src/config/config.ts @@ -1041,6 +1041,10 @@ export namespace Config { .positive() .optional() .describe("Timeout in milliseconds for model context protocol (MCP) requests"), + mcp_lazy_load: z + .boolean() + .optional() + .describe("Enable lazy loading of MCP tools (load on-demand instead of at startup)"), }) .optional(), }) diff --git a/packages/opencode/src/mcp/index.ts b/packages/opencode/src/mcp/index.ts index aca0c663152..15d940884c6 100644 --- a/packages/opencode/src/mcp/index.ts +++ b/packages/opencode/src/mcp/index.ts @@ -446,6 +446,18 @@ export namespace MCP { } } + // Skip listTools validation for lazy loading - we'll fetch tools on-demand + const config = await Config.get() + const isLazy = config.experimental?.mcp_lazy_load ?? false + + if (isLazy) { + log.info("create() skipping listTools for lazy server", { key }) + return { + mcpClient, + status, + } + } + const result = await withTimeout(mcpClient.listTools(), mcp.timeout ?? DEFAULT_TIMEOUT).catch((err) => { log.error("failed to get tools from client", { key, error: err }) return undefined @@ -539,10 +551,12 @@ export namespace MCP { s.status[name] = { status: "disabled" } } - export async function tools() { + export async function tools(sessionLoadedTools?: Record) { const result: Record = {} const s = await state() const clientsSnapshot = await clients() + const config = await Config.get() + const isLazy = config.experimental?.mcp_lazy_load ?? false for (const [clientName, client] of Object.entries(clientsSnapshot)) { // Only include tools from connected MCPs (skip disabled ones) @@ -563,7 +577,16 @@ export namespace MCP { if (!toolsResult) { continue } + + // Get the list of loaded tools for this server (if lazy loading) + const loadedToolsForServer = sessionLoadedTools?.[clientName] ?? [] + for (const mcpTool of toolsResult.tools) { + // If lazy loading is enabled, only include tools that are explicitly loaded + if (isLazy && !loadedToolsForServer.includes(mcpTool.name)) { + continue + } + const sanitizedClientName = clientName.replace(/[^a-zA-Z0-9_-]/g, "_") const sanitizedToolName = mcpTool.name.replace(/[^a-zA-Z0-9_-]/g, "_") result[sanitizedClientName + "_" + sanitizedToolName] = await convertMcpTool(mcpTool, client) @@ -572,6 +595,69 @@ export namespace MCP { return result } + export interface McpServerIndex { + toolCount: number + tools: string[] + } + + export async function index(): Promise> { + const s = await state() + const clientsSnapshot = await clients() + const result: Record = {} + + for (const [clientName, client] of Object.entries(clientsSnapshot)) { + if (s.status[clientName]?.status !== "connected") { + continue + } + + const toolsResult = await client.listTools().catch((e) => { + log.error("failed to get tools for index", { clientName, error: e.message }) + return undefined + }) + + if (toolsResult) { + result[clientName] = { + toolCount: toolsResult.tools.length, + tools: toolsResult.tools.map((t) => t.name), + } + } + } + + return result + } + + export async function loadToolsForSession( + serverName: string, + toolNames?: string[], + ): Promise<{ tools: string[]; error?: string }> { + const s = await state() + const client = s.clients[serverName] + + if (!client) { + return { tools: [], error: `Server "${serverName}" not found` } + } + + const serverStatus = s.status[serverName] + if (serverStatus?.status !== "connected") { + return { tools: [], error: `Server "${serverName}" not connected` } + } + + const toolsResult = await client.listTools().catch((e) => { + log.error("loadToolsForSession failed", { serverName, error: e.message }) + return undefined + }) + + if (!toolsResult) { + return { tools: [], error: "Failed to list tools" } + } + + const tools = toolNames + ? toolsResult.tools.filter((t) => toolNames.includes(t.name)).map((t) => t.name) + : toolsResult.tools.map((t) => t.name) + + return { tools } + } + export async function prompts() { const s = await state() const clientsSnapshot = await clients() diff --git a/packages/opencode/src/session/index.ts b/packages/opencode/src/session/index.ts index a204913f77d..1592469b7da 100644 --- a/packages/opencode/src/session/index.ts +++ b/packages/opencode/src/session/index.ts @@ -72,6 +72,7 @@ export namespace Session { diff: z.string().optional(), }) .optional(), + mcpLoadedTools: z.record(z.string(), z.array(z.string())).optional(), }) .meta({ ref: "Session", diff --git a/packages/opencode/src/session/prompt.ts b/packages/opencode/src/session/prompt.ts index 34596e62902..313f28fab2a 100644 --- a/packages/opencode/src/session/prompt.ts +++ b/packages/opencode/src/session/prompt.ts @@ -593,7 +593,11 @@ export namespace SessionPrompt { agent, abort, sessionID, - system: [...(await SystemPrompt.environment()), ...(await SystemPrompt.custom())], + system: [ + ...(await SystemPrompt.environment()), + ...(await SystemPrompt.custom()), + ...(await SystemPrompt.mcpIndex()), + ], messages: [ ...MessageV2.toModelMessage(sessionMessages), ...(isLastStep @@ -723,7 +727,7 @@ export namespace SessionPrompt { }) } - for (const [key, item] of Object.entries(await MCP.tools())) { + for (const [key, item] of Object.entries(await MCP.tools(input.session.mcpLoadedTools))) { const execute = item.execute if (!execute) continue diff --git a/packages/opencode/src/session/system.ts b/packages/opencode/src/session/system.ts index fff90808864..a7f73a94feb 100644 --- a/packages/opencode/src/session/system.ts +++ b/packages/opencode/src/session/system.ts @@ -2,6 +2,7 @@ import { Ripgrep } from "../file/ripgrep" import { Global } from "../global" import { Filesystem } from "../util/filesystem" import { Config } from "../config/config" +import { MCP } from "../mcp" import { Instance } from "../project/instance" import path from "path" @@ -135,4 +136,38 @@ export namespace SystemPrompt { ) return Promise.all([...foundFiles, ...foundUrls]).then((result) => result.filter(Boolean)) } + + export async function mcpIndex(): Promise { + const config = await Config.get() + const isLazy = config.experimental?.mcp_lazy_load ?? false + + if (!isLazy) { + return [] + } + + const index = await MCP.index() + const servers = Object.entries(index) + + if (servers.length === 0) { + return [] + } + + const lines = [ + "## Available MCP Servers", + "", + "Use mcp_load_tools(server_name) to load tools from these servers:", + "", + ] + + for (const [name, info] of servers) { + lines.push(`- **${name}** (${info.toolCount} tools): ${info.tools.join(", ")}`) + } + + lines.push("") + lines.push( + 'To use these tools, first load them with mcp_load_tools("server") or mcp_load_tools("server", ["tool1", "tool2"]).', + ) + + return [lines.join("\n")] + } } diff --git a/packages/opencode/src/tool/mcp-load-tools.ts b/packages/opencode/src/tool/mcp-load-tools.ts new file mode 100644 index 00000000000..a9e4cc093dc --- /dev/null +++ b/packages/opencode/src/tool/mcp-load-tools.ts @@ -0,0 +1,67 @@ +import z from "zod" +import { Tool } from "./tool" +import { MCP } from "../mcp" +import { Session } from "../session" + +export const McpLoadToolsTool = Tool.define("mcp_load_tools", { + description: `Loads tools from an MCP server, making them available for use in this session. +Call this before using any tools from an MCP server. +Returns the list of loaded tool names so you know what's available.`, + parameters: z.object({ + serverName: z.string().describe("Name of the MCP server to load tools from"), + toolNames: z + .array(z.string()) + .optional() + .describe("Specific tools to load. If omitted, loads all tools from the server."), + }), + async execute(args, ctx) { + const { serverName, toolNames } = args + + // Load tools from MCP + const { tools, error } = await MCP.loadToolsForSession(serverName, toolNames) + + if (error) { + return { + title: `Failed to load tools from ${serverName}`, + metadata: {}, + output: `Error: ${error}`, + } + } + + if (tools.length === 0) { + const output = toolNames + ? `No matching tools found. Requested: ${toolNames.join(", ")}` + : `Server "${serverName}" has no tools available.` + return { + title: `No tools loaded from ${serverName}`, + metadata: {}, + output, + } + } + + // Update session state with loaded tools + const session = await Session.get(ctx.sessionID) + const currentLoaded = session.mcpLoadedTools ?? {} + const serverLoaded = new Set(currentLoaded[serverName] ?? []) + + for (const toolName of tools) { + serverLoaded.add(toolName) + } + + await Session.update(ctx.sessionID, (draft) => { + draft.mcpLoadedTools = { + ...currentLoaded, + [serverName]: Array.from(serverLoaded), + } + }) + + // Format tool names with server prefix for clarity + const fullToolNames = tools.map((t) => `${serverName}_${t}`) + + return { + title: `Loaded ${tools.length} tools from ${serverName}`, + metadata: {}, + output: `Loaded ${tools.length} tools from "${serverName}". You can now use: ${fullToolNames.join(", ")}`, + } + }, +}) diff --git a/packages/opencode/src/tool/registry.ts b/packages/opencode/src/tool/registry.ts index 82bf7f56328..182beec3394 100644 --- a/packages/opencode/src/tool/registry.ts +++ b/packages/opencode/src/tool/registry.ts @@ -11,6 +11,7 @@ import { WebFetchTool } from "./webfetch" import { WriteTool } from "./write" import { InvalidTool } from "./invalid" import { SkillTool } from "./skill" +import { McpLoadToolsTool } from "./mcp-load-tools" import type { Agent } from "../agent/agent" import { Tool } from "./tool" import { Instance } from "../project/instance" @@ -109,6 +110,7 @@ export namespace ToolRegistry { SkillTool, ...(Flag.OPENCODE_EXPERIMENTAL_LSP_TOOL ? [LspTool] : []), ...(config.experimental?.batch_tool === true ? [BatchTool] : []), + ...(config.experimental?.mcp_lazy_load === true ? [McpLoadToolsTool] : []), ...custom, ] } diff --git a/packages/opencode/test/mcp/lazy-loading.test.ts b/packages/opencode/test/mcp/lazy-loading.test.ts new file mode 100644 index 00000000000..394e187d5c8 --- /dev/null +++ b/packages/opencode/test/mcp/lazy-loading.test.ts @@ -0,0 +1,355 @@ +import { describe, expect, test, mock, beforeEach } from "bun:test" +import path from "path" +import { tmpdir } from "../fixture/fixture" +import { Instance } from "../../src/project/instance" +import { Log } from "../../src/util/log" + +Log.init({ print: false }) + +// Mock MCP client to avoid actual connections +const mockTools = [ + { name: "tool_one", description: "First tool", inputSchema: { type: "object", properties: {} } }, + { name: "tool_two", description: "Second tool", inputSchema: { type: "object", properties: {} } }, + { name: "tool_three", description: "Third tool", inputSchema: { type: "object", properties: {} } }, +] + +let mockClientStatus = "connected" + +mock.module("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({ + StreamableHTTPClientTransport: class MockStreamableHTTP { + async start() { + if (mockClientStatus === "connected") return + throw new Error("Mock transport cannot connect") + } + }, +})) + +mock.module("@modelcontextprotocol/sdk/client/sse.js", () => ({ + SSEClientTransport: class MockSSE { + async start() { + if (mockClientStatus === "connected") return + throw new Error("Mock transport cannot connect") + } + }, +})) + +mock.module("@modelcontextprotocol/sdk/client/index.js", () => ({ + Client: class MockClient { + async connect() {} + async listTools() { + return { tools: mockTools } + } + async listPrompts() { + return { prompts: [] } + } + async listResources() { + return { resources: [] } + } + async callTool() { + return { content: [] } + } + }, +})) + +beforeEach(() => { + mockClientStatus = "connected" +}) + +// Import modules after mocking +const { Config } = await import("../../src/config/config") +const { Session } = await import("../../src/session") +const { ToolRegistry } = await import("../../src/tool/registry") +const { SystemPrompt } = await import("../../src/session/system") +const { MCP } = await import("../../src/mcp/index") + +describe("MCP lazy loading", () => { + describe("config", () => { + test("mcp_lazy_load defaults to undefined (off)", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const config = await Config.get() + expect(config.experimental?.mcp_lazy_load).toBeUndefined() + }, + }) + }) + + test("mcp_lazy_load can be enabled", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + experimental: { + mcp_lazy_load: true, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const config = await Config.get() + expect(config.experimental?.mcp_lazy_load).toBe(true) + }, + }) + }) + }) + + describe("session.mcpLoadedTools", () => { + test("session can store mcpLoadedTools", async () => { + await using tmp = await tmpdir({ git: true }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const session = await Session.create({}) + + // Update session with loaded tools + await Session.update(session.id, (draft) => { + draft.mcpLoadedTools = { + "test-server": ["tool_one", "tool_two"], + } + }) + + const updated = await Session.get(session.id) + expect(updated.mcpLoadedTools).toEqual({ + "test-server": ["tool_one", "tool_two"], + }) + + await Session.remove(session.id) + }, + }) + }) + + test("mcpLoadedTools can accumulate tools", async () => { + await using tmp = await tmpdir({ git: true }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const session = await Session.create({}) + + // First update + await Session.update(session.id, (draft) => { + draft.mcpLoadedTools = { + "server-a": ["tool1"], + } + }) + + // Second update - add more tools + const current = await Session.get(session.id) + await Session.update(session.id, (draft) => { + draft.mcpLoadedTools = { + ...current.mcpLoadedTools, + "server-a": [...(current.mcpLoadedTools?.["server-a"] ?? []), "tool2"], + "server-b": ["toolX"], + } + }) + + const updated = await Session.get(session.id) + expect(updated.mcpLoadedTools).toEqual({ + "server-a": ["tool1", "tool2"], + "server-b": ["toolX"], + }) + + await Session.remove(session.id) + }, + }) + }) + }) + + describe("tool registry", () => { + test("mcp_load_tools is registered when lazy loading is enabled", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + experimental: { + mcp_lazy_load: true, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const ids = await ToolRegistry.ids() + expect(ids).toContain("mcp_load_tools") + }, + }) + }) + + test("mcp_load_tools is NOT registered when lazy loading is disabled", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + experimental: { + mcp_lazy_load: false, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const ids = await ToolRegistry.ids() + expect(ids).not.toContain("mcp_load_tools") + }, + }) + }) + }) + + describe("SystemPrompt.mcpIndex", () => { + test("returns empty array when lazy loading is disabled", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + experimental: { + mcp_lazy_load: false, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const result = await SystemPrompt.mcpIndex() + expect(result).toEqual([]) + }, + }) + }) + + test("returns empty array when no MCP servers configured", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + experimental: { + mcp_lazy_load: true, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const result = await SystemPrompt.mcpIndex() + expect(result).toEqual([]) + }, + }) + }) + }) + + describe("MCP.loadToolsForSession", () => { + test("returns error for non-existent server", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + experimental: { + mcp_lazy_load: true, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const result = await MCP.loadToolsForSession("non-existent-server") + expect(result.error).toBeDefined() + expect(result.error).toContain("not found") + expect(result.tools).toEqual([]) + }, + }) + }) + }) + + describe("MCP.index", () => { + test("returns empty object when no servers connected", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + experimental: { + mcp_lazy_load: true, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const index = await MCP.index() + expect(index).toEqual({}) + }, + }) + }) + }) + + describe("MCP.tools", () => { + test("returns empty object when no servers connected", async () => { + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "opencode.json"), + JSON.stringify({ + $schema: "https://opencode.ai/config.json", + experimental: { + mcp_lazy_load: true, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const tools = await MCP.tools({}) + expect(tools).toEqual({}) + }, + }) + }) + }) +})