hono-agents 0.0.0-cccbd0f → 0.0.0-cd8b7fd

package/dist/index.js

@@ -1,3 +1,4 @@
+import { AsyncLocalStorage } from "node:async_hooks";
 import "partysocket";
 import { nanoid } from "nanoid";
 import { CfWorkerJsonSchemaValidator } from "@modelcontextprotocol/sdk/validation/cfworker-provider.js";
@@ -5,14 +6,17 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
 import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
 import { ElicitRequestSchema, PromptListChangedNotificationSchema, ResourceListChangedNotificationSchema, ToolListChangedNotificationSchema } from "@modelcontextprotocol/sdk/types.js";
-import { AsyncLocalStorage } from "node:async_hooks";
 import { parseCronExpression } from "cron-schedule";
 import "cloudflare:email";
 import { Server, routePartykitRequest } from "partyserver";
 import { env } from "hono/adapter";
 import { createMiddleware } from "hono/factory";
 
-//#region ../agents/dist/ai-types-CrMqkwc_.js
+//#region ../agents/dist/context-BkKbAa1R.js
+const agentContext = new AsyncLocalStorage();
+
+//#endregion
+//#region ../agents/dist/ai-types-CwgHzwUb.js
 /**
 * Enum for message types to improve type safety and maintainability
 */
@@ -34,7 +38,7 @@ let MessageType = /* @__PURE__ */ function(MessageType$1) {
 }({});
 
 //#endregion
-//#region ../agents/dist/client-B3SR12TQ.js
+//#region ../agents/dist/client-CcyhkGfN.js
 /**
 * Convert a camelCase string to a kebab-case string
 * @param str The string to convert
@@ -48,12 +52,14 @@ function camelCaseToKebabCase(str) {
 }
 
 //#endregion
-//#region ../agents/dist/do-oauth-client-provider-Cs9QpXYp.js
+//#region ../agents/dist/do-oauth-client-provider-B1fVIshX.js
+const STATE_EXPIRATION_MS = 600 * 1e3;
 var DurableObjectOAuthClientProvider = class {
 	constructor(storage, clientName, baseRedirectUrl) {
 		this.storage = storage;
 		this.clientName = clientName;
 		this.baseRedirectUrl = baseRedirectUrl;
+		if (!storage) throw new Error("DurableObjectOAuthClientProvider requires a valid DurableObjectStorage instance");
 	}
 	get clientMetadata() {
 		return {
@@ -69,7 +75,7 @@ var DurableObjectOAuthClientProvider = class {
 		return new URL(this.redirectUrl).origin;
 	}
 	get redirectUrl() {
-		return `${this.baseRedirectUrl}/${this.serverId}`;
+		return this.baseRedirectUrl;
 	}
 	get clientId() {
 		if (!this._clientId_) throw new Error("Trying to access clientId before it was set");
@@ -112,15 +118,72 @@ var DurableObjectOAuthClientProvider = class {
 	get authUrl() {
 		return this._authUrl_;
 	}
-	/**
-	* Because this operates on the server side (but we need browser auth), we send this url back to the user
-	* and require user interact to initiate the redirect flow
-	*/
-	redirectToAuthorization(authUrl) {
-		const stateToken = nanoid();
-		authUrl.searchParams.set("state", stateToken);
+	stateKey(nonce) {
+		return `/${this.clientName}/${this.serverId}/state/${nonce}`;
+	}
+	async state() {
+		const nonce = nanoid();
+		const state = `${nonce}.${this.serverId}`;
+		const storedState = {
+			nonce,
+			serverId: this.serverId,
+			createdAt: Date.now()
+		};
+		await this.storage.put(this.stateKey(nonce), storedState);
+		return state;
+	}
+	async checkState(state) {
+		const parts = state.split(".");
+		if (parts.length !== 2) return {
+			valid: false,
+			error: "Invalid state format"
+		};
+		const [nonce, serverId] = parts;
+		const key = this.stateKey(nonce);
+		const storedState = await this.storage.get(key);
+		if (!storedState) return {
+			valid: false,
+			error: "State not found or already used"
+		};
+		if (storedState.serverId !== serverId) {
+			await this.storage.delete(key);
+			return {
+				valid: false,
+				error: "State serverId mismatch"
+			};
+		}
+		if (Date.now() - storedState.createdAt > STATE_EXPIRATION_MS) {
+			await this.storage.delete(key);
+			return {
+				valid: false,
+				error: "State expired"
+			};
+		}
+		return {
+			valid: true,
+			serverId
+		};
+	}
+	async consumeState(state) {
+		const parts = state.split(".");
+		if (parts.length !== 2) {
+			console.warn(`[OAuth] consumeState called with invalid state format: ${state.substring(0, 20)}...`);
+			return;
+		}
+		const [nonce] = parts;
+		await this.storage.delete(this.stateKey(nonce));
+	}
+	async redirectToAuthorization(authUrl) {
 		this._authUrl_ = authUrl.toString();
 	}
+	async invalidateCredentials(scope) {
+		if (!this._clientId_) return;
+		const deleteKeys = [];
+		if (scope === "all" || scope === "client") deleteKeys.push(this.clientInfoKey(this.clientId));
+		if (scope === "all" || scope === "tokens") deleteKeys.push(this.tokenKey(this.clientId));
+		if (scope === "all" || scope === "verifier") deleteKeys.push(this.codeVerifierKey(this.clientId));
+		if (deleteKeys.length > 0) await this.storage.delete(deleteKeys);
+	}
 	codeVerifierKey(clientId) {
 		return `${this.keyPrefix(clientId)}/code_verifier`;
 	}
@@ -134,10 +197,13 @@ var DurableObjectOAuthClientProvider = class {
 		if (!codeVerifier) throw new Error("No code verifier found");
 		return codeVerifier;
 	}
+	async deleteCodeVerifier() {
+		await this.storage.delete(this.codeVerifierKey(this.clientId));
+	}
 };
 
 //#endregion
-//#region ../agents/dist/client-BZ-xTxF5.js
+//#region ../agents/dist/client-QZa2Rq0l.js
 function toDisposable(fn) {
 	return { dispose: fn };
 }
@@ -185,6 +251,22 @@ function isTransportNotImplemented(error) {
 	const msg = toErrorMessage(error);
 	return msg.includes("404") || msg.includes("405") || msg.includes("Not Implemented") || msg.includes("not implemented");
 }
