one-code-extension 0.4.1 → 0.6.0

package/extensions/mcp/announce.ts

@@ -1,22 +1,20 @@
 /**
- * What to tell the model when MCP state changes after the first request (pure).
+ * What to tell the model when MCP state changes inside an epoch (pure).
  *
  * The MCP instructions and failed-servers reminders are `first-prepend` blocks
- * on message 1. Before the first request they are simply (re)written; after it
- * message 1 is part of every request's cached prefix, so they are FROZEN and
- * changes ride one-shot notices where the model reads next (same rule as the
- * deferred-tools listing, tool-search/announce.ts). `McpAnnounced` records what
- * the model has been told; `mcpDelta` compares it with the live state.
+ * on message 1. A request that opens an epoch renders them from live state;
+ * inside an epoch message 1 is the record (extensions/ledger), so changes ride
+ * notices on the request's tail, decided when the request goes out from what
+ * the epoch has told (the same rule as the deferred-tools listing). What was
+ * told is a whole-text block fact per server (`mcp-instructed:<name>`,
+ * `mcp-failed:<name>`): the hash of the instructions or error last told, or of
+ * the empty text once a drop or a recovery was announced.
  */
 
 import { mcpFailuresReminder, mcpInstructionsReminder } from "./client.ts";
-
-export interface McpAnnounced {
-	/** Servers whose instructions the model has been given, with the text it saw. */
-	instructed: Map<string, string>;
-	/** Servers the model has been told failed, with the error it saw. */
-	failed: Map<string, string>;
-}
+import type { BlockFact, ToldFact } from "../lib/ledger/records.ts";
+import { toldBlockFact } from "../lib/reminders.ts";
+import { countNoun } from "../lib/tui-render.ts";
 
 export interface McpSnapshot {
 	connected: Array<{ name: string; instructions?: string }>;
@@ -32,47 +30,83 @@ export interface McpDelta {
 	recovered: string[];
 	/** Servers whose instructions the model saw that are no longer connected and not reported failed. */
 	dropped: string[];
+	/** What announcing this delta tells, for the request that carries it. */
+	facts: ToldFact[];
 }
 
-export function emptyAnnounced(): McpAnnounced {
-	return { instructed: new Map(), failed: new Map() };
+const INSTRUCTED = "mcp-instructed:";
+const FAILED = "mcp-failed:";
+/** The told hash of "nothing": a drop or a recovery was announced. */
+const NONE = toldBlockFact("", "").hash;
+
+/**
+ * Store and compare the exact bounded values the reminder puts on the wire.
+ * A changing tail beyond either cap is not a model-visible capability change
+ * and must not become a false notice.
+ */
+function displayedInstructions(name: string, instructions: string): string {
+	const rendered = mcpInstructionsReminder([{ server: { name }, instructions }]);
+	const prefix = "# MCP Server Instructions\n\nThe following MCP servers have provided instructions for how to use their tools and resources:\n\n" + `## ${name}\n`;
+	return rendered?.startsWith(prefix) ? rendered.slice(prefix.length) : instructions;
+}
+function displayedFailure(error: string): string {
+	return error.split("\n")[0].slice(0, 300);
 }
 
-/** What message 1 says after it was (re)written from `snapshot`, before the first request. */
-export function announcedFrom(snapshot: McpSnapshot): McpAnnounced {
-	return {
-		instructed: new Map(
-			snapshot.connected.filter((c) => c.instructions).map((c) => [c.name, c.instructions as string]),
-		),
-		failed: new Map(snapshot.failed.map((f) => [f.name, f.error])),
-	};
+const instructedFact = (name: string, instructions: string): BlockFact => toldBlockFact(INSTRUCTED + name, displayedInstructions(name, instructions));
+const failedFact = (name: string, error: string): BlockFact => toldBlockFact(FAILED + name, displayedFailure(error));
+
+/** What message 1's two MCP blocks rendered from `snapshot` tell. */
+export function mcpFacts(snapshot: McpSnapshot): ToldFact[] {
+	return [
+		...snapshot.connected.flatMap((c) => (c.instructions ? [instructedFact(c.name, c.instructions)] : [])),
+		...snapshot.failed.map((f) => failedFact(f.name, f.error)),
+	];
 }
 
-export function mcpDelta(announced: McpAnnounced, snapshot: McpSnapshot): McpDelta {
+/**
+ * What changed since the epoch told `told` (whole-text block hashes by key),
+ * and the facts announcing it tells. While the startup connections are still
+ * `settling`, a server not connected yet is not a drop, and neither is one in
+ * `connecting` (a reconnect in flight).
+ */
+export function mcpDelta(told: ReadonlyMap<string, string>, snapshot: McpSnapshot, settling = false, connecting: ReadonlySet<string> = new Set()): McpDelta {
 	const connectedNames = new Set(snapshot.connected.map((c) => c.name));
-	const failedNames = new Set(snapshot.failed.map((f) => f.name));
+	// One failure per server (the latest): two kept errors would alternate as the told fact.
+	const failed = [...new Map(snapshot.failed.map((f) => [f.name, f])).values()];
+	const failedNames = new Set(failed.map((f) => f.name));
+	const toldNames = (prefix: string) => [...told].filter(([key, hash]) => key.startsWith(prefix) && hash !== NONE).map(([key]) => key.slice(prefix.length));
+	const newInstructions = snapshot.connected
+		.filter((c) => c.instructions && told.get(INSTRUCTED + c.name) !== instructedFact(c.name, c.instructions).hash)
+		.map((c) => ({ name: c.name, instructions: displayedInstructions(c.name, c.instructions as string) }));
+	const newFailures = failed
+		.filter((f) => told.get(FAILED + f.name) !== failedFact(f.name, f.error).hash)
+		.map((f) => ({ name: f.name, error: displayedFailure(f.error) }));
+	const recovered = toldNames(FAILED).filter((name) => connectedNames.has(name));
+	const dropped = settling ? [] : toldNames(INSTRUCTED).filter((name) => !connectedNames.has(name) && !failedNames.has(name) && !connecting.has(name));
 	return {
-		newInstructions: snapshot.connected
-			.filter((c) => c.instructions && announced.instructed.get(c.name) !== c.instructions)
-			.map((c) => ({ name: c.name, instructions: c.instructions as string })),
-		newFailures: snapshot.failed.filter((f) => announced.failed.get(f.name) !== f.error),
-		recovered: [...announced.failed.keys()].filter((name) => connectedNames.has(name)),
-		dropped: [...announced.instructed.keys()].filter((name) => !connectedNames.has(name) && !failedNames.has(name)),
+		newInstructions,
+		newFailures,
+		recovered,
+		dropped,
+		facts: [
+			...newInstructions.map((s) => toldBlockFact(INSTRUCTED + s.name, s.instructions)),
+			...newFailures.map((f) => toldBlockFact(FAILED + f.name, f.error)),
+			...recovered.map((name) => toldBlockFact(FAILED + name, "")),
+			...dropped.map((name) => toldBlockFact(INSTRUCTED + name, "")),
+		],
 	};
 }
 
-export function applyMcpDelta(announced: McpAnnounced, delta: McpDelta): void {
-	for (const s of delta.newInstructions) announced.instructed.set(s.name, s.instructions);
-	for (const f of delta.newFailures) announced.failed.set(f.name, f.error);
-	for (const name of delta.recovered) announced.failed.delete(name);
-	for (const name of delta.dropped) announced.instructed.delete(name);
-}
-
-/** One-shot notice texts for a delta, in the order they should be queued. Empty when nothing changed. */
+/**
+ * The notices for a delta, in the order they go on the tail. Empty when
+ * nothing changed.
+ */
 export function mcpChangeNotices(delta: McpDelta): string[] {
 	const notices: string[] = [];
+	const add = (text: string) => notices.push(text);
 	if (delta.recovered.length > 0) {
-		notices.push(
+		add(
 			`MCP server${delta.recovered.length === 1 ? "" : "s"} ${delta.recovered.join(", ")} ` +
 				"connected after this conversation started (an earlier notice reported a connection failure). " +
 				"The server's tools are available as deferred tools — load them with tool_search.",
@@ -80,20 +114,33 @@ export function mcpChangeNotices(delta: McpDelta): string[] {
 	}
 	if (delta.newInstructions.length > 0) {
 		const names = delta.newInstructions.map((s) => s.name);
-		notices.push(
+		add(
 			`MCP server${names.length === 1 ? "" : "s"} ${names.join(", ")} connected after this conversation started; ` +
 				"the tools are available as deferred tools (load with tool_search). Server instructions:\n\n" +
 				mcpInstructionsReminder(delta.newInstructions.map((s) => ({ server: { name: s.name }, instructions: s.instructions }))),
 		);
 	}
 	if (delta.newFailures.length > 0) {
-		notices.push(mcpFailuresReminder(delta.newFailures));
+		add(mcpFailuresReminder(delta.newFailures));
 	}
 	if (delta.dropped.length > 0) {
-		notices.push(
+		add(
 			`MCP server${delta.dropped.length === 1 ? "" : "s"} ${delta.dropped.join(", ")} disconnected; ` +
 				'their tools answer "not connected" until the server is reconnected (the user can do that in /mcp).',
 		);
 	}
 	return notices;
 }
+
+/**
+ * The user's startup notices, Claude Code's wording (`useMcpConnectivityStatus`):
+ * `2 MCP servers failed · /mcp` and `1 MCP server needs auth · /mcp`, with each
+ * server's error left to /mcp. Printing every server's raw error ("fetch
+ * failed" per remote server when offline) buried the prompt.
+ */
+export function mcpStartupNotices(failedServers: number, needsAuth: number): string[] {
+	const out: string[] = [];
+	if (failedServers > 0) out.push(`${countNoun(failedServers, "MCP server")} failed · /mcp`);
+	if (needsAuth > 0) out.push(`${countNoun(needsAuth, "MCP server")} ${needsAuth === 1 ? "needs" : "need"} auth · /mcp`);
+	return out;
+}

package/extensions/mcp/client.ts

@@ -5,9 +5,13 @@
 
 import type { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import type { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import type { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
 import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
 import type { OAuthClientProvider } from "@modelcontextprotocol/sdk/client/auth.js";
+import { childProcessEnv } from "../lib/app-launch.mjs";
 import type { McpServer } from "./config.ts";
+import { collectPages } from "./discovery.ts";
+import { McpAuthenticationRequiredError } from "./oauth/provider.ts";
 
 /**
  * The MCP SDK costs ~70-80ms to load (findings §15), so it is imported on the
@@ -19,6 +23,7 @@ let sdkPromise:
 			Client: typeof import("@modelcontextprotocol/sdk/client/index.js").Client;
 			StdioClientTransport: typeof import("@modelcontextprotocol/sdk/client/stdio.js").StdioClientTransport;
 			StreamableHTTPClientTransport: typeof import("@modelcontextprotocol/sdk/client/streamableHttp.js").StreamableHTTPClientTransport;
+			SSEClientTransport: typeof import("@modelcontextprotocol/sdk/client/sse.js").SSEClientTransport;
 			UnauthorizedError: typeof import("@modelcontextprotocol/sdk/client/auth.js").UnauthorizedError;
 	  }>
 	| undefined;
@@ -27,6 +32,10 @@ let sdkPromise:
 // (the SDK's UnauthorizedError does not set `.name`, so a name/message check is
 // unreliable — a custom-message instance would be missed).
 let unauthorizedClass: typeof import("@modelcontextprotocol/sdk/client/auth.js").UnauthorizedError | undefined;
+let httpErrorClass: typeof import("@modelcontextprotocol/sdk/client/streamableHttp.js").StreamableHTTPError | undefined;
+let sseErrorClass: typeof import("@modelcontextprotocol/sdk/client/sse.js").SseError | undefined;
+
+type HttpTransport = StreamableHTTPClientTransport | SSEClientTransport;
 
 function loadSdk() {
 	sdkPromise ??= Promise.all([
@@ -34,12 +43,16 @@ function loadSdk() {
 		import("@modelcontextprotocol/sdk/client/stdio.js"),
 		import("@modelcontextprotocol/sdk/client/streamableHttp.js"),
 		import("@modelcontextprotocol/sdk/client/auth.js"),
-	]).then(([client, stdio, http, auth]) => {
+		import("@modelcontextprotocol/sdk/client/sse.js"),
+	]).then(([client, stdio, http, auth, sse]) => {
 		unauthorizedClass = auth.UnauthorizedError;
+		httpErrorClass = http.StreamableHTTPError;
+		sseErrorClass = sse.SseError;
 		return {
 			Client: client.Client,
 			StdioClientTransport: stdio.StdioClientTransport,
 			StreamableHTTPClientTransport: http.StreamableHTTPClientTransport,
+			SSEClientTransport: sse.SSEClientTransport,
 			UnauthorizedError: auth.UnauthorizedError,
 		};
 	});
@@ -62,11 +75,21 @@ export interface McpResource {
 	mimeType?: string;
 }
 
+export interface McpResourceTemplate {
+	uriTemplate: string;
+	name?: string;
+	description?: string;
+	mimeType?: string;
+}
+
 export interface Connection {
 	server: McpServer;
 	client: Client;
 	tools: McpTool[];
 	resources: McpResource[];
+	resourceTemplates?: McpResourceTemplate[];
+	/** Called after a tool-list refresh, or with its error while the last good list stays intact. */
+	onToolsChanged?: (error?: Error) => void;
 	/** The server's own usage instructions from its initialize result, if any. */
 	instructions?: string;
 	/** Non-fatal problems after a successful connect (e.g. listTools failed). */
@@ -126,8 +149,11 @@ function withTimeout<T>(promise: Promise<T>, ms: number, what: string): Promise<
 	});
 }
 
-/** True when the error is the SDK's "authorization required" signal (an HTTP 401). */
+/** True when the server requires authentication, including an expired stored grant. */
 export function isUnauthorized(error: unknown): boolean {
+	if (error instanceof McpAuthenticationRequiredError) return true;
+	if (httpErrorClass && error instanceof httpErrorClass && error.code === 401) return true;
+	if (sseErrorClass && error instanceof sseErrorClass && error.code === 401) return true;
 	// The class is captured on the first loadSdk(); every caller runs after a
 	// connect (which awaits it), so it is set. The message fallback only covers
 	// the impossible pre-load case.
@@ -140,61 +166,140 @@ function buildTransport(
 	server: McpServer,
 	sdk: Awaited<ReturnType<typeof loadSdk>>,
 	authProvider?: OAuthClientProvider,
+	keepOAuthOpen?: () => boolean,
 ): { transport: Transport; stderrTail: () => string } {
 	const tail = createTailBuffer();
 	if (server.kind === "stdio") {
 		const transport = new sdk.StdioClientTransport({
 			command: server.command,
 			args: server.args,
-			env: { ...(process.env as Record<string, string>), ...(server.env ?? {}) },
+			// The user's own environment under the bundled app (lib/app-launch.mjs), then the server's.
+			env: { ...(childProcessEnv(process.env) as Record<string, string>), ...(server.env ?? {}) },
+			...(server.cwd ? { cwd: server.cwd } : {}),
 			stderr: "pipe",
 		});
 		// The PassThrough exists before start(); drain it from the first byte.
 		transport.stderr?.on("data", (chunk: string | Uint8Array) => tail.push(chunk));
 		return { transport, stderrTail: tail.text };
 	}
-	const transport = new sdk.StreamableHTTPClientTransport(new URL(server.url), {
+	const lifetime = new AbortController();
+	const RemoteTransport = server.transport === "sse" ? sdk.SSEClientTransport : sdk.StreamableHTTPClientTransport;
+	const transport = new RemoteTransport(new URL(server.url), {
 		authProvider,
 		requestInit: server.headers ? { headers: server.headers } : undefined,
+		// The SDK's close signal covers MCP requests, but not the discovery and
+		// token fetches it starts during OAuth. Those must stop with the client too.
+		fetch: (input, init) => fetch(input, {
+			...init,
+			signal: init?.signal ? AbortSignal.any([lifetime.signal, init.signal]) : lifetime.signal,
+		}),
 	});
+	const closeTransport = transport.close.bind(transport);
+	transport.close = async () => {
+		if (!keepOAuthOpen?.()) lifetime.abort();
+		await closeTransport();
+	};
 	return { transport, stderrTail: tail.text };
 }
 
-/** List a connected client's tools and resources into a Connection. */
+/** Bound the complete list, not each page, and cancel its last request on timeout. */
+async function discoverPages<T>(
+	fetchPage: (cursor: string | undefined, signal: AbortSignal) => Promise<{ items: T[]; nextCursor?: string }>,
+	what: string,
+): Promise<T[]> {
+	const controller = new AbortController();
+	try {
+		return await withTimeout(collectPages((cursor) => fetchPage(cursor, controller.signal), controller.signal), CONNECT_TIMEOUT_MS, what);
+	} catch (error) {
+		controller.abort();
+		throw error;
+	}
+}
+
+/** A fresh, complete tool snapshot. */
+export function discoverTools(client: Client, serverName: string): Promise<McpTool[]> {
+	return discoverPages(async (cursor, signal) => {
+		const result = await client.listTools(cursor === undefined ? undefined : { cursor }, { signal });
+		return { items: result.tools, nextCursor: result.nextCursor };
+	}, `listing tools of "${serverName}"`);
+}
+
+/** List a connected client's tools, resources and URI templates into a Connection. */
 async function finalizeConnection(client: Client, server: McpServer, stderrTail: () => string = () => ""): Promise<Connection> {
 	const warnings: string[] = [];
 
-	// Each list call can take up to CONNECT_TIMEOUT_MS on a slow server; run
-	// them concurrently instead of stacking their worst-case latencies.
-	const [toolsResult, resourcesResult] = await Promise.allSettled([
-		withTimeout(client.listTools(), CONNECT_TIMEOUT_MS, `listing tools of "${server.name}"`),
-		withTimeout(client.listResources(), CONNECT_TIMEOUT_MS, `listing resources of "${server.name}"`),
+	// Each complete list has one deadline; run them concurrently instead of
+	// stacking their worst-case latencies or granting each page a fresh timeout.
+	const [toolsResult, resourcesResult, templatesResult] = await Promise.allSettled([
+		discoverTools(client, server.name),
+		discoverPages(async (cursor, signal) => {
+			const result = await client.listResources(cursor === undefined ? undefined : { cursor }, { signal });
+			return { items: result.resources, nextCursor: result.nextCursor };
+		}, `listing resources of "${server.name}"`),
+		discoverPages(async (cursor, signal) => {
+			const result = await client.listResourceTemplates(cursor === undefined ? undefined : { cursor }, { signal });
+			return { items: result.resourceTemplates, nextCursor: result.nextCursor };
+		}, `listing resource templates of "${server.name}"`),
 	]);
 
 	let tools: McpTool[] = [];
 	if (toolsResult.status === "fulfilled") {
-		tools = (toolsResult.value.tools ?? []) as McpTool[];
+		tools = toolsResult.value;
 	} else {
 		// A server that connects but fails listTools would otherwise register
 		// zero tools with no signal anywhere — surface it as a warning.
 		warnings.push(`connected, but listing tools failed: ${(toolsResult.reason as Error).message}`);
 	}
 
-	let resources: McpResource[] = [];
-	if (resourcesResult.status === "fulfilled") {
-		resources = (resourcesResult.value.resources ?? []) as McpResource[];
-	} else {
-		// Resources are optional in MCP; "method not found" is normal. Anything
-		// else (a timeout, a protocol error) is worth a warning.
-		const message = (resourcesResult.reason as Error).message;
-		if (!/method not found|-32601/i.test(message)) {
-			warnings.push(`connected, but listing resources failed: ${message}`);
-		}
-	}
-
+	// Resources and templates are optional in MCP; "method not found" is
+	// normal. A timeout or protocol error must still surface as a warning.
+	const optionalList = <T>(result: PromiseSettledResult<T[]>, label: string): T[] => {
+		if (result.status === "fulfilled") return result.value;
+		const message = (result.reason as Error).message;
+		if (!/method not found|-32601/i.test(message)) warnings.push(`connected, but listing ${label} failed: ${message}`);
+		return [];
+	};
+	const resources = optionalList(resourcesResult, "resources");
+	const resourceTemplates = optionalList(templatesResult, "resource templates");
 	const instructions = client.getInstructions()?.trim() || undefined;
 
-	return { server, client, tools, resources, instructions, warnings, stderrTail };
+	const connection: Connection = { server, client, tools, resources, resourceTemplates, instructions, warnings, stderrTail };
+	const { ToolListChangedNotificationSchema } = await import("@modelcontextprotocol/sdk/types.js");
+	let refreshing = false;
+	let pending = false;
+	client.setNotificationHandler(ToolListChangedNotificationSchema, async () => {
+		if (connection.closing) return;
+		pending = true;
+		if (refreshing) return;
+		refreshing = true;
+		try {
+			// A burst during an in-flight list gets one follow-up refresh, never
+			// concurrent snapshots that could overwrite a newer list with an old one.
+			while (pending && !connection.closing) {
+				pending = false;
+				let refreshed: McpTool[];
+				try {
+					refreshed = await discoverTools(client, server.name);
+				} catch (cause) {
+					if (connection.closing) return;
+					const error = cause instanceof Error ? cause : new Error(String(cause));
+					const warning = `connected, but refreshing tools failed: ${error.message}`;
+					if (!connection.warnings.includes(warning)) connection.warnings.push(warning);
+					connection.onToolsChanged?.(error);
+					continue;
+				}
+				if (connection.closing) return;
+				connection.tools = refreshed;
+				connection.onToolsChanged?.();
+			}
+		} finally {
+			refreshing = false;
+		}
+	});
+	// A server can exit after initialize but before discovery finishes. Its
+	// onclose already fired, before index.ts can adopt the connection.
+	if (!client.transport) throw new Error(withStderr(`connection to "${server.name}" closed during discovery`, stderrTail()));
+	return connection;
 }
 
 /**
@@ -208,10 +313,11 @@ async function finalizeConnection(client: Client, server: McpServer, stderrTail:
 async function openClient(
 	server: McpServer,
 	authProvider?: OAuthClientProvider,
+	keepOAuthOpen?: () => boolean,
 ): Promise<{ client: Client; transport: Transport; stderrTail: () => string }> {
 	const sdk = await loadSdk();
 	const client = new sdk.Client({ name: "one-code", version: "0.1.0" }, { capabilities: {} });
-	return { client, ...buildTransport(server, sdk, authProvider) };
+	return { client, ...buildTransport(server, sdk, authProvider, keepOAuthOpen) };
 }
 
 /**
@@ -254,22 +360,35 @@ export async function connect(server: McpServer, authProvider?: OAuthClientProvi
 export async function beginInteractiveAuth(
 	server: McpServer,
 	authProvider: OAuthClientProvider,
-): Promise<{ transport: StreamableHTTPClientTransport } | { connection: Connection }> {
+): Promise<{ transport: HttpTransport } | { connection: Connection }> {
 	if (server.kind !== "http") throw new Error(`OAuth is only available for http MCP servers; "${server.name}" is stdio.`);
-	const { client, transport: baseTransport, stderrTail } = await openClient(server, authProvider);
-	const transport = baseTransport as StreamableHTTPClientTransport;
+	// Client.connect closes the transport after the expected UnauthorizedError.
+	// Keep OAuth fetches usable until the redirect has been handed to the caller.
+	let awaitingRedirect = true;
+	const { client, transport: baseTransport, stderrTail } = await openClient(server, authProvider, () => awaitingRedirect);
+	const transport = baseTransport as HttpTransport;
 	try {
 		// Not connectOrClose: on a 401 the transport must stay open — it carries
 		// the OAuth flow state the caller finishes with finishAuth(code).
 		await withTimeout(client.connect(transport), CONNECT_TIMEOUT_MS, `connecting to "${server.name}"`);
 		return { connection: await finalizeConnection(client, server, stderrTail) };
 	} catch (error) {
-		if (isUnauthorized(error)) return { transport };
+		// Only the SDK's redirect signal means a browser callback is coming.
+		// A raw 401 after refreshing tokens is an auth failure, not a redirect.
+		if (unauthorizedClass && error instanceof unauthorizedClass) return { transport };
+		awaitingRedirect = false;
 		await transport.close().catch(() => {});
 		throw error;
+	} finally {
+		awaitingRedirect = false;
 	}
 }
 
+/** Bound the token exchange too; the callback deadline only covers waiting for the browser. */
+export async function finishInteractiveAuth(transport: HttpTransport, code: string): Promise<void> {
+	await withTimeout(transport.finishAuth(code), CONNECT_TIMEOUT_MS, "exchanging the authorization code");
+}
+
 const INSTRUCTIONS_CAP = 3000;
 
 /**
@@ -322,21 +441,15 @@ export async function callTool(
 	connection: Connection,
 	toolName: string,
 	args: Record<string, unknown>,
-): Promise<{ content?: unknown[]; isError?: boolean }> {
-	const result = await withTimeout(
-		connection.client.callTool({ name: toolName, arguments: args }),
-		CALL_TIMEOUT_MS,
-		`calling "${toolName}"`,
-	);
-	return result as { content?: unknown[]; isError?: boolean };
+	signal?: AbortSignal,
+): Promise<{ content?: unknown[]; isError?: boolean; structuredContent?: Record<string, unknown> }> {
+	return await connection.client.callTool({ name: toolName, arguments: args }, undefined, { signal, timeout: CALL_TIMEOUT_MS }) as {
+		content?: unknown[]; isError?: boolean; structuredContent?: Record<string, unknown>;
+	};
 }
 
-export async function readResource(connection: Connection, uri: string): Promise<{ contents?: unknown[] }> {
-	return (await withTimeout(
-		connection.client.readResource({ uri }),
-		CALL_TIMEOUT_MS,
-		`reading resource "${uri}"`,
-	)) as { contents?: unknown[] };
+export async function readResource(connection: Connection, uri: string, signal?: AbortSignal): Promise<{ contents?: unknown[] }> {
+	return await connection.client.readResource({ uri }, { signal, timeout: CALL_TIMEOUT_MS });
 }
 
 /**
@@ -345,13 +458,13 @@ export async function readResource(connection: Connection, uri: string): Promise
  * The SDK only calls `safeParse` on the result schema, and the response shape is
  * server-defined, so a passthrough schema beats pinning one the spec hasn't settled.
  */
-export async function readResourceDir(connection: Connection, uri: string): Promise<Record<string, unknown>> {
+export async function readResourceDir(connection: Connection, uri: string, signal?: AbortSignal): Promise<Record<string, unknown>> {
 	const passthrough = { safeParse: (data: unknown) => ({ success: true as const, data }) };
-	return (await withTimeout(
-		connection.client.request({ method: "resources/directory/read", params: { uri } }, passthrough as never),
-		CALL_TIMEOUT_MS,
-		`listing directory "${uri}"`,
-	)) as Record<string, unknown>;
+	return await connection.client.request(
+		{ method: "resources/directory/read", params: { uri } },
+		passthrough as never,
+		{ signal, timeout: CALL_TIMEOUT_MS },
+	) as Record<string, unknown>;
 }
 
 export async function close(connection: Connection): Promise<void> {