diff --git a/src/mcp/register.ts b/src/mcp/register.ts index 9d87967..5fd219c 100644 --- a/src/mcp/register.ts +++ b/src/mcp/register.ts @@ -1,17 +1,33 @@ import type * as VSCode from "vscode"; import { resolvePatchloomStatus } from "../binary/patchloom.js"; -import { getPatchloomLog } from "../logging/outputChannel.js"; +import { getPatchloomLog, getPatchloomRuntimeConfig } from "../logging/outputChannel.js"; +import { isAllowedPatchloomEnvKey } from "../util.js"; /** Plain descriptor used to construct vscode.McpStdioServerDefinition at register time. */ export interface McpServerBinaryDescriptor { readonly label: string; readonly command: string; readonly args: readonly string[]; + readonly env?: Record; +} + +function patchloomOnlyEnv(extra: Record | undefined): Record { + const env: Record = {}; + if (extra === undefined) { + return env; + } + for (const [key, value] of Object.entries(extra)) { + if (isAllowedPatchloomEnvKey(key)) { + env[key] = value; + } + } + return env; } /** Pure helper for native MCP definitions (no vscode). Empty when binary unknown. */ export function mcpServerDefinitionsForBinary( - binaryPath: string | undefined + binaryPath: string | undefined, + env?: Record ): readonly McpServerBinaryDescriptor[] { if (!binaryPath) { return []; @@ -20,12 +36,18 @@ export function mcpServerDefinitionsForBinary( { label: "Patchloom MCP", command: binaryPath, - args: ["mcp-server"] + args: ["mcp-server"], + env: patchloomOnlyEnv(env) } ]; } -type McpStdioServerDefinitionCtor = new (...args: readonly unknown[]) => unknown; +type McpStdioServerDefinitionCtor = new ( + label: string, + command: string, + args: string[], + env: Record +) => unknown; interface VsCodeLmWithMcp { registerMcpServerDefinitionProvider?( @@ -52,23 +74,16 @@ function mcpStdioCtor(vscode: VsCodeWithMcpApi): McpStdioServerDefinitionCtor | return typeof ctor === "function" ? ctor : undefined; } -function createMcpStdioServerDefinition( +export function createMcpStdioServerDefinition( Ctor: McpStdioServerDefinitionCtor, descriptor: McpServerBinaryDescriptor ): unknown | undefined { const args = [...descriptor.args]; + const env = descriptor.env ?? {}; try { - return new Ctor({ - label: descriptor.label, - command: descriptor.command, - args - }); + return new Ctor(descriptor.label, descriptor.command, args, env); } catch { - try { - return new Ctor(descriptor.label, descriptor.command, args); - } catch { - return undefined; - } + return undefined; } } @@ -102,8 +117,9 @@ export async function registerMcpServerProviderWithBinary(context: VSCode.Extens const disposable = vscode.lm.registerMcpServerDefinitionProvider("patchloom", { onDidChangeMcpServerDefinitions: emitter.event, provideMcpServerDefinitions: async () => { + const runtime = await getPatchloomRuntimeConfig(); const definitions: unknown[] = []; - for (const descriptor of mcpServerDefinitionsForBinary(resolvedBinaryPath)) { + for (const descriptor of mcpServerDefinitionsForBinary(resolvedBinaryPath, runtime.extraEnv)) { const definition = createMcpStdioServerDefinition(Ctor, descriptor); if (definition !== undefined) { definitions.push(definition); diff --git a/test/unit/mcpRegister.test.ts b/test/unit/mcpRegister.test.ts index 44e3242..e274013 100644 --- a/test/unit/mcpRegister.test.ts +++ b/test/unit/mcpRegister.test.ts @@ -1,6 +1,18 @@ import assert from "node:assert/strict"; import test from "node:test"; -import { mcpServerDefinitionsForBinary } from "../../src/mcp/register.js"; +import { + createMcpStdioServerDefinition, + mcpServerDefinitionsForBinary +} from "../../src/mcp/register.js"; + +class FakeStdio { + constructor( + public label: string, + public command: string, + public args: string[], + public env: Record = {} + ) {} +} test("mcpServerDefinitionsForBinary returns empty when binary is undefined", () => { assert.deepEqual(mcpServerDefinitionsForBinary(undefined), []); @@ -12,6 +24,37 @@ test("mcpServerDefinitionsForBinary returns one stdio definition for a path", () assert.deepEqual(defs[0], { label: "Patchloom MCP", command: "/opt/patchloom", - args: ["mcp-server"] + args: ["mcp-server"], + env: {} }); }); + +test("mcpServerDefinitionsForBinary includes filtered PATCHLOOM env", () => { + const defs = mcpServerDefinitionsForBinary("/opt/patchloom", { + PATCHLOOM_MCP_SURFACE: "core", + PATH: "/tmp/evil" + }); + assert.deepEqual(defs[0]?.env, { PATCHLOOM_MCP_SURFACE: "core" }); +}); + +test("createMcpStdioServerDefinition constructs positionally so command is the binary", () => { + const instance = createMcpStdioServerDefinition(FakeStdio, { + label: "Patchloom MCP", + command: "/opt/patchloom", + args: ["mcp-server"] + }) as FakeStdio; + assert.equal(instance.command, "/opt/patchloom"); + assert.equal(instance.label, "Patchloom MCP"); + assert.deepEqual(instance.args, ["mcp-server"]); + assert.deepEqual(instance.env, {}); +}); + +test("createMcpStdioServerDefinition forwards PATCHLOOM env", () => { + const instance = createMcpStdioServerDefinition(FakeStdio, { + label: "Patchloom MCP", + command: "/opt/patchloom", + args: ["mcp-server"], + env: { PATCHLOOM_MCP_SURFACE: "core" } + }) as FakeStdio; + assert.equal(instance.env.PATCHLOOM_MCP_SURFACE, "core"); +});