agents 0.0.0-fbaa8f7 → 0.0.0-fce47ef

package/src/index.ts

@@ -1,19 +1,30 @@
+import { AsyncLocalStorage } from "node:async_hooks";
+import type { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import type { SSEClientTransportOptions } from "@modelcontextprotocol/sdk/client/sse.js";
+
+import type {
+  Prompt,
+  Resource,
+  ServerCapabilities,
+  Tool,
+} from "@modelcontextprotocol/sdk/types.js";
+import { parseCronExpression } from "cron-schedule";
+import { nanoid } from "nanoid";
 import {
-  Server,
-  routePartykitRequest,
-  type PartyServerOptions,
-  getServerByName,
   type Connection,
   type ConnectionContext,
+  getServerByName,
+  type PartyServerOptions,
+  routePartykitRequest,
+  Server,
   type WSMessage,
 } from "partyserver";
+import { camelCaseToKebabCase } from "./client";
+import { MCPClientManager } from "./mcp/client";
+// import type { MCPClientConnection } from "./mcp/client-connection";
+import { DurableObjectOAuthClientProvider } from "./mcp/do-oauth-client-provider";
 
-import { parseCronExpression } from "cron-schedule";
-import { nanoid } from "nanoid";
-
-export type { Connection, WSMessage, ConnectionContext } from "partyserver";
-
-import { WorkflowEntrypoint as CFWorkflowEntrypoint } from "cloudflare:workers";
+export type { Connection, ConnectionContext, WSMessage } from "partyserver";
 
 /**
  * RPC request message from client
@@ -97,7 +108,6 @@ export type CallableMetadata = {
   streaming?: boolean;
 };
 
-// biome-ignore lint/complexity/noBannedTypes: <explanation>
 const callableMetadata = new Map<Function, CallableMetadata>();
 
 /**
@@ -107,6 +117,7 @@ const callableMetadata = new Map<Function, CallableMetadata>();
 export function unstable_callable(metadata: CallableMetadata = {}) {
   return function callableDecorator<This, Args extends unknown[], Return>(
     target: (this: This, ...args: Args) => Return,
+    // biome-ignore lint/correctness/noUnusedFunctionParameters: later
     context: ClassMethodDecoratorContext
   ) {
     if (!callableMetadata.has(target)) {
@@ -117,11 +128,6 @@ export function unstable_callable(metadata: CallableMetadata = {}) {
   };
 }
 
-/**
- * A class for creating workflow entry points that can be used with Cloudflare Workers
- */
-export class WorkflowEntrypoint extends CFWorkflowEntrypoint {}
-
 /**
  * Represents a scheduled task within an Agent
  * @template T Type of the payload data
@@ -163,18 +169,95 @@ function getNextCronTime(cron: string) {
   return interval.getNextDate();
 }
 
+/**
+ * MCP Server state update message from server -> Client
+ */
+export type MCPServerMessage = {
+  type: "cf_agent_mcp_servers";
+  mcp: MCPServersState;
+};
+
+export type MCPServersState = {
+  servers: {
+    [id: string]: MCPServer;
+  };
+  tools: Tool[];
+  prompts: Prompt[];
+  resources: Resource[];
+};
+
+export type MCPServer = {
+  name: string;
+  server_url: string;
+  auth_url: string | null;
+  // This state is specifically about the temporary process of getting a token (if needed).
+  // Scope outside of that can't be relied upon because when the DO sleeps, there's no way
+  // to communicate a change to a non-ready state.
+  state: "authenticating" | "connecting" | "ready" | "discovering" | "failed";
+  instructions: string | null;
+  capabilities: ServerCapabilities | null;
+};
+
+/**
+ * MCP Server data stored in DO SQL for resuming MCP Server connections
+ */
+type MCPServerRow = {
+  id: string;
+  name: string;
+  server_url: string;
+  client_id: string | null;
+  auth_url: string | null;
+  callback_url: string;
+  server_options: string;
+};
+
 const STATE_ROW_ID = "cf_state_row_id";
 const STATE_WAS_CHANGED = "cf_state_was_changed";
 
 const DEFAULT_STATE = {} as unknown;
 
+const agentContext = new AsyncLocalStorage<{
+  agent: Agent<unknown>;
+  connection: Connection | undefined;
+  request: Request | undefined;
+}>();
+
+export function getCurrentAgent<
+  T extends Agent<unknown, unknown> = Agent<unknown, unknown>,
+>(): {
+  agent: T | undefined;
+  connection: Connection | undefined;
+  request: Request<unknown, CfProperties<unknown>> | undefined;
+} {
+  const store = agentContext.getStore() as
+    | {
+        agent: T;
+        connection: Connection | undefined;
+        request: Request<unknown, CfProperties<unknown>> | undefined;
+      }
+    | undefined;
+  if (!store) {
+    return {
+      agent: undefined,
+      connection: undefined,
+      request: undefined,
+    };
+  }
+  return store;
+}
+
 /**
  * Base class for creating Agent implementations
  * @template Env Environment type containing bindings
  * @template State State type to store within the Agent
  */
 export class Agent<Env, State = unknown> extends Server<Env> {
-  #state = DEFAULT_STATE as State;
+  private _state = DEFAULT_STATE as State;
+
+  private _ParentClass: typeof Agent<Env, State> =
+    Object.getPrototypeOf(this).constructor;
+
+  mcp: MCPClientManager = new MCPClientManager(this._ParentClass.name, "0.0.1");
 
   /**
    * Initial state for the Agent
@@ -186,9 +269,9 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    * Current state of the Agent
    */
   get state(): State {
-    if (this.#state !== DEFAULT_STATE) {
+    if (this._state !== DEFAULT_STATE) {
       // state was previously set, and populated internal state
-      return this.#state;
+      return this._state;
     }
     // looks like this is the first time the state is being accessed
     // check if the state was set in a previous life
@@ -208,8 +291,8 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     ) {
       const state = result[0]?.state as string; // could be null?
 
-      this.#state = JSON.parse(state);
-      return this.#state;
+      this._state = JSON.parse(state);
+      return this._state;
     }
 
     // ok, this is the first time the state is being accessed
@@ -256,7 +339,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       return [...this.ctx.storage.sql.exec(query, ...values)] as T[];
     } catch (e) {
       console.error(`failed to execute sql query: ${query}`, e);
-      throw e;
+      throw this.onError(e);
     }
   }
   constructor(ctx: AgentContext, env: Env) {
@@ -270,7 +353,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     `;
 
     void this.ctx.blockConcurrencyWhile(async () => {
-      try {
+      return this._tryCatch(async () => {
         // Create alarms table if it doesn't exist
         this.sql`
         CREATE TABLE IF NOT EXISTS cf_agents_schedules (
@@ -287,102 +370,201 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
         // execute any pending alarms and schedule the next alarm
         await this.alarm();
-      } catch (e) {
-        console.error(e);
-        throw e;
-      }
+      });
     });
 
