agents 0.0.0-5908188 → 0.0.0-59ac254

package/dist/{src-CTtjSFyX.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 { a as DisposableStore, t as MCPClientManager } from "./client-CIvp_OWw.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,8 +152,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,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,
@@ -227,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({
@@ -263,8 +263,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 +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 {
@@ -340,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
@@ -367,27 +369,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;
-          `;
+				await this._tryCatch(async () => {
+					await this.mcp.restoreConnectionsFromStorage(this.name);
 					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();
-							});
-						});
-					}
 					return _onStart(props);
 				});
 			});
@@ -732,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)
@@ -772,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: {
@@ -793,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
     `;
@@ -807,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(),
@@ -839,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;
@@ -850,89 +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 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) {
+		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 = {
@@ -942,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 = {
@@ -972,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;
@@ -995,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)
@@ -1003,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();
@@ -1032,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;
 }
 /**
@@ -1228,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-CTtjSFyX.js.map
+//# sourceMappingURL=src-BmbDclOA.js.map