agents 0.5.1 → 0.7.0

package/dist/mcp/client.d.ts

@@ -1,6 +1,6 @@
 import { RetryOptions } from "../retries.js";
-import { d as Event, i as MCPTransportOptions, l as TransportType, n as MCPClientConnection, r as MCPConnectionState, t as MCPServerRow, u as Emitter } from "../client-storage-D633wI1S.js";
-import { t as MCPObservabilityEvent } from "../mcp-DA0kDE7K.js";
+import { E as Event, T as Emitter, i as MCPTransportOptions, n as MCPClientConnection, r as MCPConnectionState, t as MCPServerRow, w as TransportType } from "../client-storage-tusTuoSF.js";
+import { n as MCPObservabilityEvent } from "../agent-DnmmRjyv.js";
 import { AgentMcpOAuthProvider } from "./do-oauth-client-provider.js";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { CallToolRequest, CallToolResultSchema, CompatibilityCallToolResultSchema, GetPromptRequest, Prompt, ReadResourceRequest, Resource, ResourceTemplate, Tool } from "@modelcontextprotocol/sdk/types.js";
@@ -39,7 +39,7 @@ type MCPOAuthCallbackResult = {
 type RegisterServerOptions = {
   url: string;
   name: string;
-  callbackUrl: string;
+  callbackUrl?: string;
   client?: ConstructorParameters<typeof Client>[1];
   transport?: MCPTransportOptions;
   authUrl?: string;
@@ -102,6 +102,7 @@ declare class MCPClientManager {
   private _storage;
   private _createAuthProviderFn?;
   private _isRestored;
+  private _pendingConnections;
   /** @internal Protected for testing purposes. */
   protected readonly _onObservabilityEvent: Emitter<MCPObservabilityEvent>;
   readonly onObservabilityEvent: Event<MCPObservabilityEvent>;
@@ -133,13 +134,46 @@ declare class MCPClientManager {
    * @internal
    */
   private createAuthProvider;
+  /**
+   * Get saved RPC servers from storage (servers with rpc:// URLs).
+   * These are restored separately by the Agent class since they need env bindings.
+   */
+  getRpcServersFromStorage(): MCPServerRow[];
+  /**
+   * Save an RPC server to storage for hibernation recovery.
+   * The bindingName is stored in server_options so the Agent can look up
+   * the namespace from env during restore.
+   */
+  saveRpcServerToStorage(id: string, name: string, normalizedName: string, bindingName: string, props?: Record<string, unknown>): void;
   /**
    * Restore MCP server connections from storage
-   * This method is called on Agent initialization to restore previously connected servers
+   * This method is called on Agent initialization to restore previously connected servers.
+   * RPC servers (rpc:// URLs) are skipped here -- they are restored by the Agent class
+   * which has access to env bindings.
    *
    * @param clientName Name to use for OAuth client (typically the agent instance name)
    */
   restoreConnectionsFromStorage(clientName: string): Promise<void>;
+  /**
+   * Track a pending connection promise for a server.
+   * The promise is removed from the map when it settles.
+   */
+  private _trackConnection;
+  /**
+   * Wait for all in-flight connection and discovery operations to settle.
+   * This is useful when you need MCP tools to be available before proceeding,
+   * e.g. before calling getAITools() after the agent wakes from hibernation.
+   *
+   * Returns once every pending connection has either connected and discovered,
+   * failed, or timed out. Never rejects.
+   *
+   * @param options.timeout - Maximum time in milliseconds to wait.
+   *   `0` returns immediately without waiting.
+   *   `undefined` (default) waits indefinitely.
+   */
+  waitForConnections(options?: {
+    timeout?: number;
+  }): Promise<void>;
   /**
    * Internal method to restore a single server connection and discovery
    */
@@ -218,11 +252,14 @@ declare class MCPClientManager {
     timeoutMs?: number;
   }): Promise<MCPDiscoverResult | undefined>;
   /**
-   * Establish connection in the background after OAuth completion
-   * This method connects to the server and discovers its capabilities
+   * Establish connection in the background after OAuth completion.
+   * This method connects to the server and discovers its capabilities.
+   * The connection is automatically tracked so that `waitForConnections()`
+   * will include it.
    * @param serverId The server ID to establish connection for
    */
   establishConnection(serverId: string): Promise<void>;
+  private _doEstablishConnection;
   /**
    * Configure OAuth callback handling
    * @param config OAuth callback configuration
@@ -333,7 +370,14 @@ declare class MCPClientManager {
         priority?: number | undefined;
         lastModified?: string | undefined;
       } | undefined;
-      _meta?: Record<string, unknown> | undefined;
+      _meta
+      /**
+       * We need to delay loading ai sdk, because putting it in module scope is
+       * causing issues with startup time.
+       * The only place it's used is in getAITools, which only matters after
+       * .connect() is called on at least one server.
+       * So it's safe to delay loading it until .connect() is called.
+       */?: Record<string, unknown> | undefined;
     } | {
       type: "resource";
       resource: {

package/dist/mcp/client.js

@@ -1,5 +1,5 @@
 import { tryN } from "../retries.js";
-import { a as Emitter, i as DisposableStore, n as MCPConnectionState, r as toErrorMessage, t as MCPClientConnection } from "../client-connection-CGMuV62J.js";
+import { a as RPC_DO_PREFIX, c as Emitter, n as MCPConnectionState, o as toErrorMessage, s as DisposableStore, t as MCPClientConnection } from "../client-connection-D3Wcd6Q6.js";
 import { DurableObjectOAuthClientProvider } from "./do-oauth-client-provider.js";
 import { nanoid } from "nanoid";
 import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker-provider.js";
@@ -7,6 +7,47 @@ import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validatio
 //#region src/mcp/client.ts
 const defaultClientOptions = { jsonSchemaValidator: new CfWorkerJsonSchemaValidator() };
 /**
+* Blocked hostname patterns for SSRF protection.
+* Prevents MCP client from connecting to internal/private network addresses.
+*/
+const BLOCKED_HOSTNAMES = new Set([
+	"localhost",
+	"0.0.0.0",
+	"[::1]",
+	"[::]",
+	"metadata.google.internal"
+]);
+/**
+* Check whether a hostname looks like a private/internal IP address.
+* Blocks RFC 1918, link-local, loopback, and cloud metadata endpoints.
+*/
+function isBlockedUrl(url) {
+	let parsed;
+	try {
+		parsed = new URL(url);
+	} catch {
+		return true;
+	}
+	const hostname = parsed.hostname;
+	if (BLOCKED_HOSTNAMES.has(hostname)) return true;
+	const ipv4Parts = hostname.split(".");
+	if (ipv4Parts.length === 4 && ipv4Parts.every((p) => /^\d{1,3}$/.test(p))) {
+		const [a, b] = ipv4Parts.map(Number);
+		if (a === 10) return true;
+		if (a === 172 && b >= 16 && b <= 31) return true;
+		if (a === 192 && b === 168) return true;
+		if (a === 127) return true;
+		if (a === 169 && b === 254) return true;
+		if (a === 0) return true;
+	}
+	if (hostname.startsWith("[") && hostname.endsWith("]")) {
+		const addr = hostname.slice(1, -1).toLowerCase();
+		if (addr.startsWith("fc") || addr.startsWith("fd")) return true;
+		if (addr.startsWith("fe80")) return true;
+	}
+	return false;
+}
+/**
 * Utility class that aggregates multiple MCP clients into one
 */
 var MCPClientManager = class {
@@ -22,6 +63,7 @@ var MCPClientManager = class {
 		this._didWarnAboutUnstableGetAITools = false;
 		this._connectionDisposables = /* @__PURE__ */ new Map();
 		this._isRestored = false;
+		this._pendingConnections = /* @__PURE__ */ new Map();
 		this._onObservabilityEvent = new Emitter();
 		this.onObservabilityEvent = this._onObservabilityEvent.event;
 		this._onServerStateChanged = new Emitter();
@@ -80,8 +122,36 @@ var MCPClientManager = class {
 		return authProvider;
 	}
 	/**
+	* Get saved RPC servers from storage (servers with rpc:// URLs).
+	* These are restored separately by the Agent class since they need env bindings.
+	*/
+	getRpcServersFromStorage() {
+		return this.getServersFromStorage().filter((s) => s.server_url.startsWith(RPC_DO_PREFIX));
+	}
+	/**
+	* Save an RPC server to storage for hibernation recovery.
+	* The bindingName is stored in server_options so the Agent can look up
+	* the namespace from env during restore.
+	*/
+	saveRpcServerToStorage(id, name, normalizedName, bindingName, props) {
+		this.saveServerToStorage({
+			id,
+			name,
+			server_url: `${RPC_DO_PREFIX}${normalizedName}`,
+			client_id: null,
+			auth_url: null,
+			callback_url: "",
+			server_options: JSON.stringify({
+				bindingName,
+				props
+			})
+		});
+	}
+	/**
 	* Restore MCP server connections from storage
-	* This method is called on Agent initialization to restore previously connected servers
+	* This method is called on Agent initialization to restore previously connected servers.
+	* RPC servers (rpc:// URLs) are skipped here -- they are restored by the Agent class
+	* which has access to env bindings.
 	*
 	* @param clientName Name to use for OAuth client (typically the agent instance name)
 	*/
@@ -93,6 +163,7 @@ var MCPClientManager = class {
 			return;
 		}
 		for (const server of servers) {
+			if (server.server_url.startsWith(RPC_DO_PREFIX)) continue;
 			const existingConn = this.mcpConnections[server.id];
 			if (existingConn) {
 				if (existingConn.connectionState === MCPConnectionState.READY) {
@@ -112,9 +183,12 @@ var MCPClientManager = class {
 				}
 			}
 			const parsedOptions = server.server_options ? JSON.parse(server.server_options) : null;
-			const authProvider = this._createAuthProviderFn ? this._createAuthProviderFn(server.callback_url) : this.createAuthProvider(server.id, server.callback_url, clientName, server.client_id ?? void 0);
-			authProvider.serverId = server.id;
-			if (server.client_id) authProvider.clientId = server.client_id;
+			let authProvider;
+			if (server.callback_url) {
+				authProvider = this._createAuthProviderFn ? this._createAuthProviderFn(server.callback_url) : this.createAuthProvider(server.id, server.callback_url, clientName, server.client_id ?? void 0);
+				authProvider.serverId = server.id;
+				if (server.client_id) authProvider.clientId = server.client_id;
+			}
 			const conn = this.createConnection(server.id, server.server_url, {
 				client: parsedOptions?.client ?? {},
 				transport: {
@@ -127,11 +201,46 @@ var MCPClientManager = class {
 				conn.connectionState = MCPConnectionState.AUTHENTICATING;
 				continue;
 			}
-			this._restoreServer(server.id, parsedOptions?.retry);
+			this._trackConnection(server.id, this._restoreServer(server.id, parsedOptions?.retry));
 		}
 		this._isRestored = true;
 	}
 	/**
+	* Track a pending connection promise for a server.
+	* The promise is removed from the map when it settles.
+	*/
+	_trackConnection(serverId, promise) {
+		const tracked = promise.finally(() => {
+			if (this._pendingConnections.get(serverId) === tracked) this._pendingConnections.delete(serverId);
+		});
+		this._pendingConnections.set(serverId, tracked);
+	}
+	/**
+	* Wait for all in-flight connection and discovery operations to settle.
+	* This is useful when you need MCP tools to be available before proceeding,
+	* e.g. before calling getAITools() after the agent wakes from hibernation.
+	*
+	* Returns once every pending connection has either connected and discovered,
+	* failed, or timed out. Never rejects.
+	*
+	* @param options.timeout - Maximum time in milliseconds to wait.
+	*   `0` returns immediately without waiting.
+	*   `undefined` (default) waits indefinitely.
+	*/
+	async waitForConnections(options) {
+		if (this._pendingConnections.size === 0) return;
+		if (options?.timeout != null && options.timeout <= 0) return;
+		const settled = Promise.allSettled(this._pendingConnections.values());
+		if (options?.timeout != null && options.timeout > 0) {
+			let timerId;
+			const timer = new Promise((resolve) => {
+				timerId = setTimeout(resolve, options.timeout);
+			});
+			await Promise.race([settled, timer]);
+			clearTimeout(timerId);
+		} else await settled;
+	}
+	/**
 	* Internal method to restore a single server connection and discovery
 	*/
 	async _restoreServer(serverId, retry) {
@@ -171,6 +280,7 @@ var MCPClientManager = class {
 			options.transport.authProvider.serverId = id;
 			if (options.reconnect?.oauthClientId) options.transport.authProvider.clientId = options.reconnect?.oauthClientId;
 		}
+		if (isBlockedUrl(url)) throw new Error(`Blocked URL: ${url} — MCP client connections to private/internal addresses are not allowed`);
 		if (!options.reconnect?.oauthCode || !this.mcpConnections[id]) {
 			const normalizedTransport = {
 				...options.transport,
@@ -198,15 +308,13 @@ var MCPClientManager = class {
 		} catch (error) {
 			this._onObservabilityEvent.fire({
 				type: "mcp:client:connect",
-				displayMessage: `Failed to complete OAuth reconnection for ${id} for ${url}`,
 				payload: {
 					url,
 					transport: options.transport?.type ?? "auto",
 					state: this.mcpConnections[id].connectionState,
 					error: toErrorMessage(error)
 				},
-				timestamp: Date.now(),
-				id
+				timestamp: Date.now()
 			});
 			throw error;
 		}
@@ -259,6 +367,7 @@ var MCPClientManager = class {
 	* @returns Server ID
 	*/
 	async registerServer(id, options) {
+		if (isBlockedUrl(options.url)) throw new Error(`Blocked URL: ${options.url} — MCP client connections to private/internal addresses are not allowed`);
 		this.createConnection(id, options.url, {
 			client: options.client,
 			transport: {
@@ -271,7 +380,7 @@ var MCPClientManager = class {
 			id,
 			name: options.name,
 			server_url: options.url,
-			callback_url: options.callbackUrl,
+			callback_url: options.callbackUrl ?? "",
 			client_id: options.clientId ?? null,
 			auth_url: options.authUrl ?? null,
 			server_options: JSON.stringify({
@@ -324,6 +433,15 @@ var MCPClientManager = class {
 					});
 					this._onServerStateChanged.fire();
 				}
+				this._onObservabilityEvent.fire({
+					type: "mcp:client:authorize",
+					payload: {
+						serverId: id,
+						authUrl,
+						clientId
+					},
+					timestamp: Date.now()
+				});
 				return {
 					state: conn.connectionState,
 					authUrl,
@@ -459,10 +577,8 @@ var MCPClientManager = class {
 		if (!conn) {
 			this._onObservabilityEvent.fire({
 				type: "mcp:client:discover",
-				displayMessage: `Connection not found for ${serverId}`,
 				payload: {},
-				timestamp: Date.now(),
-				id: nanoid()
+				timestamp: Date.now()
 			});
 			return;
 		}
@@ -474,33 +590,36 @@ var MCPClientManager = class {
 		};
 	}
 	/**
-	* Establish connection in the background after OAuth completion
-	* This method connects to the server and discovers its capabilities
+	* Establish connection in the background after OAuth completion.
+	* This method connects to the server and discovers its capabilities.
+	* The connection is automatically tracked so that `waitForConnections()`
+	* will include it.
 	* @param serverId The server ID to establish connection for
 	*/
 	async establishConnection(serverId) {
+		const promise = this._doEstablishConnection(serverId);
+		this._trackConnection(serverId, promise);
+		return promise;
+	}
+	async _doEstablishConnection(serverId) {
 		const conn = this.mcpConnections[serverId];
 		if (!conn) {
 			this._onObservabilityEvent.fire({
 				type: "mcp:client:preconnect",
-				displayMessage: `Connection not found for serverId: ${serverId}`,
 				payload: { serverId },
-				timestamp: Date.now(),
-				id: nanoid()
+				timestamp: Date.now()
 			});
 			return;
 		}
 		if (conn.connectionState === MCPConnectionState.DISCOVERING || conn.connectionState === MCPConnectionState.READY) {
 			this._onObservabilityEvent.fire({
 				type: "mcp:client:connect",
-				displayMessage: `establishConnection skipped for ${serverId}, already in ${conn.connectionState} state`,
 				payload: {
 					url: conn.url.toString(),
 					transport: conn.options.transport.type || "unknown",
 					state: conn.connectionState
 				},
-				timestamp: Date.now(),
-				id: nanoid()
+				timestamp: Date.now()
 			});
 			return;
 		}
@@ -513,14 +632,12 @@ var MCPClientManager = class {
 		if (connectResult.state === MCPConnectionState.CONNECTED) await this.discoverIfConnected(serverId);
 		this._onObservabilityEvent.fire({
 			type: "mcp:client:connect",
-			displayMessage: `establishConnection completed for ${serverId}, final state: ${conn.connectionState}`,
 			payload: {
 				url: conn.url.toString(),
 				transport: conn.options.transport.type || "unknown",
 				state: conn.connectionState
 			},
-			timestamp: Date.now(),
-			id: nanoid()
+			timestamp: Date.now()
 		});
 	}
 	/**
@@ -566,8 +683,10 @@ var MCPClientManager = class {
 	getAITools() {
 		if (!this.jsonSchema) throw new Error("jsonSchema not initialized.");
 		for (const [id, conn] of Object.entries(this.mcpConnections)) if (conn.connectionState !== MCPConnectionState.READY && conn.connectionState !== MCPConnectionState.AUTHENTICATING) console.warn(`[getAITools] WARNING: Reading tools from connection ${id} in state "${conn.connectionState}". Tools may not be loaded yet.`);
-		return Object.fromEntries(getNamespacedData(this.mcpConnections, "tools").map((tool) => {
-			return [`tool_${tool.serverId.replace(/-/g, "")}_${tool.name}`, {
+		const entries = [];
+		for (const tool of getNamespacedData(this.mcpConnections, "tools")) try {
+			const toolKey = `tool_${tool.serverId.replace(/-/g, "")}_${tool.name}`;
+			entries.push([toolKey, {
 				description: tool.description,
 				execute: async (args) => {
 					const result = await this.callTool({
@@ -582,10 +701,13 @@ var MCPClientManager = class {
 					}
 					return result;
 				},
-				inputSchema: this.jsonSchema(tool.inputSchema),
+				inputSchema: tool.inputSchema ? this.jsonSchema(tool.inputSchema) : this.jsonSchema({ type: "object" }),
 				outputSchema: tool.outputSchema ? this.jsonSchema(tool.outputSchema) : void 0
-			}];
-		}));
+			}]);
+		} catch (e) {
+			console.warn(`[getAITools] Skipping tool "${tool.name}" from "${tool.serverId}": ${e}`);
+		}
+		return Object.fromEntries(entries);
 	}
 	/**
 	* @deprecated this has been renamed to getAITools(), and unstable_getAITools will be removed in the next major version
@@ -610,6 +732,7 @@ var MCPClientManager = class {
 	*/
 	async closeAllConnections() {
 		const ids = Object.keys(this.mcpConnections);
+		this._pendingConnections.clear();
 		for (const id of ids) this.mcpConnections[id].cancelDiscovery();
 		await Promise.all(ids.map(async (id) => {
 			await this.mcpConnections[id].client.close();
@@ -628,6 +751,7 @@ var MCPClientManager = class {
 	async closeConnection(id) {
 		if (!this.mcpConnections[id]) throw new Error(`Connection with id "${id}" does not exist.`);
 		this.mcpConnections[id].cancelDiscovery();
+		this._pendingConnections.delete(id);
 		await this.mcpConnections[id].client.close();
 		delete this.mcpConnections[id];
 		const store = this._connectionDisposables.get(id);
@@ -683,9 +807,10 @@ var MCPClientManager = class {
 	* Namespaced version of callTool
 	*/
 	async callTool(params, resultSchema, options) {
-		const unqualifiedName = params.name.replace(`${params.serverId}.`, "");
-		return this.mcpConnections[params.serverId].client.callTool({
-			...params,
+		const { serverId, ...mcpParams } = params;
+		const unqualifiedName = mcpParams.name.replace(`${serverId}.`, "");
+		return this.mcpConnections[serverId].client.callTool({
+			...mcpParams,
 			name: unqualifiedName
 		}, resultSchema, options);
 	}