agents 0.0.0-ccbea72 → 0.0.0-cd8b7fd

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

@@ -1,8 +1,8 @@
-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 { AsyncLocalStorage } from "node:async_hooks";
+import { t as agentContext } from "./context-BkKbAa1R.js";
+import { t as MessageType } from "./ai-types-CwgHzwUb.js";
+import { r as camelCaseToKebabCase } from "./client-CcyhkGfN.js";
+import { i as DisposableStore, r as MCPConnectionState, t as MCPClientManager } from "./client-QZa2Rq0l.js";
+import { t as DurableObjectOAuthClientProvider } from "./do-oauth-client-provider-B1fVIshX.js";
 import { parseCronExpression } from "cron-schedule";
 import { nanoid } from "nanoid";
 import { EmailMessage } from "cloudflare:email";
@@ -72,7 +72,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 {
@@ -153,9 +152,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,27 +191,37 @@ 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();
 		};
 		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.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,
@@ -228,34 +236,25 @@ 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())
       )
     `;
+		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({
@@ -264,17 +263,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 +277,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,7 +342,7 @@ 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
@@ -378,7 +370,7 @@ var Agent = class Agent extends Server {
 				email: void 0
 			}, async () => {
 				await this._tryCatch(async () => {
-					await this._ensureMcpStateRestored();
+					await this.mcp.restoreConnectionsFromStorage(this.name);
 					this.broadcastMcpServers();
 					return _onStart(props);
 				});
@@ -724,10 +716,7 @@ var Agent = class Agent 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)
@@ -764,11 +753,12 @@ var Agent = class Agent 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: {
@@ -785,7 +775,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 +789,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 +816,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
 	*
@@ -868,7 +824,8 @@ var Agent = class Agent 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;
@@ -879,29 +836,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) {
+		await this.mcp.ensureJsonSchema();
+		const id = nanoid(8);
 		const authProvider = new DurableObjectOAuthClientProvider(this.ctx.storage, this.name, callbackUrl);
-		if (reconnect) {
-			authProvider.serverId = reconnect.id;
-			if (reconnect.oauthClientId) authProvider.clientId = reconnect.oauthClientId;
-		}
+		authProvider.serverId = id;
 		const transportType = options?.transport?.type ?? "auto";
 		let headerTransportOpts = {};
 		if (options?.transport?.headers) headerTransportOpts = {
@@ -911,28 +849,33 @@ 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);
+		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 {
-			authUrl,
-			clientId,
-			id
+			id,
+			state: MCPConnectionState.READY
 		};
 	}
 	async removeMcpServer(id) {
-		this.mcp.closeConnection(id);
-		this.mcp.unregisterCallbackUrl(id);
-		this.sql`
-      DELETE FROM cf_agents_mcp_servers WHERE id = ${id};
-    `;
-		this.broadcastMcpServers();
+		await this.mcp.removeServer(id);
 	}
 	getMcpServers() {
 		const mcpState = {
@@ -941,18 +884,18 @@ var Agent = class Agent extends Server {
 			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 = 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;
@@ -964,6 +907,28 @@ var Agent = class Agent extends Server {
 		}));
 	}
 	/**
+	* 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 (!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] Connection establishment failed:", error);
+		});
+		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 +937,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 +977,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 +1178,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-BmbDclOA.js.map