agents 0.0.0-ccbea72 → 0.0.0-cccbd0f

package/dist/{src-Dz0H9hSU.js → src-C6rC6ZpH.js}

@@ -1,13 +1,81 @@
-import { t as MessageType } from "./ai-types-B3aQaFv3.js";
-import { r as camelCaseToKebabCase } from "./client-BfiZ3HQd.js";
-import { o as DisposableStore, r as MCPClientConnection, t as MCPClientManager } from "./client-9Ld2_lnt.js";
-import { t as DurableObjectOAuthClientProvider } from "./do-oauth-client-provider-CswoD5Lu.js";
+import { t as MessageType } from "./ai-types-CrMqkwc_.js";
+import { r as camelCaseToKebabCase } from "./client-B3SR12TQ.js";
+import { r as DisposableStore, t as MCPClientManager } from "./client-BZ-xTxF5.js";
+import { t as DurableObjectOAuthClientProvider } from "./do-oauth-client-provider-Cs9QpXYp.js";
 import { AsyncLocalStorage } from "node:async_hooks";
 import { parseCronExpression } from "cron-schedule";
 import { nanoid } from "nanoid";
 import { EmailMessage } from "cloudflare:email";
 import { Server, getServerByName, routePartykitRequest } from "partyserver";
 
+//#region src/mcp/client-storage.ts
+/**
+* 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);
+	}
+};
+
+//#endregion
 //#region src/observability/index.ts
 /**
 * A generic observability implementation that logs events to the console.
@@ -153,9 +221,8 @@ var Agent = class Agent extends Server {
 		super(ctx, env);
 		this._state = DEFAULT_STATE;
 		this._disposables = new DisposableStore();
-		this._mcpStateRestored = false;
+		this._destroyed = false;
 		this._ParentClass = Object.getPrototypeOf(this).constructor;
-		this.mcp = new MCPClientManager(this._ParentClass.name, "0.0.1");
 		this.initialState = DEFAULT_STATE;
 		this.observability = genericObservability;
 		this._flushingQueue = false;
@@ -193,28 +260,45 @@ var Agent = class Agent extends Server {
 					}
 				});
 				if (row.type === "cron") {
+					if (this._destroyed) return;
 					const nextExecutionTime = getNextCronTime(row.cron);
 					const nextTimestamp = Math.floor(nextExecutionTime.getTime() / 1e3);
 					this.sql`
           UPDATE cf_agents_schedules SET time = ${nextTimestamp} WHERE id = ${row.id}
         `;
-				} else this.sql`
+				} else {
+					if (this._destroyed) return;
+					this.sql`
           DELETE FROM cf_agents_schedules WHERE id = ${row.id}
         `;
+				}
 			}
+			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.onConnected(async () => {
-			this.broadcastMcpServers();
+		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,
+          name TEXT NOT NULL,
+          server_url TEXT NOT NULL,
+          callback_url TEXT NOT NULL,
+          client_id TEXT,
+          auth_url TEXT,
+          server_options TEXT
+        )
+      `;
+		this.sql`
       CREATE TABLE IF NOT EXISTS cf_agents_state (
         id TEXT PRIMARY KEY NOT NULL,
         state TEXT
@@ -228,32 +312,16 @@ var Agent = class Agent extends Server {
         created_at INTEGER DEFAULT (unixepoch())
       )
     `;
-		this.ctx.blockConcurrencyWhile(async () => {
-			return this._tryCatch(async () => {
-				this.sql`
-        CREATE TABLE IF NOT EXISTS cf_agents_schedules (
-          id TEXT PRIMARY KEY NOT NULL DEFAULT (randomblob(9)),
-          callback TEXT,
-          payload TEXT,
-          type TEXT NOT NULL CHECK(type IN ('scheduled', 'delayed', 'cron')),
-          time INTEGER,
-          delayInSeconds INTEGER,
-          cron TEXT,
-          created_at INTEGER DEFAULT (unixepoch())
-        )
-      `;
-				await this.alarm();
-			});
-		});
 		this.sql`
-      CREATE TABLE IF NOT EXISTS cf_agents_mcp_servers (
-        id TEXT PRIMARY KEY NOT NULL,
-        name TEXT NOT NULL,
-        server_url TEXT NOT NULL,
-        callback_url TEXT NOT NULL,
-        client_id TEXT,
-        auth_url TEXT,
-        server_options TEXT
+      CREATE TABLE IF NOT EXISTS cf_agents_schedules (
+        id TEXT PRIMARY KEY NOT NULL DEFAULT (randomblob(9)),
+        callback TEXT,
+        payload TEXT,
+        type TEXT NOT NULL CHECK(type IN ('scheduled', 'delayed', 'cron')),
+        time INTEGER,
+        delayInSeconds INTEGER,
+        cron TEXT,
+        created_at INTEGER DEFAULT (unixepoch())
       )
     `;
 		const _onRequest = this.onRequest.bind(this);
@@ -264,17 +332,9 @@ var Agent = class Agent extends Server {
 				request,
 				email: void 0
 			}, async () => {
-				await this._ensureMcpStateRestored();
-				if (this.mcp.isCallbackRequest(request)) {
-					const result = await this.mcp.handleCallbackRequest(request);
-					this.broadcastMcpServers();
-					if (result.authSuccess) this.mcp.establishConnection(result.serverId).catch((error) => {
-						console.error("Background connection failed:", error);
-					}).finally(() => {
-						this.broadcastMcpServers();
-					});
-					return this.handleOAuthCallbackResponse(result, request);
-				}
+				await this.mcp.ensureJsonSchema();
+				const oauthResponse = await this.handleMcpOAuthCallback(request);
+				if (oauthResponse) return oauthResponse;
 				return this._tryCatch(() => _onRequest(request));
 			});
 		};
@@ -286,6 +346,7 @@ var Agent = class Agent extends Server {
 				request: void 0,
 				email: void 0
 			}, async () => {
+				await this.mcp.ensureJsonSchema();
 				if (typeof message !== "string") return this._tryCatch(() => _onMessage(connection, message));
 				let parsed;
 				try {
@@ -350,13 +411,13 @@ var Agent = class Agent extends Server {
 				connection,
 				request: ctx$1.request,
 				email: void 0
-			}, () => {
+			}, async () => {
 				if (this.state) connection.send(JSON.stringify({
 					state: this.state,
 					type: MessageType.CF_AGENT_STATE
 				}));
 				connection.send(JSON.stringify({
-					mcp: this.getMcpServers(),
+					mcp: await this.getMcpServers(),
 					type: MessageType.CF_AGENT_MCP_SERVERS
 				}));
 				this.observability?.emit({
@@ -378,8 +439,8 @@ var Agent = class Agent extends Server {
 				email: void 0
 			}, async () => {
 				await this._tryCatch(async () => {
-					await this._ensureMcpStateRestored();
-					this.broadcastMcpServers();
+					await this.mcp.restoreConnectionsFromStorage(this.name);
+					await this.broadcastMcpServers();
 					return _onStart(props);
 				});
 			});
@@ -785,7 +846,7 @@ var Agent = class Agent extends Server {
 	async _scheduleNextAlarm() {
 		const result = this.sql`
       SELECT time FROM cf_agents_schedules
-      WHERE time > ${Math.floor(Date.now() / 1e3)}
+      WHERE time >= ${Math.floor(Date.now() / 1e3)}
       ORDER BY time ASC
       LIMIT 1
     `;
@@ -799,15 +860,18 @@ var Agent = class Agent extends Server {
 	* Destroy the Agent, removing all state and scheduled tasks
 	*/
 	async destroy() {
+		this.sql`DROP TABLE IF EXISTS cf_agents_mcp_servers`;
 		this.sql`DROP TABLE IF EXISTS cf_agents_state`;
 		this.sql`DROP TABLE IF EXISTS cf_agents_schedules`;
-		this.sql`DROP TABLE IF EXISTS cf_agents_mcp_servers`;
 		this.sql`DROP TABLE IF EXISTS cf_agents_queues`;
 		await this.ctx.storage.deleteAlarm();
 		await this.ctx.storage.deleteAll();
 		this._disposables.dispose();
-		await this.mcp.dispose?.();
-		this.ctx.abort("destroyed");
+		await this.mcp.dispose();
+		this._destroyed = true;
+		setTimeout(() => {
+			this.ctx.abort("destroyed");
+		}, 0);
 		this.observability?.emit({
 			displayMessage: "Agent destroyed",
 			id: nanoid(),
@@ -823,43 +887,6 @@ var Agent = class Agent extends Server {
 	_isCallable(method) {
 		return callableMetadata.has(this[method]);
 	}
-	async _ensureMcpStateRestored() {
-		if (this._mcpStateRestored) return;
-		this._mcpStateRestored = true;
-		const servers = this.sql`
-        SELECT id, name, server_url, client_id, auth_url, callback_url, server_options
-        FROM cf_agents_mcp_servers
-      `;
-		if (!servers || !Array.isArray(servers) || servers.length === 0) return;
-		for (const server of servers) if (server.callback_url) this.mcp.registerCallbackUrl(`${server.callback_url}/${server.id}`);
-		for (const server of servers) if (!!server.auth_url) {
-			const authProvider = new DurableObjectOAuthClientProvider(this.ctx.storage, this.name, server.callback_url);
-			authProvider.serverId = server.id;
-			if (server.client_id) authProvider.clientId = server.client_id;
-			const parsedOptions = server.server_options ? JSON.parse(server.server_options) : void 0;
-			const conn = new MCPClientConnection(new URL(server.server_url), {
-				name: this.name,
-				version: "1.0.0"
-			}, {
-				client: parsedOptions?.client ?? {},
-				transport: {
-					...parsedOptions?.transport ?? {},
-					type: parsedOptions?.transport?.type ?? "auto",
-					authProvider
-				}
-			});
-			conn.connectionState = "authenticating";
-			this.mcp.mcpConnections[server.id] = conn;
-		} else {
-			const parsedOptions = server.server_options ? JSON.parse(server.server_options) : void 0;
-			this._connectToMcpServerInternal(server.name, server.server_url, server.callback_url, parsedOptions, {
-				id: server.id,
-				oauthClientId: server.client_id ?? void 0
-			}).catch((error) => {
-				console.error(`Error restoring ${server.id}:`, error);
-			});
-		}
-	}
 	/**
 	* Connect to a new MCP Server
 	*
@@ -879,29 +906,10 @@ var Agent = class Agent extends Server {
 			resolvedCallbackHost = `${requestUrl.protocol}//${requestUrl.host}`;
 		}
 		const callbackUrl = `${resolvedCallbackHost}/${agentsPrefix}/${camelCaseToKebabCase(this._ParentClass.name)}/${this.name}/callback`;
-		const result = await this._connectToMcpServerInternal(serverName, url, callbackUrl, options);
-		this.sql`
-        INSERT
-        OR REPLACE INTO cf_agents_mcp_servers (id, name, server_url, client_id, auth_url, callback_url, server_options)
-      VALUES (
-        ${result.id},
-        ${serverName},
-        ${url},
-        ${result.clientId ?? null},
-        ${result.authUrl ?? null},
-        ${callbackUrl},
-        ${options ? JSON.stringify(options) : null}
-        );
-    `;
-		this.broadcastMcpServers();
-		return result;
-	}
-	async _connectToMcpServerInternal(_serverName, url, callbackUrl, options, reconnect) {
-		const authProvider = new DurableObjectOAuthClientProvider(this.ctx.storage, this.name, callbackUrl);
-		if (reconnect) {
-			authProvider.serverId = reconnect.id;
-			if (reconnect.oauthClientId) authProvider.clientId = reconnect.oauthClientId;
-		}
+		await this.mcp.ensureJsonSchema();
+		const id = nanoid(8);
+		const authProvider = new DurableObjectOAuthClientProvider(this.ctx.storage.kv, this.name, callbackUrl);
+		authProvider.serverId = id;
 		const transportType = options?.transport?.type ?? "auto";
 		let headerTransportOpts = {};
 		if (options?.transport?.headers) headerTransportOpts = {
@@ -911,59 +919,83 @@ var Agent = class Agent extends Server {
 			}) },
 			requestInit: { headers: options?.transport?.headers }
 		};
-		const { id, authUrl, clientId } = await this.mcp.connect(url, {
+		await this.mcp.registerServer(id, {
+			url,
+			name: serverName,
+			callbackUrl,
 			client: options?.client,
-			reconnect,
 			transport: {
 				...headerTransportOpts,
 				authProvider,
 				type: transportType
 			}
 		});
+		const result = await this.mcp.connectToServer(id);
 		return {
-			authUrl,
-			clientId,
-			id
+			id,
+			authUrl: result.state === "authenticating" ? result.authUrl : void 0
 		};
 	}
 	async removeMcpServer(id) {
-		this.mcp.closeConnection(id);
-		this.mcp.unregisterCallbackUrl(id);
-		this.sql`
-      DELETE FROM cf_agents_mcp_servers WHERE id = ${id};
-    `;
-		this.broadcastMcpServers();
+		if (this.mcp.mcpConnections[id]) await this.mcp.closeConnection(id);
+		await this.mcp.removeServer(id);
 	}
-	getMcpServers() {
+	async getMcpServers() {
 		const mcpState = {
 			prompts: this.mcp.listPrompts(),
 			resources: this.mcp.listResources(),
 			servers: {},
 			tools: this.mcp.listTools()
 		};
-		const servers = this.sql`
-      SELECT id, name, server_url, client_id, auth_url, callback_url, server_options FROM cf_agents_mcp_servers;
-    `;
+		const servers = await 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";
+			if (!serverConn && server.auth_url) defaultState = "authenticating";
 			mcpState.servers[server.id] = {
 				auth_url: server.auth_url,
 				capabilities: serverConn?.serverCapabilities ?? null,
 				instructions: serverConn?.instructions ?? null,
 				name: server.name,
 				server_url: server.server_url,
-				state: serverConn?.connectionState ?? "authenticating"
+				state: serverConn?.connectionState ?? defaultState
 			};
 		}
 		return mcpState;
 	}
-	broadcastMcpServers() {
+	async broadcastMcpServers() {
 		this.broadcast(JSON.stringify({
-			mcp: this.getMcpServers(),
+			mcp: await this.getMcpServers(),
 			type: MessageType.CF_AGENT_MCP_SERVERS
 		}));
 	}
 	/**
+	* Handle MCP OAuth callback request if it's an OAuth callback.
+	*
+	* This method encapsulates the entire OAuth callback flow:
+	* 1. Checks if the request is an MCP OAuth callback
+	* 2. Processes the OAuth code exchange
+	* 3. Establishes the connection if successful
+	* 4. Broadcasts MCP server state updates
+	* 5. Returns the appropriate HTTP response
+	*
+	* @param request The incoming HTTP request
+	* @returns Response if this was an OAuth callback, null otherwise
+	*/
+	async handleMcpOAuthCallback(request) {
+		if (!await 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();
+			});
+		}
+		return this.handleOAuthCallbackResponse(result, request);
+	}
+	/**
 	* Handle OAuth callback response using MCPClientManager configuration
 	* @param result OAuth callback result
 	* @param request The original request (needed for base URL)
@@ -972,10 +1004,21 @@ var Agent = class Agent extends Server {
 	handleOAuthCallbackResponse(result, request) {
 		const config = this.mcp.getOAuthCallbackConfig();
 		if (config?.customHandler) return config.customHandler(result);
-		if (config?.successRedirect && result.authSuccess) return Response.redirect(config.successRedirect);
-		if (config?.errorRedirect && !result.authSuccess) return Response.redirect(`${config.errorRedirect}?error=${encodeURIComponent(result.authError || "Unknown error")}`);
-		const baseUrl = new URL(request.url).origin;
-		return Response.redirect(baseUrl);
+		const baseOrigin = new URL(request.url).origin;
+		if (config?.successRedirect && result.authSuccess) try {
+			return Response.redirect(new URL(config.successRedirect, baseOrigin).href);
+		} catch (e) {
+			console.error("Invalid successRedirect URL:", config.successRedirect, e);
+			return Response.redirect(baseOrigin);
+		}
+		if (config?.errorRedirect && !result.authSuccess) try {
+			const errorUrl = `${config.errorRedirect}?error=${encodeURIComponent(result.authError || "Unknown error")}`;
+			return Response.redirect(new URL(errorUrl, baseOrigin).href);
+		} catch (e) {
+			console.error("Invalid errorRedirect URL:", config.errorRedirect, e);
+			return Response.redirect(baseOrigin);
+		}
+		return Response.redirect(baseOrigin);
 	}
 };
 const wrappedClasses = /* @__PURE__ */ new Set();
@@ -1001,10 +1044,15 @@ async function routeAgentRequest(request, env, options) {
 		prefix: "agents",
 		...options
 	});
-	if (response && corsHeaders && request.headers.get("upgrade")?.toLowerCase() !== "websocket" && request.headers.get("Upgrade")?.toLowerCase() !== "websocket") response = new Response(response.body, { headers: {
-		...response.headers,
-		...corsHeaders
-	} });
+	if (response && corsHeaders && request.headers.get("upgrade")?.toLowerCase() !== "websocket" && request.headers.get("Upgrade")?.toLowerCase() !== "websocket") {
+		const newHeaders = new Headers(response.headers);
+		for (const [key, value] of Object.entries(corsHeaders)) newHeaders.set(key, value);
+		response = new Response(response.body, {
+			status: response.status,
+			statusText: response.statusText,
+			headers: newHeaders
+		});
+	}
 	return response;
 }
 /**
@@ -1197,4 +1245,4 @@ var StreamingResponse = class {
 
 //#endregion
 export { createCatchAllEmailResolver as a, getCurrentAgent as c, unstable_callable as d, genericObservability as f, createAddressBasedEmailResolver as i, routeAgentEmail as l, StreamingResponse as n, createHeaderBasedEmailResolver as o, callable as r, getAgentByName as s, Agent as t, routeAgentRequest as u };
-//# sourceMappingURL=src-Dz0H9hSU.js.map
+//# sourceMappingURL=src-C6rC6ZpH.js.map