agents 0.0.0-a215bb2 → 0.0.0-a315e86

package/dist/{src-L3cHuAag.js → src-Dk8lwxHf.js}

@@ -1,13 +1,81 @@
-import { MessageType } from "./ai-types-UZlfLOYP.js";
-import { camelCaseToKebabCase } from "./client-DjR-lC16.js";
-import { DisposableStore, MCPClientManager } from "./client-CZBVDDoO.js";
-import { DurableObjectOAuthClientProvider } from "./do-oauth-client-provider-B2jr6UNq.js";
+import { t as MessageType } from "./ai-types-B3aQaFv3.js";
+import { r as camelCaseToKebabCase } from "./client-BfiZ3HQd.js";
+import { r as DisposableStore, t as MCPClientManager } from "./client-DpkZyXgJ.js";
+import { t as DurableObjectOAuthClientProvider } from "./do-oauth-client-provider-D2P1lSft.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,8 +221,8 @@ var Agent = class Agent extends Server {
 		super(ctx, env);
 		this._state = DEFAULT_STATE;
 		this._disposables = new DisposableStore();
+		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;
@@ -192,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
@@ -227,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);
@@ -263,8 +332,9 @@ var Agent = class Agent extends Server {
 				request,
 				email: void 0
 			}, async () => {
-				const callbackResult = await this._handlePotentialOAuthCallback(request);
-				if (callbackResult) return callbackResult;
+				await this.mcp.ensureJsonSchema();
+				const oauthResponse = await this.handleMcpOAuthCallback(request);
+				if (oauthResponse) return oauthResponse;
 				return this._tryCatch(() => _onRequest(request));
 			});
 		};
@@ -276,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 {
@@ -340,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({
@@ -367,27 +438,9 @@ var Agent = class Agent extends Server {
 				request: void 0,
 				email: void 0
 			}, async () => {
-				await this._tryCatch(() => {
-					const servers = this.sql`
-            SELECT id, name, server_url, client_id, auth_url, callback_url, server_options FROM cf_agents_mcp_servers;
-          `;
-					this.broadcastMcpServers();
-					if (servers && Array.isArray(servers) && servers.length > 0) {
-						servers.forEach((server) => {
-							if (server.callback_url) this.mcp.registerCallbackUrl(`${server.callback_url}/${server.id}`);
-						});
-						servers.forEach((server) => {
-							this._connectToMcpServerInternal(server.name, server.server_url, server.callback_url, server.server_options ? JSON.parse(server.server_options) : void 0, {
-								id: server.id,
-								oauthClientId: server.client_id ?? void 0
-							}).then(() => {
-								this.broadcastMcpServers();
-							}).catch((error) => {
-								console.error(`Error connecting to MCP server: ${server.name} (${server.server_url})`, error);
-								this.broadcastMcpServers();
-							});
-						});
-					}
+				await this._tryCatch(async () => {
+					await this.mcp.restoreConnectionsFromStorage(this.name);
+					await this.broadcastMcpServers();
 					return _onStart(props);
 				});
 			});
@@ -793,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
     `;
@@ -807,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(),
@@ -850,89 +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 serverId = nanoid(8);
-		this.sql`
-        INSERT OR REPLACE INTO cf_agents_mcp_servers (id, name, server_url, client_id, auth_url, callback_url, server_options)
-      VALUES (
-        ${serverId},
-        ${serverName},
-        ${url},
-        ${null},
-        ${null},
-        ${callbackUrl},
-        ${options ? JSON.stringify(options) : null}
-        );
-    `;
-		const result = await this._connectToMcpServerInternal(serverName, url, callbackUrl, options, { id: serverId });
-		if (result.clientId || result.authUrl) this.sql`
-        UPDATE cf_agents_mcp_servers
-        SET client_id = ${result.clientId ?? null}, auth_url = ${result.authUrl ?? null}
-        WHERE id = ${serverId}
-      `;
-		this.broadcastMcpServers();
-		return result;
-	}
-	/**
-	* Handle potential OAuth callback requests after DO hibernation.
-	* Detects OAuth callbacks, restores state from database, and processes the callback.
-	* Returns a Response if this was an OAuth callback, otherwise returns undefined.
-	*/
-	async _handlePotentialOAuthCallback(request) {
-		if (request.method !== "GET") return;
-		const url = new URL(request.url);
-		if (!(url.pathname.includes("/callback/") && url.searchParams.has("code"))) return;
-		const pathParts = url.pathname.split("/");
-		const callbackIndex = pathParts.indexOf("callback");
-		const serverId = callbackIndex !== -1 ? pathParts[callbackIndex + 1] : null;
-		if (!serverId) return new Response("Invalid callback URL: missing serverId", { status: 400 });
-		if (this.mcp.isCallbackRequest(request) && this.mcp.mcpConnections[serverId]) return this._processOAuthCallback(request);
-		try {
-			const server = this.sql`
-        SELECT id, name, server_url, client_id, auth_url, callback_url, server_options
-        FROM cf_agents_mcp_servers
-        WHERE id = ${serverId}
-      `.find((s) => s.id === serverId);
-			if (!server) return new Response(`OAuth callback failed: Server ${serverId} not found in database`, { status: 404 });
-			if (!server.callback_url) return new Response(`OAuth callback failed: No callback URL stored for server ${serverId}`, { status: 500 });
-			this.mcp.registerCallbackUrl(`${server.callback_url}/${server.id}`);
-			if (!this.mcp.mcpConnections[serverId]) {
-				let parsedOptions;
-				try {
-					parsedOptions = server.server_options ? JSON.parse(server.server_options) : void 0;
-				} catch {
-					return new Response(`OAuth callback failed: Invalid server options in database for ${serverId}`, { status: 500 });
-				}
-				await this._connectToMcpServerInternal(server.name, server.server_url, server.callback_url, parsedOptions, {
-					id: server.id,
-					oauthClientId: server.client_id ?? void 0
-				});
-			}
-			return this._processOAuthCallback(request);
-		} catch (error) {
-			const errorMsg = error instanceof Error ? error.message : "Unknown error";
-			console.error(`Failed to restore MCP state for ${serverId}:`, error);
-			return new Response(`OAuth callback failed during state restoration: ${errorMsg}`, { status: 500 });
-		}
-	}
-	/**
-	* Process an OAuth callback request (assumes state is already restored)
-	*/
-	async _processOAuthCallback(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);
-	}
-	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 = {
@@ -942,59 +919,78 @@ 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.mcp.establishConnection(result.serverId).catch((error) => {
+			console.error("[Agent handleMcpOAuthCallback] Background connection failed:", error);
+		});
+		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)
@@ -1003,10 +999,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();
@@ -1032,10 +1039,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;
 }
 /**
@@ -1227,5 +1239,5 @@ var StreamingResponse = class {
 };
 
 //#endregion
-export { Agent, StreamingResponse, callable, createAddressBasedEmailResolver, createCatchAllEmailResolver, createHeaderBasedEmailResolver, genericObservability, getAgentByName, getCurrentAgent, routeAgentEmail, routeAgentRequest, unstable_callable };
-//# sourceMappingURL=src-L3cHuAag.js.map
+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-Dk8lwxHf.js.map