+/**
+* Connection state machine for MCP client connections.
+*
+* State transitions:
+* - Non-OAuth: init() → CONNECTING → DISCOVERING → READY
+* - OAuth: init() → AUTHENTICATING → (callback) → CONNECTING → DISCOVERING → READY
+* - Any state can transition to FAILED on error
+*/
+const MCPConnectionState = {
+	AUTHENTICATING: "authenticating",
+	CONNECTING: "connecting",
+	CONNECTED: "connected",
+	DISCOVERING: "discovering",
+	READY: "ready",
+	FAILED: "failed"
+};
 var MCPClientConnection = class {
 	constructor(url, info, options = {
 		client: {},
@@ -192,7 +274,7 @@ var MCPClientConnection = class {
 	}) {
 		this.url = url;
 		this.options = options;
-		this.connectionState = "connecting";
+		this.connectionState = MCPConnectionState.CONNECTING;
 		this.tools = [];
 		this.prompts = [];
 		this.resources = [];
@@ -208,36 +290,49 @@ var MCPClientConnection = class {
 		});
 	}
 	/**
-	* Initialize a client connection
+	* Initialize a client connection, if authentication is required, the connection will be in the AUTHENTICATING state
+	* Sets connection state based on the result and emits observability events
 	*
-	* @returns
+	* @returns Error message if connection failed, undefined otherwise
 	*/
 	async init() {
 		const transportType = this.options.transport.type;
 		if (!transportType) throw new Error("Transport type must be specified");
-		try {
-			await this.tryConnect(transportType);
-		} catch (e) {
-			if (isUnauthorized(e)) {
-				this.connectionState = "authenticating";
-				return;
-			}
+		const res = await this.tryConnect(transportType);
+		this.connectionState = res.state;
+		if (res.state === MCPConnectionState.CONNECTED && res.transport) {
+			this.client.setRequestHandler(ElicitRequestSchema, async (request) => {
+				return await this.handleElicitationRequest(request);
+			});
+			this.lastConnectedTransport = res.transport;
+			this._onObservabilityEvent.fire({
+				type: "mcp:client:connect",
+				displayMessage: `Connected successfully using ${res.transport} transport for ${this.url.toString()}`,
+				payload: {
+					url: this.url.toString(),
+					transport: res.transport,
+					state: this.connectionState
+				},
+				timestamp: Date.now(),
+				id: nanoid()
+			});
+			return;
+		} else if (res.state === MCPConnectionState.FAILED && res.error) {
+			const errorMessage = toErrorMessage(res.error);
 			this._onObservabilityEvent.fire({
 				type: "mcp:client:connect",
-				displayMessage: `Connection initialization failed for ${this.url.toString()}`,
+				displayMessage: `Failed to connect to ${this.url.toString()}: ${errorMessage}`,
 				payload: {
 					url: this.url.toString(),
 					transport: transportType,
 					state: this.connectionState,
-					error: toErrorMessage(e)
+					error: errorMessage
 				},
 				timestamp: Date.now(),
 				id: nanoid()
 			});
-			this.connectionState = "failed";
-			return;
+			return errorMessage;
 		}
-		await this.discoverAndRegister();
 	}
 	/**
 	* Finish OAuth by probing transports based on configured type.
@@ -269,113 +364,189 @@ var MCPClientConnection = class {
 	* Complete OAuth authorization
 	*/
 	async completeAuthorization(code) {
-		if (this.connectionState !== "authenticating") throw new Error("Connection must be in authenticating state to complete authorization");
+		if (this.connectionState !== MCPConnectionState.AUTHENTICATING) throw new Error("Connection must be in authenticating state to complete authorization");
 		try {
 			await this.finishAuthProbe(code);
-			this.connectionState = "connecting";
+			this.connectionState = MCPConnectionState.CONNECTING;
 		} catch (error) {
-			this.connectionState = "failed";
+			this.connectionState = MCPConnectionState.FAILED;
 			throw error;
 		}
 	}
 	/**
-	* Establish connection after successful authorization
+	* Discover server capabilities and register tools, resources, prompts, and templates.
+	* This method does the work but does not manage connection state - that's handled by discover().
 	*/
-	async establishConnection() {
-		if (this.connectionState !== "connecting") throw new Error("Connection must be in connecting state to establish connection");
+	async discoverAndRegister() {
+		this.serverCapabilities = this.client.getServerCapabilities();
+		if (!this.serverCapabilities) throw new Error("The MCP Server failed to return server capabilities");
+		const operations = [];
+		const operationNames = [];
+		operations.push(Promise.resolve(this.client.getInstructions()));
+		operationNames.push("instructions");
+		if (this.serverCapabilities.tools) {
+			operations.push(this.registerTools());
+			operationNames.push("tools");
+		}
+		if (this.serverCapabilities.resources) {
+			operations.push(this.registerResources());
+			operationNames.push("resources");
+		}
+		if (this.serverCapabilities.prompts) {
+			operations.push(this.registerPrompts());
+			operationNames.push("prompts");
+		}
+		if (this.serverCapabilities.resources) {
+			operations.push(this.registerResourceTemplates());
+			operationNames.push("resource templates");
+		}
 		try {
-			const transportType = this.options.transport.type;
-			if (!transportType) throw new Error("Transport type must be specified");
-			await this.tryConnect(transportType);
-			await this.discoverAndRegister();
+			const results = await Promise.all(operations);
+			for (let i = 0; i < results.length; i++) {
+				const result = results[i];
+				switch (operationNames[i]) {
+					case "instructions":
+						this.instructions = result;
+						break;
+					case "tools":
+						this.tools = result;
+						break;
+					case "resources":
+						this.resources = result;
+						break;
+					case "prompts":
+						this.prompts = result;
+						break;
+					case "resource templates":
+						this.resourceTemplates = result;
+						break;
+				}
+			}
 		} catch (error) {
-			this.connectionState = "failed";
+			this._onObservabilityEvent.fire({
+				type: "mcp:client:discover",
+				displayMessage: `Failed to discover capabilities for ${this.url.toString()}: ${toErrorMessage(error)}`,
+				payload: {
+					url: this.url.toString(),
+					error: toErrorMessage(error)
+				},
+				timestamp: Date.now(),
+				id: nanoid()
+			});
 			throw error;
 		}
 	}
 	/**
-	* Discover server capabilities and register tools, resources, prompts, and templates
+	* Discover server capabilities with timeout and cancellation support.
+	* If called while a previous discovery is in-flight, the previous discovery will be aborted.
+	*
+	* @param options Optional configuration
+	* @param options.timeoutMs Timeout in milliseconds (default: 15000)
+	* @returns Result indicating success/failure with optional error message
 	*/
-	async discoverAndRegister() {
-		this.connectionState = "discovering";
-		this.serverCapabilities = this.client.getServerCapabilities();
-		if (!this.serverCapabilities) throw new Error("The MCP Server failed to return server capabilities");
-		const [instructionsResult, toolsResult, resourcesResult, promptsResult, resourceTemplatesResult] = await Promise.allSettled([
-			this.client.getInstructions(),
-			this.registerTools(),
-			this.registerResources(),
-			this.registerPrompts(),
-			this.registerResourceTemplates()
-		]);
-		const operations = [
-			{
-				name: "instructions",
-				result: instructionsResult
-			},
-			{
-				name: "tools",
-				result: toolsResult
-			},
-			{
-				name: "resources",
-				result: resourcesResult
-			},
-			{
-				name: "prompts",
-				result: promptsResult
-			},
-			{
-				name: "resource templates",
-				result: resourceTemplatesResult
-			}
-		];
-		for (const { name, result } of operations) if (result.status === "rejected") {
-			const url = this.url.toString();
+	async discover(options = {}) {
+		const { timeoutMs = 15e3 } = options;
+		if (this.connectionState !== MCPConnectionState.CONNECTED && this.connectionState !== MCPConnectionState.READY) {
 			this._onObservabilityEvent.fire({
 				type: "mcp:client:discover",
-				displayMessage: `Failed to discover ${name} for ${url}`,
+				displayMessage: `Discovery skipped for ${this.url.toString()}, state is ${this.connectionState}`,
 				payload: {
-					url,
-					capability: name,
-					error: result.reason
+					url: this.url.toString(),
+					state: this.connectionState
 				},
 				timestamp: Date.now(),
 				id: nanoid()
 			});
+			return {
+				success: false,
+				error: `Discovery skipped - connection in ${this.connectionState} state`
+			};
+		}
+		if (this._discoveryAbortController) {
+			this._discoveryAbortController.abort();
+			this._discoveryAbortController = void 0;
+		}
+		const abortController = new AbortController();
+		this._discoveryAbortController = abortController;
+		this.connectionState = MCPConnectionState.DISCOVERING;
+		let timeoutId;
+		try {
+			const timeoutPromise = new Promise((_, reject) => {
+				timeoutId = setTimeout(() => reject(/* @__PURE__ */ new Error(`Discovery timed out after ${timeoutMs}ms`)), timeoutMs);
+			});
+			if (abortController.signal.aborted) throw new Error("Discovery was cancelled");
+			const abortPromise = new Promise((_, reject) => {
+				abortController.signal.addEventListener("abort", () => {
+					reject(/* @__PURE__ */ new Error("Discovery was cancelled"));
+				});
+			});
+			await Promise.race([
+				this.discoverAndRegister(),
+				timeoutPromise,
+				abortPromise
+			]);
+			if (timeoutId !== void 0) clearTimeout(timeoutId);
+			this.connectionState = MCPConnectionState.READY;
+			this._onObservabilityEvent.fire({
+				type: "mcp:client:discover",
+				displayMessage: `Discovery completed for ${this.url.toString()}`,
+				payload: { url: this.url.toString() },
+				timestamp: Date.now(),
+				id: nanoid()
+			});
+			return { success: true };
+		} catch (e) {
+			if (timeoutId !== void 0) clearTimeout(timeoutId);
+			this.connectionState = MCPConnectionState.CONNECTED;
+			return {
+				success: false,
+				error: e instanceof Error ? e.message : String(e)
+			};
+		} finally {
+			this._discoveryAbortController = void 0;
+		}
+	}
+	/**
+	* Cancel any in-flight discovery operation.
+	* Called when closing the connection.
+	*/
+	cancelDiscovery() {
+		if (this._discoveryAbortController) {
+			this._discoveryAbortController.abort();
+			this._discoveryAbortController = void 0;
 		}
-		this.instructions = instructionsResult.status === "fulfilled" ? instructionsResult.value : void 0;
-		this.tools = toolsResult.status === "fulfilled" ? toolsResult.value : [];
-		this.resources = resourcesResult.status === "fulfilled" ? resourcesResult.value : [];
-		this.prompts = promptsResult.status === "fulfilled" ? promptsResult.value : [];
-		this.resourceTemplates = resourceTemplatesResult.status === "fulfilled" ? resourceTemplatesResult.value : [];
-		this.connectionState = "ready";
 	}
 	/**
-	* Notification handler registration
+	* Notification handler registration for tools
+	* Should only be called if serverCapabilities.tools exists
 	*/
 	async registerTools() {
-		if (!this.serverCapabilities || !this.serverCapabilities.tools) return [];
-		if (this.serverCapabilities.tools.listChanged) this.client.setNotificationHandler(ToolListChangedNotificationSchema, async (_notification) => {
+		if (this.serverCapabilities?.tools?.listChanged) this.client.setNotificationHandler(ToolListChangedNotificationSchema, async (_notification) => {
 			this.tools = await this.fetchTools();
 		});
 		return this.fetchTools();
 	}
+	/**
+	* Notification handler registration for resources
+	* Should only be called if serverCapabilities.resources exists
+	*/
 	async registerResources() {
-		if (!this.serverCapabilities || !this.serverCapabilities.resources) return [];
-		if (this.serverCapabilities.resources.listChanged) this.client.setNotificationHandler(ResourceListChangedNotificationSchema, async (_notification) => {
+		if (this.serverCapabilities?.resources?.listChanged) this.client.setNotificationHandler(ResourceListChangedNotificationSchema, async (_notification) => {
 			this.resources = await this.fetchResources();
 		});
 		return this.fetchResources();
 	}
+	/**
+	* Notification handler registration for prompts
+	* Should only be called if serverCapabilities.prompts exists
+	*/
 	async registerPrompts() {
-		if (!this.serverCapabilities || !this.serverCapabilities.prompts) return [];
-		if (this.serverCapabilities.prompts.listChanged) this.client.setNotificationHandler(PromptListChangedNotificationSchema, async (_notification) => {
+		if (this.serverCapabilities?.prompts?.listChanged) this.client.setNotificationHandler(PromptListChangedNotificationSchema, async (_notification) => {
 			this.prompts = await this.fetchPrompts();
 		});
 		return this.fetchPrompts();
 	}
 	async registerResourceTemplates() {
-		if (!this.serverCapabilities || !this.serverCapabilities.resources) return [];
 		return this.fetchResourceTemplates();
 	}
 	async fetchTools() {
@@ -441,44 +612,24 @@ var MCPClientConnection = class {
 			const transport = this.getTransport(currentTransportType);
 			try {
 				await this.client.connect(transport);
-				this.lastConnectedTransport = currentTransportType;
-				const url = this.url.toString();
-				this._onObservabilityEvent.fire({
-					type: "mcp:client:connect",
-					displayMessage: `Connected successfully using ${currentTransportType} transport for ${url}`,
-					payload: {
-						url,
-						transport: currentTransportType,
-						state: this.connectionState
-					},
-					timestamp: Date.now(),
-					id: nanoid()
-				});
-				break;
+				return {
+					state: MCPConnectionState.CONNECTED,
+					transport: currentTransportType
+				};
 			} catch (e) {
 				const error = e instanceof Error ? e : new Error(String(e));
-				if (isUnauthorized(error)) throw e;
-				if (hasFallback && isTransportNotImplemented(error)) {
-					const url = this.url.toString();
-					this._onObservabilityEvent.fire({
-						type: "mcp:client:connect",
-						displayMessage: `${currentTransportType} transport not available, trying ${transports[transports.indexOf(currentTransportType) + 1]} for ${url}`,
-						payload: {
-							url,
-							transport: currentTransportType,
-							state: this.connectionState
-						},
-						timestamp: Date.now(),
-						id: nanoid()
-					});
-					continue;
-				}
-				throw e;
+				if (isUnauthorized(error)) return { state: MCPConnectionState.AUTHENTICATING };
+				if (isTransportNotImplemented(error) && hasFallback) continue;
+				return {
+					state: MCPConnectionState.FAILED,
+					error
+				};
 			}
 		}
-		this.client.setRequestHandler(ElicitRequestSchema, async (request) => {
-			return await this.handleElicitationRequest(request);
-		});
+		return {
+			state: MCPConnectionState.FAILED,
+			error: /* @__PURE__ */ new Error("No transports available")
+		};
 	}
 	_capabilityErrorHandler(empty, method) {
 		return (e) => {
@@ -522,13 +673,32 @@ var MCPClientManager = class {
 		this.onObservabilityEvent = this._onObservabilityEvent.event;
 		this._onServerStateChanged = new Emitter();
 		this.onServerStateChanged = this._onServerStateChanged.event;
+		if (!options.storage) throw new Error("MCPClientManager requires a valid DurableObjectStorage instance");
 		this._storage = options.storage;
 	}
+	sql(query, ...bindings) {
+		return [...this._storage.sql.exec(query, ...bindings)];
+	}
+	saveServerToStorage(server) {
+		this.sql(`INSERT OR REPLACE INTO cf_agents_mcp_servers (
+        id, name, server_url, client_id, auth_url, callback_url, server_options
+      ) VALUES (?, ?, ?, ?, ?, ?, ?)`, server.id, server.name, server.server_url, server.client_id ?? null, server.auth_url ?? null, server.callback_url, server.server_options ?? null);
+	}
+	removeServerFromStorage(serverId) {
+		this.sql("DELETE FROM cf_agents_mcp_servers WHERE id = ?", serverId);
+	}
+	getServersFromStorage() {
+		return this.sql("SELECT id, name, server_url, client_id, auth_url, callback_url, server_options FROM cf_agents_mcp_servers");
+	}
+	clearServerAuthUrl(serverId) {
+		this.sql("UPDATE cf_agents_mcp_servers SET auth_url = NULL WHERE id = ?", serverId);
+	}
 	/**
 	* Create an auth provider for a server
 	* @internal
 	*/
 	createAuthProvider(serverId, callbackUrl, clientName, clientId) {
+		if (!this._storage) throw new Error("Cannot create auth provider: storage is not initialized");
 		const authProvider = new DurableObjectOAuthClientProvider(this._storage, clientName, callbackUrl);
 		authProvider.serverId = serverId;
 		if (clientId) authProvider.clientId = clientId;
@@ -542,7 +712,7 @@ var MCPClientManager = class {
 	*/
 	async restoreConnectionsFromStorage(clientName) {
 		if (this._isRestored) return;
-		const servers = await this._storage.listServers();
+		const servers = this.getServersFromStorage();
 		if (!servers || servers.length === 0) {
 			this._isRestored = true;
 			return;
@@ -550,12 +720,12 @@ var MCPClientManager = class {
 		for (const server of servers) {
 			const existingConn = this.mcpConnections[server.id];
 			if (existingConn) {
-				if (existingConn.connectionState === "ready") {
+				if (existingConn.connectionState === MCPConnectionState.READY) {
 					console.warn(`[MCPClientManager] Server ${server.id} already has a ready connection. Skipping recreation.`);
 					continue;
 				}
-				if (existingConn.connectionState === "authenticating" || existingConn.connectionState === "connecting" || existingConn.connectionState === "discovering") continue;
-				if (existingConn.connectionState === "failed") {
+				if (existingConn.connectionState === MCPConnectionState.AUTHENTICATING || existingConn.connectionState === MCPConnectionState.CONNECTING || existingConn.connectionState === MCPConnectionState.DISCOVERING) continue;
+				if (existingConn.connectionState === MCPConnectionState.FAILED) {
 					try {
 						await existingConn.client.close();
 					} catch (error) {
@@ -568,7 +738,7 @@ var MCPClientManager = class {
 			}
 			const parsedOptions = server.server_options ? JSON.parse(server.server_options) : null;
 			const authProvider = this.createAuthProvider(server.id, server.callback_url, clientName, server.client_id ?? void 0);
-			this.createConnection(server.id, server.server_url, {
+			const conn = this.createConnection(server.id, server.server_url, {
 				client: parsedOptions?.client ?? {},
 				transport: {
 					...parsedOptions?.transport ?? {},
@@ -576,13 +746,27 @@ var MCPClientManager = class {
 					authProvider
 				}
 			});
-			await this.connectToServer(server.id).catch((error) => {
-				console.error(`Error restoring ${server.id}:`, error);
-			});
+			if (server.auth_url) {
+				conn.connectionState = MCPConnectionState.AUTHENTICATING;
+				continue;
+			}
+			this._restoreServer(server.id);
 		}
 		this._isRestored = true;
 	}
 	/**
+	* Internal method to restore a single server connection and discovery
+	*/
+	async _restoreServer(serverId) {
+		if ((await this.connectToServer(serverId).catch((error) => {
+			console.error(`Error connecting to ${serverId}:`, error);
+			return null;
+		}))?.state === MCPConnectionState.CONNECTED) {
+			const discoverResult = await this.discoverIfConnected(serverId);
+			if (discoverResult && !discoverResult.success) console.error(`Error discovering ${serverId}:`, discoverResult.error);
+		}
+	}
+	/**
 	* Connect to and register an MCP server
 	*
 	* @deprecated This method is maintained for backward compatibility.
@@ -629,7 +813,7 @@ var MCPClientManager = class {
 		await this.mcpConnections[id].init();
 		if (options.reconnect?.oauthCode) try {
 			await this.mcpConnections[id].completeAuthorization(options.reconnect.oauthCode);
-			await this.mcpConnections[id].establishConnection();
+			await this.mcpConnections[id].init();
 		} catch (error) {
 			this._onObservabilityEvent.fire({
 				type: "mcp:client:connect",
@@ -646,19 +830,22 @@ var MCPClientManager = class {
 			throw error;
 		}
 		const authUrl = options.transport?.authProvider?.authUrl;
-		if (this.mcpConnections[id].connectionState === "authenticating" && authUrl && options.transport?.authProvider?.redirectUrl) return {
+		if (this.mcpConnections[id].connectionState === MCPConnectionState.AUTHENTICATING && authUrl && options.transport?.authProvider?.redirectUrl) return {
 			authUrl,
 			clientId: options.transport?.authProvider?.clientId,
 			id
 		};
+		const discoverResult = await this.discoverIfConnected(id);
+		if (discoverResult && !discoverResult.success) throw new Error(`Failed to discover server capabilities: ${discoverResult.error}`);
 		return { id };
 	}
 	/**
 	* Create an in-memory connection object and set up observability
 	* Does NOT save to storage - use registerServer() for that
+	* @returns The connection object (existing or newly created)
 	*/
 	createConnection(id, url, options) {
-		if (this.mcpConnections[id]) return;
+		if (this.mcpConnections[id]) return this.mcpConnections[id];
 		const normalizedTransport = {
 			...options.transport,
 			type: options.transport?.type ?? "auto"
@@ -680,6 +867,7 @@ var MCPClientManager = class {
 		store.add(this.mcpConnections[id].onObservabilityEvent((event) => {
 			this._onObservabilityEvent.fire(event);
 		}));
+		return this.mcpConnections[id];
 	}
 	/**
 	* Register an MCP server connection without connecting
@@ -697,7 +885,8 @@ var MCPClientManager = class {
 				type: options.transport?.type ?? "auto"
 			}
 		});
-		await this._storage.saveServer({
+		const { authProvider: _, ...transportWithoutAuth } = options.transport ?? {};
+		this.saveServerToStorage({
 			id,
 			name: options.name,
 			server_url: options.url,
@@ -706,7 +895,7 @@ var MCPClientManager = class {
 			auth_url: options.authUrl ?? null,
 			server_options: JSON.stringify({
 				client: options.client,
-				transport: options.transport
+				transport: transportWithoutAuth
 			})
 		});
 		this._onServerStateChanged.fire();
@@ -715,14 +904,13 @@ var MCPClientManager = class {
 	/**
 	* Connect to an already registered MCP server and initialize the connection.
 	*
-	* For OAuth servers, this returns `{ state: "authenticating", authUrl, clientId? }`
-	* without establishing the connection. The user must complete the OAuth flow via
-	* the authUrl, which will trigger a callback handled by `handleCallbackRequest()`.
-	*
-	* For non-OAuth servers, this establishes the connection immediately and returns
-	* `{ state: "ready" }`.
+	* For OAuth servers, returns `{ state: "authenticating", authUrl, clientId? }`.
+	* The user must complete the OAuth flow via the authUrl, which triggers a
+	* callback handled by `handleCallbackRequest()`.
 	*
-	* Updates storage with auth URL and client ID after connection.
+	* For non-OAuth servers, establishes the transport connection and returns
+	* `{ state: "connected" }`. Call `discoverIfConnected()` afterwards to
+	* discover capabilities and transition to "ready" state.
 	*
 	* @param id Server ID (must be registered first via registerServer())
 	* @returns Connection result with current state and OAuth info (if applicable)
@@ -730,70 +918,99 @@ var MCPClientManager = class {
 	async connectToServer(id) {
 		const conn = this.mcpConnections[id];
 		if (!conn) throw new Error(`Server ${id} is not registered. Call registerServer() first.`);
-		await conn.init();
-		const authUrl = conn.options.transport.authProvider?.authUrl;
-		if (conn.connectionState === "authenticating" && authUrl && conn.options.transport.authProvider?.redirectUrl) {
-			const clientId = conn.options.transport.authProvider?.clientId;
-			const serverRow = (await this._storage.listServers()).find((s) => s.id === id);
-			if (serverRow) await this._storage.saveServer({
-				...serverRow,
-				auth_url: authUrl,
-				client_id: clientId ?? null
-			});
-			this._onServerStateChanged.fire();
-			return {
-				state: "authenticating",
-				authUrl,
-				clientId
+		const error = await conn.init();
+		this._onServerStateChanged.fire();
+		switch (conn.connectionState) {
+			case MCPConnectionState.FAILED: return {
+				state: conn.connectionState,
+				error: error ?? "Unknown connection error"
+			};
+			case MCPConnectionState.AUTHENTICATING: {
+				const authUrl = conn.options.transport.authProvider?.authUrl;
+				const redirectUrl = conn.options.transport.authProvider?.redirectUrl;
+				if (!authUrl || !redirectUrl) return {
+					state: MCPConnectionState.FAILED,
+					error: `OAuth configuration incomplete: missing ${!authUrl ? "authUrl" : "redirectUrl"}`
+				};
+				const clientId = conn.options.transport.authProvider?.clientId;
+				const serverRow = this.getServersFromStorage().find((s) => s.id === id);
+				if (serverRow) {
+					this.saveServerToStorage({
+						...serverRow,
+						auth_url: authUrl,
+						client_id: clientId ?? null
+					});
+					this._onServerStateChanged.fire();
+				}
+				return {
+					state: conn.connectionState,
+					authUrl,
+					clientId
+				};
+			}
+			case MCPConnectionState.CONNECTED: return { state: conn.connectionState };
+			default: return {
+				state: MCPConnectionState.FAILED,
+				error: `Unexpected connection state after init: ${conn.connectionState}`
 			};
 		}
-		if (conn.connectionState === "ready") this._onServerStateChanged.fire();
-		return { state: "ready" };
 	}
-	async isCallbackRequest(req) {
+	extractServerIdFromState(state) {
+		if (!state) return null;
+		const parts = state.split(".");
+		return parts.length === 2 ? parts[1] : null;
+	}
+	isCallbackRequest(req) {
 		if (req.method !== "GET") return false;
 		if (!req.url.includes("/callback")) return false;
-		return (await this._storage.listServers()).some((server) => server.callback_url && req.url.startsWith(server.callback_url));
+		const state = new URL(req.url).searchParams.get("state");
+		const serverId = this.extractServerIdFromState(state);
+		if (!serverId) return false;
+		return this.getServersFromStorage().some((server) => server.id === serverId);
 	}
 	async handleCallbackRequest(req) {
 		const url = new URL(req.url);
-		const matchingServer = (await this._storage.listServers()).find((server) => {
-			return server.callback_url && req.url.startsWith(server.callback_url);
-		});
-		if (!matchingServer) throw new Error(`No callback URI match found for the request url: ${req.url}. Was the request matched with \`isCallbackRequest()\`?`);
-		const serverId = matchingServer.id;
 		const code = url.searchParams.get("code");
 		const state = url.searchParams.get("state");
 		const error = url.searchParams.get("error");
 		const errorDescription = url.searchParams.get("error_description");
+		if (!state) throw new Error("Unauthorized: no state provided");
+		const serverId = this.extractServerIdFromState(state);
+		if (!serverId) throw new Error("No serverId found in state parameter. Expected format: {nonce}.{serverId}");
+		if (!this.getServersFromStorage().some((server) => server.id === serverId)) throw new Error(`No server found with id "${serverId}". Was the request matched with \`isCallbackRequest()\`?`);
+		if (this.mcpConnections[serverId] === void 0) throw new Error(`Could not find serverId: ${serverId}`);
+		const conn = this.mcpConnections[serverId];
+		if (!conn.options.transport.authProvider) throw new Error("Trying to finalize authentication for a server connection without an authProvider");
+		const authProvider = conn.options.transport.authProvider;
+		authProvider.serverId = serverId;
+		const stateValidation = await authProvider.checkState(state);
+		if (!stateValidation.valid) throw new Error(`Invalid state: ${stateValidation.error}`);
 		if (error) return {
 			serverId,
 			authSuccess: false,
 			authError: errorDescription || error
 		};
 		if (!code) throw new Error("Unauthorized: no code provided");
-		if (!state) throw new Error("Unauthorized: no state provided");
-		if (this.mcpConnections[serverId] === void 0) throw new Error(`Could not find serverId: ${serverId}`);
-		if (this.mcpConnections[serverId].connectionState === "ready") return {
-			serverId,
-			authSuccess: true
-		};
-		if (this.mcpConnections[serverId].connectionState !== "authenticating") throw new Error(`Failed to authenticate: the client is in "${this.mcpConnections[serverId].connectionState}" state, expected "authenticating"`);
-		const conn = this.mcpConnections[serverId];
-		if (!conn.options.transport.authProvider) throw new Error("Trying to finalize authentication for a server connection without an authProvider");
-		const clientId = conn.options.transport.authProvider.clientId || state;
-		conn.options.transport.authProvider.clientId = clientId;
-		conn.options.transport.authProvider.serverId = serverId;
+		if (this.mcpConnections[serverId].connectionState === MCPConnectionState.READY || this.mcpConnections[serverId].connectionState === MCPConnectionState.CONNECTED) {
+			this.clearServerAuthUrl(serverId);
+			return {
+				serverId,
+				authSuccess: true
+			};
+		}
+		if (this.mcpConnections[serverId].connectionState !== MCPConnectionState.AUTHENTICATING) throw new Error(`Failed to authenticate: the client is in "${this.mcpConnections[serverId].connectionState}" state, expected "authenticating"`);
 		try {
+			await authProvider.consumeState(state);
 			await conn.completeAuthorization(code);
-			await this._storage.clearAuthUrl(serverId);
+			await authProvider.deleteCodeVerifier();
+			this.clearServerAuthUrl(serverId);
 			this._onServerStateChanged.fire();
 			return {
 				serverId,
 				authSuccess: true
 			};
-		} catch (error$1) {
-			const errorMessage = error$1 instanceof Error ? error$1.message : String(error$1);
+		} catch (authError) {
+			const errorMessage = authError instanceof Error ? authError.message : String(authError);
 			this._onServerStateChanged.fire();
 			return {
 				serverId,
@@ -803,8 +1020,40 @@ var MCPClientManager = class {
 		}
 	}
 	/**
+	* Discover server capabilities if connection is in CONNECTED or READY state.
+	* Transitions to DISCOVERING then READY (or CONNECTED on error).
+	* Can be called to refresh server capabilities (e.g., from a UI refresh button).
+	*
+	* If called while a previous discovery is in-flight for the same server,
+	* the previous discovery will be aborted.
+	*
+	* @param serverId The server ID to discover
+	* @param options Optional configuration
+	* @param options.timeoutMs Timeout in milliseconds (default: 30000)
+	* @returns Result with current state and optional error, or undefined if connection not found
+	*/
+	async discoverIfConnected(serverId, options = {}) {
+		const conn = this.mcpConnections[serverId];
+		if (!conn) {
+			this._onObservabilityEvent.fire({
+				type: "mcp:client:discover",
+				displayMessage: `Connection not found for ${serverId}`,
+				payload: {},
+				timestamp: Date.now(),
+				id: nanoid()
+			});
+			return;
+		}
+		const result = await conn.discover(options);
+		this._onServerStateChanged.fire();
+		return {
+			...result,
+			state: conn.connectionState
+		};
+	}
+	/**
 	* Establish connection in the background after OAuth completion
-	* This method is called asynchronously and doesn't block the OAuth callback response
+	* This method connects to the server and discovers its capabilities
 	* @param serverId The server ID to establish connection for
 	*/
 	async establishConnection(serverId) {
@@ -819,25 +1068,34 @@ var MCPClientManager = class {
 			});
 			return;
 		}
-		try {
-			await conn.establishConnection();
-			this._onServerStateChanged.fire();
-		} catch (error) {
-			const url = conn.url.toString();
+		if (conn.connectionState === MCPConnectionState.DISCOVERING || conn.connectionState === MCPConnectionState.READY) {
 			this._onObservabilityEvent.fire({
 				type: "mcp:client:connect",
-				displayMessage: `Failed to establish connection to server ${serverId} with url ${url}`,
+				displayMessage: `establishConnection skipped for ${serverId}, already in ${conn.connectionState} state`,
 				payload: {
-					url,
-					transport: conn.options.transport.type ?? "auto",
-					state: conn.connectionState,
-					error: toErrorMessage(error)
+					url: conn.url.toString(),
+					transport: conn.options.transport.type || "unknown",
+					state: conn.connectionState
 				},
 				timestamp: Date.now(),
 				id: nanoid()
 			});
-			this._onServerStateChanged.fire();
+			return;
 		}
+		const connectResult = await this.connectToServer(serverId);
+		this._onServerStateChanged.fire();
+		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()
+		});
 	}
 	/**
 	* Configure OAuth callback handling
@@ -881,7 +1139,7 @@ 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 !== "ready" && conn.connectionState !== "authenticating") console.warn(`[getAITools] WARNING: Reading tools from connection ${id} in state "${conn.connectionState}". Tools may not be loaded yet.`);
+		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}`, {
 				description: tool.description,
@@ -911,10 +1169,18 @@ var MCPClientManager = class {
 		return this.getAITools();
 	}
 	/**
-	* Closes all connections to MCP servers
+	* Closes all active in-memory connections to MCP servers.
+	*
+	* Note: This only closes the transport connections - it does NOT remove
+	* servers from storage. Servers will still be listed and their callback
+	* URLs will still match incoming OAuth requests.
+	*
+	* Use removeServer() instead if you want to fully clean up a server
+	* (closes connection AND removes from storage).
 	*/
 	async closeAllConnections() {
 		const ids = Object.keys(this.mcpConnections);
+		for (const id of ids) this.mcpConnections[id].cancelDiscovery();
 		await Promise.all(ids.map(async (id) => {
 			await this.mcpConnections[id].client.close();
 		}));
@@ -931,6 +1197,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();
 		await this.mcpConnections[id].client.close();
 		delete this.mcpConnections[id];
 		const store = this._connectionDisposables.get(id);
@@ -938,17 +1205,20 @@ var MCPClientManager = class {
 		this._connectionDisposables.delete(id);
 	}
 	/**
-	* Remove an MCP server from storage
+	* Remove an MCP server - closes connection if active and removes from storage.
 	*/
 	async removeServer(serverId) {
-		await this._storage.removeServer(serverId);
+		if (this.mcpConnections[serverId]) try {
+			await this.closeConnection(serverId);
+		} catch (_e) {}
+		this.removeServerFromStorage(serverId);
 		this._onServerStateChanged.fire();
 	}
 	/**
 	* List all MCP servers from storage
 	*/
-	async listServers() {
-		return await this._storage.listServers();
+	listServers() {
+		return this.getServersFromStorage();
 	}
 	/**
 	* Dispose the manager and all resources.
@@ -1019,72 +1289,7 @@ function getNamespacedData(mcpClients, type) {
 }
 
 //#endregion
-//#region ../agents/dist/src-C6rC6ZpH.js
-/**
-* SQL-based storage adapter that wraps SQL operations
-* Used by Agent class to provide SQL access to MCPClientManager
-*/
-var AgentMCPClientStorage = class {
-	constructor(sql, kv) {
-		this.sql = sql;
-		this.kv = kv;
-	}
-	async saveServer(server) {
-		this.sql`
-      INSERT OR REPLACE INTO cf_agents_mcp_servers (
-        id,
-        name,
-        server_url,
-        client_id,
-        auth_url,
-        callback_url,
-        server_options
-      )
-      VALUES (
-        ${server.id},
-        ${server.name},
-        ${server.server_url},
-        ${server.client_id ?? null},
-        ${server.auth_url ?? null},
-        ${server.callback_url},
-        ${server.server_options ?? null}
-      )
-    `;
-	}
-	async removeServer(serverId) {
-		this.sql`
-      DELETE FROM cf_agents_mcp_servers WHERE id = ${serverId}
-    `;
-	}
-	async listServers() {
-		return this.sql`
-      SELECT id, name, server_url, client_id, auth_url, callback_url, server_options
-      FROM cf_agents_mcp_servers
-    `;
-	}
-	async getServerByCallbackUrl(callbackUrl) {
-		const results = this.sql`
-      SELECT id, name, server_url, client_id, auth_url, callback_url, server_options
-      FROM cf_agents_mcp_servers
-      WHERE callback_url = ${callbackUrl}
-      LIMIT 1
-    `;
-		return results.length > 0 ? results[0] : null;
-	}
-	async clearAuthUrl(serverId) {
-		this.sql`
-      UPDATE cf_agents_mcp_servers
-      SET auth_url = NULL
-      WHERE id = ${serverId}
-    `;
-	}
-	async get(key) {
-		return this.kv.get(key);
-	}
-	async put(key, value) {
-		return this.kv.put(key, value);
-	}
-};
+//#region ../agents/dist/src-BmbDclOA.js
 /**
 * A generic observability implementation that logs events to the console.
 */
@@ -1122,7 +1327,6 @@ function getNextCronTime(cron) {
 const STATE_ROW_ID = "cf_state_row_id";
 const STATE_WAS_CHANGED = "cf_state_was_changed";
 const DEFAULT_STATE = {};
-const agentContext = new AsyncLocalStorage();
 function getCurrentAgent() {
 	const store = agentContext.getStore();
 	if (!store) return {
@@ -1258,17 +1462,10 @@ var Agent = class Agent$1 extends Server {
 			if (this._destroyed) return;
 			await this._scheduleNextAlarm();
 		};
-		this.mcp = new MCPClientManager(this._ParentClass.name, "0.0.1", { storage: new AgentMCPClientStorage(this.sql.bind(this), this.ctx.storage.kv) });
 		if (!wrappedClasses.has(this.constructor)) {
 			this._autoWrapCustomMethods();
 			wrappedClasses.add(this.constructor);
 		}
-		this._disposables.add(this.mcp.onServerStateChanged(async () => {
-			await this.broadcastMcpServers();
-		}));
-		this._disposables.add(this.mcp.onObservabilityEvent((event) => {
-			this.observability?.emit(event);
-		}));
 		this.sql`
         CREATE TABLE IF NOT EXISTS cf_agents_mcp_servers (
           id TEXT PRIMARY KEY NOT NULL,
@@ -1306,6 +1503,13 @@ var Agent = class Agent$1 extends Server {
         created_at INTEGER DEFAULT (unixepoch())
       )
     `;
+		this.mcp = new MCPClientManager(this._ParentClass.name, "0.0.1", { storage: this.ctx.storage });
+		this._disposables.add(this.mcp.onServerStateChanged(async () => {
+			this.broadcastMcpServers();
+		}));
+		this._disposables.add(this.mcp.onObservabilityEvent((event) => {
+			this.observability?.emit(event);
+		}));
 		const _onRequest = this.onRequest.bind(this);
 		this.onRequest = (request) => {
 			return agentContext.run({
@@ -1399,7 +1603,7 @@ var Agent = class Agent$1 extends Server {
 					type: MessageType.CF_AGENT_STATE
 				}));
 				connection.send(JSON.stringify({
-					mcp: await this.getMcpServers(),
+					mcp: this.getMcpServers(),
 					type: MessageType.CF_AGENT_MCP_SERVERS
 				}));
 				this.observability?.emit({
@@ -1422,7 +1626,7 @@ var Agent = class Agent$1 extends Server {
 			}, async () => {
 				await this._tryCatch(async () => {
 					await this.mcp.restoreConnectionsFromStorage(this.name);
-					await this.broadcastMcpServers();
+					this.broadcastMcpServers();
 					return _onStart(props);
 				});
 			});
@@ -1767,10 +1971,7 @@ var Agent = class Agent$1 extends Server {
 		const result = this.sql`
       SELECT * FROM cf_agents_schedules WHERE id = ${id}
     `;
-		if (!result) {
-			console.error(`schedule ${id} not found`);
-			return;
-		}
+		if (!result || result.length === 0) return;
 		return {
 			...result[0],
 			payload: JSON.parse(result[0].payload)
@@ -1807,11 +2008,12 @@ var Agent = class Agent$1 extends Server {
 	/**
 	* Cancel a scheduled task
 	* @param id ID of the task to cancel
-	* @returns true if the task was cancelled, false otherwise
+	* @returns true if the task was cancelled, false if the task was not found
 	*/
 	async cancelSchedule(id) {
 		const schedule = await this.getSchedule(id);
-		if (schedule) this.observability?.emit({
+		if (!schedule) return false;
+		this.observability?.emit({
 			displayMessage: `Schedule ${id} cancelled`,
 			id: nanoid(),
 			payload: {
@@ -1877,7 +2079,8 @@ var Agent = class Agent$1 extends Server {
 	* @param callbackHost Base host for the agent, used for the redirect URI. If not provided, will be derived from the current request.
 	* @param agentsPrefix agents routing prefix if not using `agents`
 	* @param options MCP client and transport options
-	* @returns authUrl
+	* @returns Server id and state - either "authenticating" with authUrl, or "ready"
+	* @throws If connection or discovery fails
 	*/
 	async addMcpServer(serverName, url, callbackHost, agentsPrefix = "agents", options) {
 		let resolvedCallbackHost = callbackHost;
@@ -1890,7 +2093,7 @@ var Agent = class Agent$1 extends Server {
 		const callbackUrl = `${resolvedCallbackHost}/${agentsPrefix}/${camelCaseToKebabCase(this._ParentClass.name)}/${this.name}/callback`;
 		await this.mcp.ensureJsonSchema();
 		const id = nanoid(8);
-		const authProvider = new DurableObjectOAuthClientProvider(this.ctx.storage.kv, this.name, callbackUrl);
+		const authProvider = new DurableObjectOAuthClientProvider(this.ctx.storage, this.name, callbackUrl);
 		authProvider.serverId = id;
 		const transportType = options?.transport?.type ?? "auto";
 		let headerTransportOpts = {};
@@ -1913,23 +2116,30 @@ var Agent = class Agent$1 extends Server {
 			}
 		});
 		const result = await this.mcp.connectToServer(id);
+		if (result.state === MCPConnectionState.FAILED) throw new Error(`Failed to connect to MCP server at ${url}: ${result.error}`);
+		if (result.state === MCPConnectionState.AUTHENTICATING) return {
+			id,
+			state: result.state,
+			authUrl: result.authUrl
+		};
+		const discoverResult = await this.mcp.discoverIfConnected(id);
+		if (discoverResult && !discoverResult.success) throw new Error(`Failed to discover MCP server capabilities: ${discoverResult.error}`);
 		return {
 			id,
-			authUrl: result.state === "authenticating" ? result.authUrl : void 0
+			state: MCPConnectionState.READY
 		};
 	}
 	async removeMcpServer(id) {
-		if (this.mcp.mcpConnections[id]) await this.mcp.closeConnection(id);
 		await this.mcp.removeServer(id);
 	}
-	async getMcpServers() {
+	getMcpServers() {
 		const mcpState = {
 			prompts: this.mcp.listPrompts(),
 			resources: this.mcp.listResources(),
 			servers: {},
 			tools: this.mcp.listTools()
 		};
-		const servers = await this.mcp.listServers();
+		const servers = this.mcp.listServers();
 		if (servers && Array.isArray(servers) && servers.length > 0) for (const server of servers) {
 			const serverConn = this.mcp.mcpConnections[server.id];
 			let defaultState = "not-connected";
@@ -1945,9 +2155,9 @@ var Agent = class Agent$1 extends Server {
 		}
 		return mcpState;
 	}
-	async broadcastMcpServers() {
+	broadcastMcpServers() {
 		this.broadcast(JSON.stringify({
-			mcp: await this.getMcpServers(),
+			mcp: this.getMcpServers(),
 			type: MessageType.CF_AGENT_MCP_SERVERS
 		}));
 	}
@@ -1965,16 +2175,12 @@ var Agent = class Agent$1 extends Server {
 	* @returns Response if this was an OAuth callback, null otherwise
 	*/
 	async handleMcpOAuthCallback(request) {
-		if (!await this.mcp.isCallbackRequest(request)) return null;
+		if (!this.mcp.isCallbackRequest(request)) return null;
 		const result = await this.mcp.handleCallbackRequest(request);
-		if (result.authSuccess) {
-			this.broadcastMcpServers();
-			this.mcp.establishConnection(result.serverId).catch((error) => {
-				console.error("[Agent handleMcpOAuthCallback] Connection establishment failed:", error);
-			}).finally(() => {
-				this.broadcastMcpServers();
-			});
-		}
+		if (result.authSuccess) this.mcp.establishConnection(result.serverId).catch((error) => {
+			console.error("[Agent handleMcpOAuthCallback] Connection establishment failed:", error);
+		});
+		this.broadcastMcpServers();
 		return this.handleOAuthCallbackResponse(result, request);
 	}
 	/**