-    const _onMessage = this.onMessage.bind(this);
-    this.onMessage = async (connection: Connection, message: WSMessage) => {
-      if (typeof message !== "string") {
-        return _onMessage(connection, message);
-      }
-
-      let parsed: unknown;
-      try {
-        parsed = JSON.parse(message);
-      } catch (e) {
-        // silently fail and let the onMessage handler handle it
-        return _onMessage(connection, message);
-      }
+    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
+      )
+    `;
 
-      if (isStateUpdateMessage(parsed)) {
-        this.#setStateInternal(parsed.state as State, connection);
-        return;
-      }
+    const _onRequest = this.onRequest.bind(this);
+    this.onRequest = (request: Request) => {
+      return agentContext.run(
+        { agent: this, connection: undefined, request },
+        async () => {
+          if (this.mcp.isCallbackRequest(request)) {
+            await this.mcp.handleCallbackRequest(request);
+
+            // after the MCP connection handshake, we can send updated mcp state
+            this.broadcast(
+              JSON.stringify({
+                mcp: this.getMcpServers(),
+                type: "cf_agent_mcp_servers",
+              })
+            );
+
+            // We probably should let the user configure this response/redirect, but this is fine for now.
+            return new Response("<script>window.close();</script>", {
+              headers: { "content-type": "text/html" },
+              status: 200,
+            });
+          }
 
-      if (isRPCRequest(parsed)) {
-        try {
-          const { id, method, args } = parsed;
+          return this._tryCatch(() => _onRequest(request));
+        }
+      );
+    };
 
-          // Check if method exists and is callable
-          const methodFn = this[method as keyof this];
-          if (typeof methodFn !== "function") {
-            throw new Error(`Method ${method} does not exist`);
+    const _onMessage = this.onMessage.bind(this);
+    this.onMessage = async (connection: Connection, message: WSMessage) => {
+      return agentContext.run(
+        { agent: this, connection, request: undefined },
+        async () => {
+          if (typeof message !== "string") {
+            return this._tryCatch(() => _onMessage(connection, message));
           }
 
-          if (!this.isCallable(method)) {
-            throw new Error(`Method ${method} is not callable`);
+          let parsed: unknown;
+          try {
+            parsed = JSON.parse(message);
+          } catch (_e) {
+            // silently fail and let the onMessage handler handle it
+            return this._tryCatch(() => _onMessage(connection, message));
           }
 
-          // biome-ignore lint/complexity/noBannedTypes: <explanation>
-          const metadata = callableMetadata.get(methodFn as Function);
+          if (isStateUpdateMessage(parsed)) {
+            this._setStateInternal(parsed.state as State, connection);
+            return;
+          }
 
-          // For streaming methods, pass a StreamingResponse object
-          if (metadata?.streaming) {
-            const stream = new StreamingResponse(connection, id);
-            await methodFn.apply(this, [stream, ...args]);
+          if (isRPCRequest(parsed)) {
+            try {
+              const { id, method, args } = parsed;
+
+              // Check if method exists and is callable
+              const methodFn = this[method as keyof this];
+              if (typeof methodFn !== "function") {
+                throw new Error(`Method ${method} does not exist`);
+              }
+
+              if (!this._isCallable(method)) {
+                throw new Error(`Method ${method} is not callable`);
+              }
+
+              const metadata = callableMetadata.get(methodFn as Function);
+
+              // For streaming methods, pass a StreamingResponse object
+              if (metadata?.streaming) {
+                const stream = new StreamingResponse(connection, id);
+                await methodFn.apply(this, [stream, ...args]);
+                return;
+              }
+
+              // For regular methods, execute and send response
+              const result = await methodFn.apply(this, args);
+              const response: RPCResponse = {
+                done: true,
+                id,
+                result,
+                success: true,
+                type: "rpc",
+              };
+              connection.send(JSON.stringify(response));
+            } catch (e) {
+              // Send error response
+              const response: RPCResponse = {
+                error:
+                  e instanceof Error ? e.message : "Unknown error occurred",
+                id: parsed.id,
+                success: false,
+                type: "rpc",
+              };
+              connection.send(JSON.stringify(response));
+              console.error("RPC error:", e);
+            }
             return;
           }
 
-          // For regular methods, execute and send response
-          const result = await methodFn.apply(this, args);
-          const response: RPCResponse = {
-            type: "rpc",
-            id,
-            success: true,
-            result,
-            done: true,
-          };
-          connection.send(JSON.stringify(response));
-        } catch (e) {
-          // Send error response
-          const response: RPCResponse = {
-            type: "rpc",
-            id: parsed.id,
-            success: false,
-            error: e instanceof Error ? e.message : "Unknown error occurred",
-          };
-          connection.send(JSON.stringify(response));
-          console.error("RPC error:", e);
+          return this._tryCatch(() => _onMessage(connection, message));
         }
-        return;
-      }
-
-      return _onMessage(connection, message);
+      );
     };
 
     const _onConnect = this.onConnect.bind(this);
     this.onConnect = (connection: Connection, ctx: ConnectionContext) => {
       // TODO: This is a hack to ensure the state is sent after the connection is established
       // must fix this
-      setTimeout(() => {
-        if (this.state) {
-          connection.send(
+      return agentContext.run(
+        { agent: this, connection, request: ctx.request },
+        async () => {
+          setTimeout(() => {
+            if (this.state) {
+              connection.send(
+                JSON.stringify({
+                  state: this.state,
+                  type: "cf_agent_state",
+                })
+              );
+            }
+
+            connection.send(
+              JSON.stringify({
+                mcp: this.getMcpServers(),
+                type: "cf_agent_mcp_servers",
+              })
+            );
+
+            return this._tryCatch(() => _onConnect(connection, ctx));
+          }, 20);
+        }
+      );
+    };
+
+    const _onStart = this.onStart.bind(this);
+    this.onStart = async () => {
+      return agentContext.run(
+        { agent: this, connection: undefined, request: undefined },
+        async () => {
+          const servers = this.sql<MCPServerRow>`
+            SELECT id, name, server_url, client_id, auth_url, callback_url, server_options FROM cf_agents_mcp_servers;
+          `;
+
+          // from DO storage, reconnect to all servers not currently in the oauth flow using our saved auth information
+          await Promise.allSettled(
+            servers
+              .filter((server) => server.auth_url === null)
+              .map((server) => {
+                return this._connectToMcpServerInternal(
+                  server.name,
+                  server.server_url,
+                  server.callback_url,
+                  server.server_options
+                    ? JSON.parse(server.server_options)
+                    : undefined,
+                  {
+                    id: server.id,
+                    oauthClientId: server.client_id ?? undefined,
+                  }
+                );
+              })
+          );
+
+          this.broadcast(
             JSON.stringify({
-              type: "cf_agent_state",
-              state: this.state,
+              mcp: this.getMcpServers(),
+              type: "cf_agent_mcp_servers",
             })
           );
+
+          await this._tryCatch(() => _onStart());
         }
-        _onConnect(connection, ctx);
-      }, 20);
+      );
     };
   }
 
-  #setStateInternal(state: State, source: Connection | "server" = "server") {
-    this.#state = state;
+  private _setStateInternal(
+    state: State,
+    source: Connection | "server" = "server"
+  ) {
+    this._state = state;
     this.sql`
     INSERT OR REPLACE INTO cf_agents_state (id, state)
     VALUES (${STATE_ROW_ID}, ${JSON.stringify(state)})
