gsd-pi 2.76.0 → 2.77.0-dev.2daa994b6

package/src/resources/extensions/mcp-client/index.ts

@@ -11,7 +11,7 @@
  *   mcp_call      — Call a tool on an MCP server (lazy connect)
  */
 
-import type { ExtensionAPI } from "@gsd/pi-coding-agent";
+import type { ExtensionAPI, ExtensionContext } from "@gsd/pi-coding-agent";
 import {
 	truncateHead,
 	DEFAULT_MAX_BYTES,
@@ -33,6 +33,7 @@ import type { McpHttpAuthConfig } from "./auth.js";
 interface McpServerConfig {
 	name: string;
 	transport: "stdio" | "http" | "unknown";
+	sourcePath: string;
 	command?: string;
 	args?: string[];
 	env?: Record<string, string>;
@@ -58,8 +59,40 @@ interface ManagedConnection {
 // ─── Connection Manager ───────────────────────────────────────────────────────
 
 const connections = new Map<string, ManagedConnection>();
+const pendingConnections = new Map<string, Promise<Client>>();
 let configCache: McpServerConfig[] | null = null;
 const toolCache = new Map<string, McpToolSchema[]>();
+const trustedStdioServers = new Set<string>();
+
+const CHILD_ENV_ALLOWLIST = new Set([
+	"PATH",
+	"Path",
+	"HOME",
+	"USER",
+	"USERNAME",
+	"USERPROFILE",
+	"SHELL",
+	"TMPDIR",
+	"TEMP",
+	"TMP",
+	"SystemRoot",
+	"WINDIR",
+	"APPDATA",
+	"LOCALAPPDATA",
+	"XDG_CONFIG_HOME",
+	"XDG_CACHE_HOME",
+]);
+
+function stdioTrustKey(config: McpServerConfig): string {
+	return JSON.stringify({
+		name: config.name,
+		sourcePath: config.sourcePath,
+		command: config.command,
+		args: config.args ?? [],
+		cwd: config.cwd,
+		env: config.env ?? {},
+	});
+}
 
 function readConfigs(): McpServerConfig[] {
 	if (configCache) return configCache;
@@ -99,6 +132,7 @@ function readConfigs(): McpServerConfig[] {
 				servers.push({
 					name,
 					transport,
+					sourcePath: configPath,
 					...(hasCommand && {
 						command: config.command as string,
 						args: Array.isArray(config.args) ? (config.args as string[]) : undefined,
@@ -121,7 +155,57 @@ function readConfigs(): McpServerConfig[] {
 	return servers;
 }
 
-function getServerConfig(name: string): McpServerConfig | undefined {
+export function _buildMcpChildEnvForTest(configEnv: Record<string, string> | undefined): Record<string, string> {
+	const childEnv: Record<string, string> = {};
+	for (const key of CHILD_ENV_ALLOWLIST) {
+		const value = process.env[key];
+		if (typeof value === "string") childEnv[key] = value;
+	}
+	return {
+		...childEnv,
+		...(configEnv ? resolveEnv(configEnv) : {}),
+	};
+}
+
+export function _buildMcpTrustConfirmOptionsForTest(signal?: AbortSignal): { timeout: number; signal?: AbortSignal } {
+	return signal ? { timeout: 120_000, signal } : { timeout: 120_000 };
+}
+
+async function assertTrustedStdioServer(
+	config: McpServerConfig,
+	ctx?: ExtensionContext,
+	signal?: AbortSignal,
+): Promise<string | undefined> {
+	if (config.transport !== "stdio") return undefined;
+	const trustKey = stdioTrustKey(config);
+	if (trustedStdioServers.has(trustKey)) return undefined;
+
+	if (!ctx?.hasUI) {
+		throw new Error(
+			`MCP server "${config.name}" is a project-local stdio command from ${config.sourcePath}. ` +
+			"Run this from an interactive GSD session and approve the server before use.",
+		);
+	}
+
+	const commandLine = [config.command, ...(config.args ?? [])].filter(Boolean).join(" ");
+	const envKeys = Object.keys(config.env ?? {});
+	const envSummary = envKeys.length > 0
+		? `\n\nConfigured environment keys: ${envKeys.join(", ")}`
+		: "\n\nNo explicit environment keys configured.";
+	const approved = await ctx.ui.confirm(
+		`Trust MCP server "${config.name}"?`,
+		`Project config ${config.sourcePath} wants to start:\n\n${commandLine}${envSummary}\n\nOnly approve MCP servers you trust.`,
+		_buildMcpTrustConfirmOptionsForTest(signal),
+	);
+	if (!approved) {
+		throw new Error(`MCP server "${config.name}" was not approved by the user.`);
+	}
+	return trustKey;
+}
+
+// Exported for tests (see tests/server-name-spaces.test.ts).
+// Production call sites treat this as module-private.
+export function getServerConfig(name: string): McpServerConfig | undefined {
 	const trimmed = name.trim();
 	return readConfigs().find((s) =>
 		s.name === trimmed ||
@@ -145,7 +229,7 @@ function resolveEnv(env: Record<string, string>): Record<string, string> {
 	return resolved;
 }
 
-async function getOrConnect(name: string, signal?: AbortSignal): Promise<Client> {
+async function getOrConnect(name: string, signal?: AbortSignal, ctx?: ExtensionContext): Promise<Client> {
 	const config = getServerConfig(name);
 	if (!config) throw new Error(`Unknown MCP server: "${name}". Use mcp_servers to list available servers.`);
 
@@ -154,14 +238,29 @@ async function getOrConnect(name: string, signal?: AbortSignal): Promise<Client>
 	const existing = connections.get(config.name);
 	if (existing) return existing.client;
 
+	const pending = pendingConnections.get(config.name);
+	if (pending) return pending;
+
+	const connectionPromise = connectServer(config, signal, ctx);
+	pendingConnections.set(config.name, connectionPromise);
+	try {
+		return await connectionPromise;
+	} finally {
+		pendingConnections.delete(config.name);
+	}
+}
+
+async function connectServer(config: McpServerConfig, signal?: AbortSignal, ctx?: ExtensionContext): Promise<Client> {
 	const client = new Client({ name: "gsd", version: "1.0.0" });
 	let transport: StdioClientTransport | StreamableHTTPClientTransport;
+	let approvedTrustKey: string | undefined;
 
 	if (config.transport === "stdio" && config.command) {
+		approvedTrustKey = await assertTrustedStdioServer(config, ctx, signal);
 		transport = new StdioClientTransport({
 			command: config.command,
 			args: config.args,
-			env: config.env ? { ...process.env, ...resolveEnv(config.env) } as Record<string, string> : undefined,
+			env: _buildMcpChildEnvForTest(config.env),
 			cwd: config.cwd,
 			stderr: "pipe",
 		});
@@ -179,9 +278,24 @@ async function getOrConnect(name: string, signal?: AbortSignal): Promise<Client>
 		throw new Error(`Server "${config.name}" has unsupported transport: ${config.transport}`);
 	}
 
-	await client.connect(transport, { signal, timeout: 30000 });
-	connections.set(config.name, { client, transport });
-	return client;
+	try {
+		await client.connect(transport, { signal, timeout: 30000 });
+		if (approvedTrustKey) trustedStdioServers.add(approvedTrustKey);
+		connections.set(config.name, { client, transport });
+		return client;
+	} catch (err) {
+		try {
+			await transport.close();
+		} catch {
+			// Best-effort cleanup after a failed or aborted connection attempt.
+		}
+		try {
+			await client.close();
+		} catch {
+			// Best-effort cleanup after a failed or aborted connection attempt.
+		}
+		throw err;
+	}
 }
 
 async function closeAll(): Promise<void> {
@@ -191,9 +305,16 @@ async function closeAll(): Promise<void> {
 		} catch {
 			// Best-effort cleanup
 		}
+		try {
+			await conn.transport.close();
+		} catch {
+			// Best-effort cleanup
+		}
 		connections.delete(name);
 	});
 	await Promise.allSettled(closing);
+	pendingConnections.clear();
+	trustedStdioServers.clear();
 	toolCache.clear();
 }
 
@@ -331,7 +452,7 @@ export default function (pi: ExtensionAPI) {
 			}),
 		}),
 
-		async execute(_id, params, signal) {
+		async execute(_id, params, signal, _onUpdate, ctx) {
 			try {
 				// Return cached tools if available
 				const cached = toolCache.get(params.server);
@@ -348,7 +469,7 @@ export default function (pi: ExtensionAPI) {
 					};
 				}
 
-				const client = await getOrConnect(params.server, signal);
+				const client = await getOrConnect(params.server, signal, ctx);
 				const result = await client.listTools(undefined, { signal, timeout: 30000 });
 				const tools: McpToolSchema[] = (result.tools ?? []).map((t) => ({
 					name: t.name,
@@ -423,9 +544,9 @@ export default function (pi: ExtensionAPI) {
 			),
 		}),
 
-		async execute(_id, params, signal) {
+		async execute(_id, params, signal, _onUpdate, ctx) {
 			try {
-				const client = await getOrConnect(params.server, signal);
+				const client = await getOrConnect(params.server, signal, ctx);
 				const result = await client.callTool(
 					{ name: params.tool, arguments: params.args ?? {} },
 					undefined,

package/src/resources/extensions/mcp-client/tests/server-name-spaces.test.ts

@@ -1,55 +1,89 @@
 /**
  * Regression test for #3029 — mcp_discover fails for server names with spaces.
  *
- * The getServerConfig lookup must handle:
- *   1. Exact match (already works)
+ * getServerConfig must handle:
+ *   1. Exact match
  *   2. Names with leading/trailing whitespace (trimming)
  *   3. Case-insensitive matching (e.g. "Langgraph code" vs "langgraph Code")
  *
- * We test at the source level since getServerConfig is not exported.
+ * getOrConnect must use the canonical (config.name) as the cache key so that
+ * subsequent lookups with variant casing/whitespace hit the same connection.
+ *
+ * These are behaviour tests against the real exported getServerConfig — no
+ * source grep.
  */
 
-import test from "node:test";
+import test, { before, after } from "node:test";
 import assert from "node:assert/strict";
-import { readFileSync } from "node:fs";
-import { dirname, join } from "node:path";
-import { fileURLToPath } from "node:url";
+import { writeFileSync, unlinkSync, existsSync, mkdirSync, rmSync } from "node:fs";
+import { join } from "node:path";
+import { mkdtempSync } from "node:fs";
+import { tmpdir } from "node:os";
+
+import { getServerConfig } from "../index.js";
+
+// readConfigs() anchors to process.cwd() — run each test in a sandbox dir
+// with a purpose-built .mcp.json so the extension reads our fixture, not
+// whatever .mcp.json happens to live in the current working directory.
+let sandboxDir: string;
+let originalCwd: string;
+
+before(() => {
+	originalCwd = process.cwd();
+	sandboxDir = mkdtempSync(join(tmpdir(), "mcp-name-spaces-"));
+	const mcpConfig = {
+		mcpServers: {
+			"Langgraph Code": {
+				command: "echo",
+				args: ["test"],
+			},
+			"other-server": {
+				url: "https://example.com",
+			},
+		},
+	};
+	writeFileSync(join(sandboxDir, ".mcp.json"), JSON.stringify(mcpConfig), "utf-8");
+	process.chdir(sandboxDir);
+});
 
-const __filename = fileURLToPath(import.meta.url);
-const __dirname = dirname(__filename);
+after(() => {
+	process.chdir(originalCwd);
+	try {
+		rmSync(sandboxDir, { recursive: true, force: true });
+	} catch {
+		// Best-effort cleanup
+	}
+});
 
-const source = readFileSync(join(__dirname, "..", "index.ts"), "utf-8");
+test("#3029: getServerConfig finds exact match", () => {
+	const cfg = getServerConfig("Langgraph Code");
+	assert.ok(cfg, "exact name must resolve");
+	assert.equal(cfg?.name, "Langgraph Code");
+});
 
 test("#3029: getServerConfig trims whitespace from input name", () => {
-	assert.ok(
-		source.includes(".trim()"),
-		"getServerConfig should trim the input name before comparison",
-	);
+	const cfg = getServerConfig("   Langgraph Code  ");
+	assert.ok(cfg, "whitespace-padded name must resolve to the same server");
+	assert.equal(cfg?.name, "Langgraph Code");
 });
 
 test("#3029: getServerConfig performs case-insensitive matching", () => {
-	assert.ok(
-		source.includes(".toLowerCase()"),
-		"getServerConfig should compare names case-insensitively",
-	);
+	const cfg = getServerConfig("langgraph code");
+	assert.ok(cfg, "lower-cased name must resolve");
+	assert.equal(cfg?.name, "Langgraph Code");
+
+	const mixed = getServerConfig("LANGGRAPH CODE");
+	assert.ok(mixed, "upper-cased name must resolve");
+	assert.equal(mixed?.name, "Langgraph Code");
+});
+
+test("#3029: getServerConfig combines trim + case-insensitive", () => {
+	const cfg = getServerConfig("  LANGGRAPH code  ");
+	assert.ok(cfg, "padded + mixed-case must resolve");
+	assert.equal(cfg?.name, "Langgraph Code");
 });
 
-test("#3029: getOrConnect normalizes name for connection cache lookup", () => {
-	// The connections Map key must use the canonical (config) name, not the
-	// raw user input, so that subsequent lookups hit the cache even when the
-	// user's casing differs.
-	const getOrConnectMatch = source.match(
-		/async function getOrConnect\(name: string[\s\S]*?const existing = connections\.get\(/,
-	);
-	assert.ok(
-		getOrConnectMatch,
-		"getOrConnect function should exist",
-	);
-	// After the fix, getOrConnect should normalize the name via getServerConfig
-	// or use config.name as the canonical cache key.
-	assert.ok(
-		source.includes("connections.get(config.name") ||
-		source.includes("connections.set(config.name"),
-		"getOrConnect should use config.name (canonical) as the connections cache key",
-	);
+test("#3029: getServerConfig returns undefined for unknown name", () => {
+	const cfg = getServerConfig("does-not-exist");
+	assert.equal(cfg, undefined);
 });

package/src/resources/extensions/ollama/index.ts

@@ -45,8 +45,12 @@ let providerRegistered = false;
 /**
  * Probe Ollama and register discovered models.
  * Safe to call multiple times — re-discovers and re-registers.
+ *
+ * Exported for tests (see ollama-auth-mode.test.ts, ollama-status-indicator.test.ts)
+ * so a fake HTTP endpoint can drive the registration/unregistration paths.
+ * Production callers always go through the session_start handler below.
  */
-async function probeAndRegister(pi: ExtensionAPI): Promise<boolean> {
+export async function probeAndRegister(pi: ExtensionAPI): Promise<boolean> {
 	const running = await client.isRunning();
 	if (!running) {
 		if (providerRegistered) {

package/src/resources/extensions/ollama/ollama-auth-mode.test.ts

@@ -1,20 +1,128 @@
 /**
  * Regression test for #3440: Ollama extension must register with
- * authMode "apiKey" (not "none") to avoid streamSimple requirement.
+ * authMode "apiKey" (not "none"), otherwise the core bails out because
+ * the provider has no streamSimple.
+ *
+ * Behaviour test: spin up a fake Ollama endpoint, point OLLAMA_HOST at
+ * it, and invoke probeAndRegister with a mock pi. Assert that
+ * registerProvider was called with authMode "apiKey".
  */
-import { test } from "node:test";
+import { test, before, after, beforeEach } from "node:test";
 import assert from "node:assert/strict";
-import { readFileSync } from "node:fs";
-import { join, dirname } from "node:path";
-import { fileURLToPath } from "node:url";
-
-const __dirname = dirname(fileURLToPath(import.meta.url));
-
-test("Ollama registers with authMode apiKey, not none (#3440)", () => {
-  const src = readFileSync(join(__dirname, "index.ts"), "utf-8");
-  // Find the registerProvider call
-  const registerBlock = src.slice(src.indexOf("pi.registerProvider(\"ollama\""));
-  const authLine = registerBlock.match(/authMode:\s*"(\w+)"/);
-  assert.ok(authLine, "registerProvider must specify authMode");
-  assert.equal(authLine![1], "apiKey", "authMode must be apiKey, not none");
+import { createServer, type Server } from "node:http";
+import type { AddressInfo } from "node:net";
+
+import { probeAndRegister } from "./index.js";
+
+type RegisterCall = [string, Record<string, unknown>];
+
+function makeMockPi() {
+	const calls: { registerProvider: RegisterCall[]; unregisterProvider: string[] } = {
+		registerProvider: [],
+		unregisterProvider: [],
+	};
+	const pi = {
+		registerProvider(id: string, spec: Record<string, unknown>) {
+			calls.registerProvider.push([id, spec]);
+		},
+		unregisterProvider(id: string) {
+			calls.unregisterProvider.push(id);
+		},
+	} as unknown as Parameters<typeof probeAndRegister>[0];
+	return { pi, calls };
+}
+
+let server: Server;
+let savedHost: string | undefined;
+
+before(async () => {
+	// Fake Ollama endpoint that:
+	//   GET /          → 200 (isRunning probe)
+	//   GET /api/tags  → one model
+	//   POST /api/show → minimal capability info
+	server = createServer((req, res) => {
+		if (req.method === "GET" && req.url === "/") {
+			res.writeHead(200, { "Content-Type": "text/plain" });
+			res.end("Ollama is running");
+			return;
+		}
+		if (req.method === "GET" && req.url === "/api/tags") {
+			res.writeHead(200, { "Content-Type": "application/json" });
+			res.end(
+				JSON.stringify({
+					models: [
+						{
+							name: "llama3:latest",
+							modified_at: new Date().toISOString(),
+							size: 1_000_000,
+							digest: "abc",
+							details: {
+								parent_model: "",
+								format: "gguf",
+								family: "llama",
+								families: ["llama"],
+								parameter_size: "8B",
+								quantization_level: "Q4_0",
+							},
+						},
+					],
+				}),
+			);
+			return;
+		}
+		if (req.method === "POST" && req.url === "/api/show") {
+			res.writeHead(200, { "Content-Type": "application/json" });
+			res.end(
+				JSON.stringify({
+					modelfile: "",
+					parameters: "",
+					template: "",
+					details: {
+						parent_model: "",
+						format: "gguf",
+						family: "llama",
+						families: ["llama"],
+						parameter_size: "8B",
+						quantization_level: "Q4_0",
+					},
+					model_info: { "llama.context_length": 8192 },
+					capabilities: [],
+				}),
+			);
+			return;
+		}
+		res.writeHead(404);
+		res.end();
+	});
+	await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
+	const { port } = server.address() as AddressInfo;
+	savedHost = process.env.OLLAMA_HOST;
+	process.env.OLLAMA_HOST = `http://127.0.0.1:${port}`;
+});
+
+after(async () => {
+	await new Promise<void>((resolve) => server.close(() => resolve()));
+	if (savedHost === undefined) delete process.env.OLLAMA_HOST;
+	else process.env.OLLAMA_HOST = savedHost;
+});
+
+beforeEach(() => {
+	// Each test starts from a clean providerRegistered state in the module.
+	// probeAndRegister is idempotent — calling it is safe regardless.
+});
+
+test("Ollama registers with authMode apiKey, not none (#3440)", async () => {
+	const { pi, calls } = makeMockPi();
+
+	const found = await probeAndRegister(pi);
+	assert.equal(found, true, "probeAndRegister should return true when models are discovered");
+
+	assert.equal(calls.registerProvider.length, 1, "registerProvider should be called exactly once");
+	const [providerId, spec] = calls.registerProvider[0];
+	assert.equal(providerId, "ollama");
+	assert.equal(
+		spec.authMode,
+		"apiKey",
+		"authMode must be apiKey so the core doesn't require streamSimple for every model",
+	);
 });