@@ -393,12 +575,20 @@ export class Agent<Env, State = unknown> extends Server<Env> {
   `;
     this.broadcast(
       JSON.stringify({
-        type: "cf_agent_state",
         state: state,
+        type: "cf_agent_state",
       }),
       source !== "server" ? [source.id] : []
     );
-    this.onStateUpdate(state, source);
+    return this._tryCatch(() => {
+      const { connection, request } = agentContext.getStore() || {};
+      return agentContext.run(
+        { agent: this, connection, request },
+        async () => {
+          return this.onStateUpdate(state, source);
+        }
+      );
+    });
   }
 
   /**
@@ -406,7 +596,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    * @param state New state to set
    */
   setState(state: State) {
-    this.#setStateInternal(state, "server");
+    this._setStateInternal(state, "server");
   }
 
   /**
@@ -414,6 +604,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    * @param state Updated state
    * @param source Source of the state update ("server" or a client connection)
    */
+  // biome-ignore lint/correctness/noUnusedFunctionParameters: overridden later
   onStateUpdate(state: State | undefined, source: Connection | "server") {
     // override this to handle state updates
   }
@@ -422,8 +613,49 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    * Called when the Agent receives an email
    * @param email Email message to process
    */
+  // biome-ignore lint/correctness/noUnusedFunctionParameters: overridden later
   onEmail(email: ForwardableEmailMessage) {
-    throw new Error("Not implemented");
+    return agentContext.run(
+      { agent: this, connection: undefined, request: undefined },
+      async () => {
+        console.error("onEmail not implemented");
+      }
+    );
+  }
+
+  private async _tryCatch<T>(fn: () => T | Promise<T>) {
+    try {
+      return await fn();
+    } catch (e) {
+      throw this.onError(e);
+    }
+  }
+
+  override onError(
+    connection: Connection,
+    error: unknown
+  ): void | Promise<void>;
+  override onError(error: unknown): void | Promise<void>;
+  override onError(connectionOrError: Connection | unknown, error?: unknown) {
+    let theError: unknown;
+    if (connectionOrError && error) {
+      theError = error;
+      // this is a websocket connection error
+      console.error(
+        "Error on websocket connection:",
+        (connectionOrError as Connection).id,
+        theError
+      );
+      console.error(
+        "Override onError(connection, error) to handle websocket connection errors"
+      );
+    } else {
+      theError = connectionOrError;
+      // this is a server error
+      console.error("Error on server:", theError);
+      console.error("Override onError(error) to handle server errors");
+    }
+    throw theError;
   }
 
   /**
@@ -465,11 +697,11 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         )}, 'scheduled', ${timestamp})
       `;
 
-      await this.scheduleNextAlarm();
+      await this._scheduleNextAlarm();
 
       return {
-        id,
         callback: callback,
+        id,
         payload: payload as T,
         time: timestamp,
         type: "scheduled",
@@ -486,13 +718,13 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         )}, 'delayed', ${when}, ${timestamp})
       `;
 
-      await this.scheduleNextAlarm();
+      await this._scheduleNextAlarm();
 
       return {
-        id,
         callback: callback,
-        payload: payload as T,
         delayInSeconds: when,
+        id,
+        payload: payload as T,
         time: timestamp,
         type: "delayed",
       };
@@ -508,13 +740,13 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         )}, 'cron', ${when}, ${timestamp})
       `;
 
-      await this.scheduleNextAlarm();
+      await this._scheduleNextAlarm();
 
       return {
-        id,
         callback: callback,
-        payload: payload as T,
         cron: when,
+        id,
+        payload: payload as T,
         time: timestamp,
         type: "cron",
       };
@@ -532,7 +764,10 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     const result = this.sql<Schedule<string>>`
       SELECT * FROM cf_agents_schedules WHERE id = ${id}
     `;
-    if (!result) return undefined;
+    if (!result) {
+      console.error(`schedule ${id} not found`);
+      return undefined;
+    }
 
     return { ...result[0], payload: JSON.parse(result[0].payload) as T };
   }
@@ -545,7 +780,6 @@ export class Agent<Env, State = unknown> extends Server<Env> {
    */
   getSchedules<T = string>(
     criteria: {
-      description?: string;
       id?: string;
       type?: "scheduled" | "delayed" | "cron";
       timeRange?: { start?: Date; end?: Date };
@@ -559,11 +793,6 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       params.push(criteria.id);
     }
 
-    if (criteria.description) {
-      query += " AND description = ?";
-      params.push(criteria.description);
-    }
-
     if (criteria.type) {
       query += " AND type = ?";
       params.push(criteria.type);
@@ -598,11 +827,11 @@ export class Agent<Env, State = unknown> extends Server<Env> {
   async cancelSchedule(id: string): Promise<boolean> {
     this.sql`DELETE FROM cf_agents_schedules WHERE id = ${id}`;
 
-    await this.scheduleNextAlarm();
+    await this._scheduleNextAlarm();
     return true;
   }
 
-  private async scheduleNextAlarm() {
+  private async _scheduleNextAlarm() {
     // Find the next schedule that needs to be executed
     const result = this.sql`
       SELECT time FROM cf_agents_schedules 
@@ -619,10 +848,14 @@ export class Agent<Env, State = unknown> extends Server<Env> {
   }
 
   /**
-   * Method called when an alarm fires
-   * Executes any scheduled tasks that are due
+   * Method called when an alarm fires.
+   * Executes any scheduled tasks that are due.
+   *
+   * @remarks
+   * To schedule a task, please use the `this.schedule` method instead.
+   * See {@link https://developers.cloudflare.com/agents/api-reference/schedule-tasks/}
    */
-  async alarm() {
+  public readonly alarm = async () => {
     const now = Math.floor(Date.now() / 1000);
 
     // Get all schedules that should be executed now
@@ -636,16 +869,21 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         console.error(`callback ${row.callback} not found`);
         continue;
       }
-      try {
-        (
-          callback as (
-            payload: unknown,
-            schedule: Schedule<unknown>
-          ) => Promise<void>
-        ).bind(this)(JSON.parse(row.payload as string), row);
-      } catch (e) {
-        console.error(`error executing callback ${row.callback}`, e);
-      }
+      await agentContext.run(
+        { agent: this, connection: undefined, request: undefined },
+        async () => {
+          try {
+            await (
+              callback as (
+                payload: unknown,
+                schedule: Schedule<unknown>
+              ) => Promise<void>
+            ).bind(this)(JSON.parse(row.payload as string), row);
+          } catch (e) {
+            console.error(`error executing callback "${row.callback}"`, e);
+          }
+        }
+      );
       if (row.type === "cron") {
         // Update next execution time for cron schedules
         const nextExecutionTime = getNextCronTime(row.cron);
@@ -663,8 +901,8 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     }
 
     // Schedule the next alarm
-    await this.scheduleNextAlarm();
-  }
+    await this._scheduleNextAlarm();
+  };
 
   /**
    * Destroy the Agent, removing all state and scheduled tasks
@@ -673,20 +911,184 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     // drop all tables
     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`;
 
     // delete all alarms
     await this.ctx.storage.deleteAlarm();
     await this.ctx.storage.deleteAll();
+    this.ctx.abort("destroyed"); // enforce that the agent is evicted
   }
 
   /**
    * Get all methods marked as callable on this Agent
    * @returns A map of method names to their metadata
    */
-  private isCallable(method: string): boolean {
-    // biome-ignore lint/complexity/noBannedTypes: <explanation>
+  private _isCallable(method: string): boolean {
     return callableMetadata.has(this[method as keyof this] as Function);
   }
+
+  /**
+   * Connect to a new MCP Server
+   *
+   * @param url MCP Server SSE URL
+   * @param callbackHost Base host for the agent, used for the redirect URI.
+   * @param agentsPrefix agents routing prefix if not using `agents`
+   * @param options MCP client and transport (header) options
+   * @returns authUrl
+   */
+  async addMcpServer(
+    serverName: string,
+    url: string,
+    callbackHost: string,
+    agentsPrefix = "agents",
+    options?: {
+      client?: ConstructorParameters<typeof Client>[1];
+      transport?: {
+        headers: HeadersInit;
+      };
+    }
+  ): Promise<{ id: string; authUrl: string | undefined }> {
+    const callbackUrl = `${callbackHost}/${agentsPrefix}/${camelCaseToKebabCase(this._ParentClass.name)}/${this.name}/callback`;
+
+    const result = await this._connectToMcpServerInternal(
+      serverName,
+      url,
+      callbackUrl,
+      options
+    );
+
+    this.broadcast(
+      JSON.stringify({
+        mcp: this.getMcpServers(),
+        type: "cf_agent_mcp_servers",
+      })
+    );
+
+    return result;
+  }
+
+  async _connectToMcpServerInternal(
+    serverName: string,
+    url: string,
+    callbackUrl: string,
+    // it's important that any options here are serializable because we put them into our sqlite DB for reconnection purposes
+    options?: {
+      client?: ConstructorParameters<typeof Client>[1];
+      /**
+       * We don't expose the normal set of transport options because:
+       * 1) we can't serialize things like the auth provider or a fetch function into the DB for reconnection purposes
+       * 2) We probably want these options to be agnostic to the transport type (SSE vs Streamable)
+       *
+       * This has the limitation that you can't override fetch, but I think headers should handle nearly all cases needed (i.e. non-standard bearer auth).
+       */
+      transport?: {
+        headers?: HeadersInit;
+      };
+    },
+    reconnect?: {
+      id: string;
+      oauthClientId?: string;
+    }
+  ): Promise<{ id: string; authUrl: string | undefined }> {
+    const authProvider = new DurableObjectOAuthClientProvider(
+      this.ctx.storage,
+      this.name,
+      callbackUrl
+    );
+
+    if (reconnect) {
+      authProvider.serverId = reconnect.id;
+      if (reconnect.oauthClientId) {
+        authProvider.clientId = reconnect.oauthClientId;
+      }
+    }
+
+    // allows passing through transport headers if necessary
+    // this handles some non-standard bearer auth setups (i.e. MCP server behind CF access instead of OAuth)
+    let headerTransportOpts: SSEClientTransportOptions = {};
+    if (options?.transport?.headers) {
+      headerTransportOpts = {
+        eventSourceInit: {
+          fetch: (url, init) =>
+            fetch(url, {
+              ...init,
+              headers: options?.transport?.headers,
+            }),
+        },
+        requestInit: {
+          headers: options?.transport?.headers,
+        },
+      };
+    }
+
+    const { id, authUrl, clientId } = await this.mcp.connect(url, {
+      client: options?.client,
+      reconnect,
+      transport: {
+        ...headerTransportOpts,
+        authProvider,
+      },
+    });
+
+    this.sql`
+      INSERT OR REPLACE INTO cf_agents_mcp_servers (id, name, server_url, client_id, auth_url, callback_url, server_options)
+      VALUES (
+        ${id},
+        ${serverName},
+        ${url},
+        ${clientId ?? null},
+        ${authUrl ?? null},
+        ${callbackUrl},
+        ${options ? JSON.stringify(options) : null}
+      );
+    `;
+
+    return {
+      authUrl,
+      id,
+    };
+  }
+
+  async removeMcpServer(id: string) {
+    this.mcp.closeConnection(id);
+    this.sql`
+      DELETE FROM cf_agents_mcp_servers WHERE id = ${id};
+    `;
+    this.broadcast(
+      JSON.stringify({
+        mcp: this.getMcpServers(),
+        type: "cf_agent_mcp_servers",
+      })
+    );
+  }
+
+  getMcpServers(): MCPServersState {
+    const mcpState: MCPServersState = {
+      prompts: this.mcp.listPrompts(),
+      resources: this.mcp.listResources(),
+      servers: {},
+      tools: this.mcp.listTools(),
+    };
+
+    const servers = this.sql<MCPServerRow>`
+      SELECT id, name, server_url, client_id, auth_url, callback_url, server_options FROM cf_agents_mcp_servers;
+    `;
+
+    for (const server of servers) {
+      const serverConn = this.mcp.mcpConnections[server.id];
+      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,
+        // mark as "authenticating" because the server isn't automatically connected, so it's pending authenticating
+        state: serverConn?.connectionState ?? "authenticating",
+      };
+    }
+
+    return mcpState;
+  }
 }
 
 /**
@@ -726,9 +1128,9 @@ export async function routeAgentRequest<Env>(
   const corsHeaders =
     options?.cors === true
       ? {
-          "Access-Control-Allow-Origin": "*",
-          "Access-Control-Allow-Methods": "GET, POST, HEAD, OPTIONS",
           "Access-Control-Allow-Credentials": "true",
+          "Access-Control-Allow-Methods": "GET, POST, HEAD, OPTIONS",
+          "Access-Control-Allow-Origin": "*",
           "Access-Control-Max-Age": "86400",
         }
       : options?.cors;
@@ -756,7 +1158,8 @@ export async function routeAgentRequest<Env>(
   if (
     response &&
     corsHeaders &&
-    request.headers.get("upgrade") !== "websocket"
+    request.headers.get("upgrade")?.toLowerCase() !== "websocket" &&
+    request.headers.get("Upgrade")?.toLowerCase() !== "websocket"
   ) {
     response = new Response(response.body, {
       headers: {
@@ -775,9 +1178,9 @@ export async function routeAgentRequest<Env>(
  * @param options Routing options
  */
 export async function routeAgentEmail<Env>(
-  email: ForwardableEmailMessage,
-  env: Env,
-  options?: AgentOptions<Env>
+  _email: ForwardableEmailMessage,
+  _env: Env,
+  _options?: AgentOptions<Env>
 ): Promise<void> {}
 
 /**
@@ -789,7 +1192,7 @@ export async function routeAgentEmail<Env>(
  * @param options Options for Agent creation
  * @returns Promise resolving to an Agent instance stub
  */
-export function getAgentByName<Env, T extends Agent<Env>>(
+export async function getAgentByName<Env, T extends Agent<Env>>(
   namespace: AgentNamespace<T>,
   name: string,
   options?: {
@@ -804,13 +1207,13 @@ export function getAgentByName<Env, T extends Agent<Env>>(
  * A wrapper for streaming responses in callable methods
  */
 export class StreamingResponse {
-  #connection: Connection;
-  #id: string;
-  #closed = false;
+  private _connection: Connection;
+  private _id: string;
+  private _closed = false;
 
   constructor(connection: Connection, id: string) {
-    this.#connection = connection;
-    this.#id = id;
+    this._connection = connection;
+    this._id = id;
   }
 
   /**
@@ -818,17 +1221,17 @@ export class StreamingResponse {
    * @param chunk The data to send
    */
   send(chunk: unknown) {
-    if (this.#closed) {
+    if (this._closed) {
       throw new Error("StreamingResponse is already closed");
     }
     const response: RPCResponse = {
-      type: "rpc",
-      id: this.#id,
-      success: true,
-      result: chunk,
       done: false,
+      id: this._id,
+      result: chunk,
+      success: true,
+      type: "rpc",
     };
-    this.#connection.send(JSON.stringify(response));
+    this._connection.send(JSON.stringify(response));
   }
 
   /**
@@ -836,17 +1239,17 @@ export class StreamingResponse {
    * @param finalChunk Optional final chunk of data to send
    */
   end(finalChunk?: unknown) {
-    if (this.#closed) {
+    if (this._closed) {
       throw new Error("StreamingResponse is already closed");
     }
-    this.#closed = true;
+    this._closed = true;
     const response: RPCResponse = {
-      type: "rpc",
-      id: this.#id,
-      success: true,
-      result: finalChunk,
       done: true,
+      id: this._id,
+      result: finalChunk,
+      success: true,
+      type: "rpc",
     };
-    this.#connection.send(JSON.stringify(response));
+    this._connection.send(JSON.stringify(response));
   }
 }