agents 0.0.0-c5e3a32 → 0.0.0-c6d9bf1

package/src/index.ts

@@ -1,3 +1,4 @@
+import type { env } from "cloudflare:workers";
 import { AsyncLocalStorage } from "node:async_hooks";
 import type { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import type { SSEClientTransportOptions } from "@modelcontextprotocol/sdk/client/sse.js";
@@ -21,10 +22,13 @@ import {
   routePartykitRequest
 } from "partyserver";
 import { camelCaseToKebabCase } from "./client";
-import { MCPClientManager } from "./mcp/client";
-// import type { MCPClientConnection } from "./mcp/client-connection";
+import { MCPClientManager, type MCPClientOAuthResult } from "./mcp/client";
+import type { MCPConnectionState } from "./mcp/client-connection";
 import { DurableObjectOAuthClientProvider } from "./mcp/do-oauth-client-provider";
+import type { TransportType } from "./mcp/types";
 import { genericObservability, type Observability } from "./observability";
+import { DisposableStore } from "./core/events";
+import { MessageType } from "./ai-types";
 
 export type { Connection, ConnectionContext, WSMessage } from "partyserver";
 
@@ -42,7 +46,7 @@ export type RPCRequest = {
  * State update message from client
  */
 export type StateUpdateMessage = {
-  type: "cf_agent_state";
+  type: MessageType.CF_AGENT_STATE;
   state: unknown;
 };
 
@@ -50,7 +54,7 @@ export type StateUpdateMessage = {
  * RPC response message to client
  */
 export type RPCResponse = {
-  type: "rpc";
+  type: MessageType.RPC;
   id: string;
 } & (
   | {
@@ -77,7 +81,7 @@ function isRPCRequest(msg: unknown): msg is RPCRequest {
     typeof msg === "object" &&
     msg !== null &&
     "type" in msg &&
-    msg.type === "rpc" &&
+    msg.type === MessageType.RPC &&
     "id" in msg &&
     typeof msg.id === "string" &&
     "method" in msg &&
@@ -95,7 +99,7 @@ function isStateUpdateMessage(msg: unknown): msg is StateUpdateMessage {
     typeof msg === "object" &&
     msg !== null &&
     "type" in msg &&
-    msg.type === "cf_agent_state" &&
+    msg.type === MessageType.CF_AGENT_STATE &&
     "state" in msg
   );
 }
@@ -116,7 +120,7 @@ const callableMetadata = new Map<Function, CallableMetadata>();
  * Decorator that marks a method as callable by clients
  * @param metadata Optional metadata about the callable method
  */
-export function unstable_callable(metadata: CallableMetadata = {}) {
+export function callable(metadata: CallableMetadata = {}) {
   return function callableDecorator<This, Args extends unknown[], Return>(
     target: (this: This, ...args: Args) => Return,
     // biome-ignore lint/correctness/noUnusedFunctionParameters: later
@@ -130,6 +134,30 @@ export function unstable_callable(metadata: CallableMetadata = {}) {
   };
 }
 
+let didWarnAboutUnstableCallable = false;
+
+/**
+ * Decorator that marks a method as callable by clients
+ * @deprecated this has been renamed to callable, and unstable_callable will be removed in the next major version
+ * @param metadata Optional metadata about the callable method
+ */
+export const unstable_callable = (metadata: CallableMetadata = {}) => {
+  if (!didWarnAboutUnstableCallable) {
+    didWarnAboutUnstableCallable = true;
+    console.warn(
+      "unstable_callable is deprecated, use callable instead. unstable_callable will be removed in the next major version."
+    );
+  }
+  callable(metadata);
+};
+
+export type QueueItem<T = string> = {
+  id: string;
+  payload: T;
+  callback: keyof Agent<unknown>;
+  created_at: number;
+};
+
 /**
  * Represents a scheduled task within an Agent
  * @template T Type of the payload data
@@ -171,11 +199,13 @@ function getNextCronTime(cron: string) {
   return interval.getNextDate();
 }
 
+export type { TransportType } from "./mcp/types";
+
 /**
  * MCP Server state update message from server -> Client
  */
 export type MCPServerMessage = {
-  type: "cf_agent_mcp_servers";
+  type: MessageType.CF_AGENT_MCP_SERVERS;
   mcp: MCPServersState;
 };
 
@@ -195,7 +225,7 @@ export type MCPServer = {
   // 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";
+  state: MCPConnectionState;
   instructions: string | null;
   capabilities: ServerCapabilities | null;
 };
@@ -264,7 +294,13 @@ function withAgentContext<T extends (...args: any[]) => any>(
   method: T
 ): (this: Agent<unknown, unknown>, ...args: Parameters<T>) => ReturnType<T> {
   return function (...args: Parameters<T>): ReturnType<T> {
-    const { connection, request, email } = getCurrentAgent();
+    const { connection, request, email, agent } = getCurrentAgent();
+
+    if (agent === this) {
+      // already wrapped, so we can just call the method
+      return method.apply(this, args);
+    }
+    // not wrapped, so we need to wrap it
     return agentContext.run({ agent: this, connection, request, email }, () => {
       return method.apply(this, args);
     });
@@ -276,13 +312,21 @@ function withAgentContext<T extends (...args: any[]) => any>(
  * @template Env Environment type containing bindings
  * @template State State type to store within the Agent
  */
-export class Agent<Env, State = unknown> extends Server<Env> {
+export class Agent<
+  Env = typeof env,
+  State = unknown,
+  Props extends Record<string, unknown> = Record<string, unknown>
+> extends Server<Env, Props> {
   private _state = DEFAULT_STATE as State;
+  private _disposables = new DisposableStore();
 
   private _ParentClass: typeof Agent<Env, State> =
     Object.getPrototypeOf(this).constructor;
 
-  mcp: MCPClientManager = new MCPClientManager(this._ParentClass.name, "0.0.1");
+  readonly mcp: MCPClientManager = new MCPClientManager(
+    this._ParentClass.name,
+    "0.0.1"
+  );
 
   /**
    * Initial state for the Agent
@@ -375,6 +419,26 @@ export class Agent<Env, State = unknown> extends Server<Env> {
   constructor(ctx: AgentContext, env: Env) {
     super(ctx, env);
 
+    if (!wrappedClasses.has(this.constructor)) {
+      // Auto-wrap custom methods with agent context
+      this._autoWrapCustomMethods();
+      wrappedClasses.add(this.constructor);
+    }
+
+    // Broadcast server state after background connects (for OAuth servers)
+    this._disposables.add(
+      this.mcp.onConnected(async () => {
+        this.broadcastMcpServers();
+      })
+    );
+
+    // Emit MCP observability events
+    this._disposables.add(
+      this.mcp.onObservabilityEvent((event) => {
+        this.observability?.emit(event);
+      })
+    );
+
     this.sql`
       CREATE TABLE IF NOT EXISTS cf_agents_state (
         id TEXT PRIMARY KEY NOT NULL,
@@ -382,8 +446,14 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       )
     `;
 
-    // Auto-wrap custom methods with agent context
-    this._autoWrapCustomMethods();
+    this.sql`
+      CREATE TABLE IF NOT EXISTS cf_agents_queues (
+        id TEXT PRIMARY KEY NOT NULL,
+        payload TEXT,
+        callback TEXT,
+        created_at INTEGER DEFAULT (unixepoch())
+      )
+    `;
 
     void this.ctx.blockConcurrencyWhile(async () => {
       return this._tryCatch(async () => {
@@ -424,21 +494,24 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         { agent: this, connection: undefined, request, email: undefined },
         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"
-              })
-            );
+            const result = await this.mcp.handleCallbackRequest(request);
+            this.broadcastMcpServers();
+
+            if (result.authSuccess) {
+              // Start background connection if auth was successful
+              this.mcp
+                .establishConnection(result.serverId)
+                .catch((error) => {
+                  console.error("Background connection failed:", error);
+                })
+                .finally(() => {
+                  // Broadcast after background connection resolves (success/failure)
+                  this.broadcastMcpServers();
+                });
+            }
 
-            // 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
-            });
+            // Handle OAuth callback response using MCPClientManager configuration
+            return this.handleOAuthCallbackResponse(result, request);
           }
 
           return this._tryCatch(() => _onRequest(request));
@@ -499,10 +572,8 @@ export class Agent<Env, State = unknown> extends Server<Env> {
                   displayMessage: `RPC call to ${method}`,
                   id: nanoid(),
                   payload: {
-                    args,
                     method,
-                    streaming: metadata?.streaming,
-                    success: true
+                    streaming: metadata?.streaming
                   },
                   timestamp: Date.now(),
                   type: "rpc"
@@ -515,7 +586,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
                 id,
                 result,
                 success: true,
-                type: "rpc"
+                type: MessageType.RPC
               };
               connection.send(JSON.stringify(response));
             } catch (e) {
@@ -525,7 +596,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
                   e instanceof Error ? e.message : "Unknown error occurred",
                 id: parsed.id,
                 success: false,
-                type: "rpc"
+                type: MessageType.RPC
               };
               connection.send(JSON.stringify(response));
               console.error("RPC error:", e);
@@ -544,44 +615,42 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       // must fix this
       return agentContext.run(
         { agent: this, connection, request: ctx.request, email: undefined },
-        async () => {
-          setTimeout(() => {
-            if (this.state) {
-              connection.send(
-                JSON.stringify({
-                  state: this.state,
-                  type: "cf_agent_state"
-                })
-              );
-            }
-
+        () => {
+          if (this.state) {
             connection.send(
               JSON.stringify({
-                mcp: this.getMcpServers(),
-                type: "cf_agent_mcp_servers"
+                state: this.state,
+                type: MessageType.CF_AGENT_STATE
               })
             );
+          }
 
-            this.observability?.emit(
-              {
-                displayMessage: "Connection established",
-                id: nanoid(),
-                payload: {
-                  connectionId: connection.id
-                },
-                timestamp: Date.now(),
-                type: "connect"
+          connection.send(
+            JSON.stringify({
+              mcp: this.getMcpServers(),
+              type: MessageType.CF_AGENT_MCP_SERVERS
+            })
+          );
+
+          this.observability?.emit(
+            {
+              displayMessage: "Connection established",
+              id: nanoid(),
+              payload: {
+                connectionId: connection.id
               },
-              this.ctx
-            );
-            return this._tryCatch(() => _onConnect(connection, ctx));
-          }, 20);
+              timestamp: Date.now(),
+              type: "connect"
+            },
+            this.ctx
+          );
+          return this._tryCatch(() => _onConnect(connection, ctx));
         }
       );
     };
 
     const _onStart = this.onStart.bind(this);
-    this.onStart = async () => {
+    this.onStart = async (props?: Props) => {
       return agentContext.run(
         {
           agent: this,
@@ -590,15 +659,27 @@ export class Agent<Env, State = unknown> extends Server<Env> {
           email: undefined
         },
         async () => {
-          const servers = this.sql<MCPServerRow>`
+          await this._tryCatch(() => {
+            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
-          if (servers && Array.isArray(servers) && servers.length > 0) {
-            Promise.allSettled(
-              servers.map((server) => {
-                return this._connectToMcpServerInternal(
+            this.broadcastMcpServers();
+
+            // from DO storage, reconnect to all servers not currently in the oauth flow using our saved auth information
+            if (servers && Array.isArray(servers) && servers.length > 0) {
+              // Restore callback URLs for OAuth-enabled servers
+              servers.forEach((server) => {
+                if (server.callback_url) {
+                  // Register the full redirect URL including serverId to avoid ambiguous matches
+                  this.mcp.registerCallbackUrl(
+                    `${server.callback_url}/${server.id}`
+                  );
+                }
+              });
+
+              servers.forEach((server) => {
+                this._connectToMcpServerInternal(
                   server.name,
                   server.server_url,
                   server.callback_url,
@@ -609,18 +690,23 @@ export class Agent<Env, State = unknown> extends Server<Env> {
                     id: server.id,
                     oauthClientId: server.client_id ?? undefined
                   }
-                );
-              })
-            ).then((_results) => {
-              this.broadcast(
-                JSON.stringify({
-                  mcp: this.getMcpServers(),
-                  type: "cf_agent_mcp_servers"
-                })
-              );
-            });
-          }
-          await this._tryCatch(() => _onStart());
+                )
+                  .then(() => {
+                    // Broadcast updated MCP servers state after each server connects
+                    this.broadcastMcpServers();
+                  })
+                  .catch((error) => {
+                    console.error(
+                      `Error connecting to MCP server: ${server.name} (${server.server_url})`,
+                      error
+                    );
+                    // Still broadcast even if connection fails, so clients know about the failure
+                    this.broadcastMcpServers();
+                  });
+              });
+            }
+            return _onStart(props);
+          });
         }
       );
     };
@@ -630,7 +716,6 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     state: State,
     source: Connection | "server" = "server"
   ) {
-    const previousState = this._state;
     this._state = state;
     this.sql`
     INSERT OR REPLACE INTO cf_agents_state (id, state)
@@ -643,7 +728,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     this.broadcast(
       JSON.stringify({
         state: state,
-        type: "cf_agent_state"
+        type: MessageType.CF_AGENT_STATE
       }),
       source !== "server" ? [source.id] : []
     );
@@ -656,10 +741,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
             {
               displayMessage: "State updated",
               id: nanoid(),
-              payload: {
-                previousState,
-                state
-              },
+              payload: {},
               timestamp: Date.now(),
               type: "state:update"
             },
@@ -799,41 +881,37 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     while (proto && proto !== Object.prototype && depth < 10) {
       const methodNames = Object.getOwnPropertyNames(proto);
       for (const methodName of methodNames) {
-        // Skip if it's a private method or not a function
+        const descriptor = Object.getOwnPropertyDescriptor(proto, methodName);
+
+        // Skip if it's a private method, a base method, a getter, or not a function,
         if (
           baseMethods.has(methodName) ||
           methodName.startsWith("_") ||
-          typeof this[methodName as keyof this] !== "function"
+          !descriptor ||
+          !!descriptor.get ||
+          typeof descriptor.value !== "function"
         ) {
           continue;
         }
-        // If the method doesn't exist in base prototypes, it's a custom method
-        if (!baseMethods.has(methodName)) {
-          const descriptor = Object.getOwnPropertyDescriptor(proto, methodName);
-          if (descriptor && typeof descriptor.value === "function") {
-            // Wrap the custom method with context
-
-            const wrappedFunction = withAgentContext(
-              // biome-ignore lint/suspicious/noExplicitAny: I can't typescript
-              this[methodName as keyof this] as (...args: any[]) => any
-              // biome-ignore lint/suspicious/noExplicitAny: I can't typescript
-            ) as any;
-
-            // if the method is callable, copy the metadata from the original method
-            if (this._isCallable(methodName)) {
-              callableMetadata.set(
-                wrappedFunction,
-                callableMetadata.get(
-                  this[methodName as keyof this] as Function
-                )!
-              );
-            }
 
-            // set the wrapped function on the prototype
-            this.constructor.prototype[methodName as keyof this] =
-              wrappedFunction;
-          }
+        // Now, methodName is confirmed to be a custom method/function
+        // Wrap the custom method with context
+        const wrappedFunction = withAgentContext(
+          // biome-ignore lint/suspicious/noExplicitAny: I can't typescript
+          this[methodName as keyof this] as (...args: any[]) => any
+          // biome-ignore lint/suspicious/noExplicitAny: I can't typescript
+        ) as any;
+
+        // if the method is callable, copy the metadata from the original method
+        if (this._isCallable(methodName)) {
+          callableMetadata.set(
+            wrappedFunction,
+            callableMetadata.get(this[methodName as keyof this] as Function)!
+          );
         }
+
+        // set the wrapped function on the prototype
+        this.constructor.prototype[methodName as keyof this] = wrappedFunction;
       }
 
       proto = Object.getPrototypeOf(proto);
@@ -875,6 +953,131 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     throw new Error("Not implemented");
   }
 
+  /**
+   * Queue a task to be executed in the future
+   * @param payload Payload to pass to the callback
+   * @param callback Name of the method to call
+   * @returns The ID of the queued task
+   */
+  async queue<T = unknown>(callback: keyof this, payload: T): Promise<string> {
+    const id = nanoid(9);
+    if (typeof callback !== "string") {
+      throw new Error("Callback must be a string");
+    }
+
+    if (typeof this[callback] !== "function") {
+      throw new Error(`this.${callback} is not a function`);
+    }
+
+    this.sql`
+      INSERT OR REPLACE INTO cf_agents_queues (id, payload, callback)
+      VALUES (${id}, ${JSON.stringify(payload)}, ${callback})
+    `;
+
+    void this._flushQueue().catch((e) => {
+      console.error("Error flushing queue:", e);
+    });
+
+    return id;
+  }
+
+  private _flushingQueue = false;
+
+  private async _flushQueue() {
+    if (this._flushingQueue) {
+      return;
+    }
+    this._flushingQueue = true;
+    while (true) {
+      const result = this.sql<QueueItem<string>>`
+      SELECT * FROM cf_agents_queues
+      ORDER BY created_at ASC
+    `;
+
+      if (!result || result.length === 0) {
+        break;
+      }
+
+      for (const row of result || []) {
+        const callback = this[row.callback as keyof Agent<Env>];
+        if (!callback) {
+          console.error(`callback ${row.callback} not found`);
+          continue;
+        }
+        const { connection, request, email } = agentContext.getStore() || {};
+        await agentContext.run(
+          {
+            agent: this,
+            connection,
+            request,
+            email
+          },
+          async () => {
+            // TODO: add retries and backoff
+            await (
+              callback as (
+                payload: unknown,
+                queueItem: QueueItem<string>
+              ) => Promise<void>
+            ).bind(this)(JSON.parse(row.payload as string), row);
+            await this.dequeue(row.id);
+          }
+        );
+      }
+    }
+    this._flushingQueue = false;
+  }
+
+  /**
+   * Dequeue a task by ID
+   * @param id ID of the task to dequeue
+   */
+  async dequeue(id: string) {
+    this.sql`DELETE FROM cf_agents_queues WHERE id = ${id}`;
+  }
+
+  /**
+   * Dequeue all tasks
+   */
+  async dequeueAll() {
+    this.sql`DELETE FROM cf_agents_queues`;
+  }
+
+  /**
+   * Dequeue all tasks by callback
+   * @param callback Name of the callback to dequeue
+   */
+  async dequeueAllByCallback(callback: string) {
+    this.sql`DELETE FROM cf_agents_queues WHERE callback = ${callback}`;
+  }
+
+  /**
+   * Get a queued task by ID
+   * @param id ID of the task to get
+   * @returns The task or undefined if not found
+   */
+  async getQueue(id: string): Promise<QueueItem<string> | undefined> {
+    const result = this.sql<QueueItem<string>>`
+      SELECT * FROM cf_agents_queues WHERE id = ${id}
+    `;
+    return result
+      ? { ...result[0], payload: JSON.parse(result[0].payload) }
+      : undefined;
+  }
+
+  /**
+   * Get all queues by key and value
+   * @param key Key to filter by
+   * @param value Value to filter by
+   * @returns Array of matching QueueItem objects
+   */
+  async getQueues(key: string, value: string): Promise<QueueItem<string>[]> {
+    const result = this.sql<QueueItem<string>>`
+      SELECT * FROM cf_agents_queues
+    `;
+    return result.filter((row) => JSON.parse(row.payload)[key] === value);
+  }
+
   /**
    * Schedule a task to be executed in the future
    * @template T Type of the payload data
@@ -895,7 +1098,10 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         {
           displayMessage: `Schedule ${schedule.id} created`,
           id: nanoid(),
-          payload: schedule,
+          payload: {
+            callback: callback as string,
+            id: id
+          },
           timestamp: Date.now(),
           type: "schedule:create"
         },
@@ -1065,7 +1271,10 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         {
           displayMessage: `Schedule ${id} cancelled`,
           id: nanoid(),
-          payload: schedule,
+          payload: {
+            callback: schedule.callback,
+            id: schedule.id
+          },
           timestamp: Date.now(),
           type: "schedule:cancel"
         },
@@ -1081,9 +1290,9 @@ export class Agent<Env, State = unknown> extends Server<Env> {
   private async _scheduleNextAlarm() {
     // Find the next schedule that needs to be executed
     const result = this.sql`
-      SELECT time FROM cf_agents_schedules 
+      SELECT time FROM cf_agents_schedules
       WHERE time > ${Math.floor(Date.now() / 1000)}
-      ORDER BY time ASC 
+      ORDER BY time ASC
       LIMIT 1
     `;
     if (!result) return;
@@ -1130,7 +1339,10 @@ export class Agent<Env, State = unknown> extends Server<Env> {
                 {
                   displayMessage: `Schedule ${row.id} executed`,
                   id: nanoid(),
-                  payload: row,
+                  payload: {
+                    callback: row.callback,
+                    id: row.id
+                  },
                   timestamp: Date.now(),
                   type: "schedule:execute"
                 },
@@ -1177,10 +1389,13 @@ export class Agent<Env, State = unknown> extends Server<Env> {
     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`;
 
     // delete all alarms
     await this.ctx.storage.deleteAlarm();
     await this.ctx.storage.deleteAll();
+    this._disposables.dispose();
+    await this.mcp.dispose?.();
     this.ctx.abort("destroyed"); // enforce that the agent is evicted
 
     this.observability?.emit(
@@ -1206,25 +1421,42 @@ export class Agent<Env, State = unknown> extends Server<Env> {
   /**
    * Connect to a new MCP Server
    *
+   * @param serverName Name of the MCP server
    * @param url MCP Server SSE URL
-   * @param callbackHost Base host for the agent, used for the redirect URI.
+   * @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 (header) options
+   * @param options MCP client and transport options
    * @returns authUrl
    */
   async addMcpServer(
     serverName: string,
     url: string,
-    callbackHost: string,
+    callbackHost?: string,
     agentsPrefix = "agents",
     options?: {
       client?: ConstructorParameters<typeof Client>[1];
       transport?: {
-        headers: HeadersInit;
+        headers?: HeadersInit;
+        type?: TransportType;
       };
     }
   ): Promise<{ id: string; authUrl: string | undefined }> {
-    const callbackUrl = `${callbackHost}/${agentsPrefix}/${camelCaseToKebabCase(this._ParentClass.name)}/${this.name}/callback`;
+    // If callbackHost is not provided, derive it from the current request
+    let resolvedCallbackHost = callbackHost;
+    if (!resolvedCallbackHost) {
+      const { request } = getCurrentAgent();
+      if (!request) {
+        throw new Error(
+          "callbackHost is required when not called within a request context"
+        );
+      }
+
+      // Extract the origin from the request
+      const requestUrl = new URL(request.url);
+      resolvedCallbackHost = `${requestUrl.protocol}//${requestUrl.host}`;
+    }
+
+    const callbackUrl = `${resolvedCallbackHost}/${agentsPrefix}/${camelCaseToKebabCase(this._ParentClass.name)}/${this.name}/callback`;
 
     const result = await this._connectToMcpServerInternal(
       serverName,
@@ -1232,6 +1464,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       callbackUrl,
       options
     );
+
     this.sql`
         INSERT
         OR REPLACE INTO cf_agents_mcp_servers (id, name, server_url, client_id, auth_url, callback_url, server_options)
@@ -1246,17 +1479,12 @@ export class Agent<Env, State = unknown> extends Server<Env> {
         );
     `;
 
-    this.broadcast(
-      JSON.stringify({
-        mcp: this.getMcpServers(),
-        type: "cf_agent_mcp_servers"
-      })
-    );
+    this.broadcastMcpServers();
 
     return result;
   }
 
-  async _connectToMcpServerInternal(
+  private async _connectToMcpServerInternal(
     _serverName: string,
     url: string,
     callbackUrl: string,
@@ -1272,6 +1500,7 @@ export class Agent<Env, State = unknown> extends Server<Env> {
        */
       transport?: {
         headers?: HeadersInit;
+        type?: TransportType;
       };
     },
     reconnect?: {
@@ -1296,6 +1525,9 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       }
     }
 
+    // Use the transport type specified in options, or default to "auto"
+    const transportType: TransportType = options?.transport?.type ?? "auto";
+
     // 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 = {};
@@ -1319,7 +1551,8 @@ export class Agent<Env, State = unknown> extends Server<Env> {
       reconnect,
       transport: {
         ...headerTransportOpts,
-        authProvider
+        authProvider,
+        type: transportType
       }
     });
 
@@ -1332,15 +1565,11 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
   async removeMcpServer(id: string) {
     this.mcp.closeConnection(id);
+    this.mcp.unregisterCallbackUrl(id);
     this.sql`
       DELETE FROM cf_agents_mcp_servers WHERE id = ${id};
     `;
-    this.broadcast(
-      JSON.stringify({
-        mcp: this.getMcpServers(),
-        type: "cf_agent_mcp_servers"
-      })
-    );
+    this.broadcastMcpServers();
   }
 
   getMcpServers(): MCPServersState {
@@ -1372,8 +1601,53 @@ export class Agent<Env, State = unknown> extends Server<Env> {
 
     return mcpState;
   }
+
+  private broadcastMcpServers() {
+    this.broadcast(
+      JSON.stringify({
+        mcp: this.getMcpServers(),
+        type: MessageType.CF_AGENT_MCP_SERVERS
+      })
+    );
+  }
+
+  /**
+   * Handle OAuth callback response using MCPClientManager configuration
+   * @param result OAuth callback result
+   * @param request The original request (needed for base URL)
+   * @returns Response for the OAuth callback
+   */
+  private handleOAuthCallbackResponse(
+    result: MCPClientOAuthResult,
+    request: Request
+  ): Response {
+    const config = this.mcp.getOAuthCallbackConfig();
+
+    // Use custom handler if configured
+    if (config?.customHandler) {
+      return config.customHandler(result);
+    }
+
+    // Use redirect URLs if configured
+    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")}`
+      );
+    }
+
+    // Default behavior - redirect to base URL
+    const baseUrl = new URL(request.url).origin;
+    return Response.redirect(baseUrl);
+  }
 }
 
+// A set of classes that have been wrapped with agent context
+const wrappedClasses = new Set<typeof Agent.prototype.constructor>();
+
 /**
  * Namespace for creating Agent instances
  * @template Agentic Type of the Agent class
@@ -1550,6 +1824,13 @@ export type EmailRoutingOptions<Env> = AgentOptions<Env> & {
   resolver: EmailResolver<Env>;
 };
 
+// Cache the agent namespace map for email routing
+// This maps both kebab-case and original names to namespaces
+const agentMapCache = new WeakMap<
+  Record<string, unknown>,
+  Record<string, unknown>
+>();
+
 /**
  * Route an email to the appropriate Agent
  * @param email The email to route
@@ -1569,28 +1850,41 @@ export async function routeAgentEmail<Env>(
     return;
   }
 
-  const namespaceBinding = env[routingInfo.agentName as keyof Env];
-  if (!namespaceBinding) {
-    throw new Error(
-      `Agent namespace '${routingInfo.agentName}' not found in environment`
-    );
+  // Build a map that includes both original names and kebab-case versions
+  if (!agentMapCache.has(env as Record<string, unknown>)) {
+    const map: Record<string, unknown> = {};
+    for (const [key, value] of Object.entries(env as Record<string, unknown>)) {
+      if (
+        value &&
+        typeof value === "object" &&
+        "idFromName" in value &&
+        typeof value.idFromName === "function"
+      ) {
+        // Add both the original name and kebab-case version
+        map[key] = value;
+        map[camelCaseToKebabCase(key)] = value;
+      }
+    }
+    agentMapCache.set(env as Record<string, unknown>, map);
   }
 
-  // Type guard to check if this is actually a DurableObjectNamespace (AgentNamespace)
-  if (
-    typeof namespaceBinding !== "object" ||
-    !("idFromName" in namespaceBinding) ||
-    typeof namespaceBinding.idFromName !== "function"
-  ) {
+  const agentMap = agentMapCache.get(env as Record<string, unknown>)!;
+  const namespace = agentMap[routingInfo.agentName];
+
+  if (!namespace) {
+    // Provide helpful error message listing available agents
+    const availableAgents = Object.keys(agentMap)
+      .filter((key) => !key.includes("-")) // Show only original names, not kebab-case duplicates
+      .join(", ");
     throw new Error(
-      `Environment binding '${routingInfo.agentName}' is not an AgentNamespace (found: ${typeof namespaceBinding})`
+      `Agent namespace '${routingInfo.agentName}' not found in environment. Available agents: ${availableAgents}`
     );
   }
 
-  // Safe cast after runtime validation
-  const namespace = namespaceBinding as unknown as AgentNamespace<Agent<Env>>;
-
-  const agent = await getAgentByName(namespace, routingInfo.agentId);
+  const agent = await getAgentByName(
+    namespace as unknown as AgentNamespace<Agent<Env>>,
+    routingInfo.agentId
+  );
 
   // let's make a serialisable version of the email
   const serialisableEmail: AgentEmail = {
@@ -1669,12 +1963,17 @@ export type EmailSendOptions = {
  * @param options Options for Agent creation
  * @returns Promise resolving to an Agent instance stub
  */
-export async function getAgentByName<Env, T extends Agent<Env>>(
+export async function getAgentByName<
+  Env,
+  T extends Agent<Env>,
+  Props extends Record<string, unknown> = Record<string, unknown>
+>(
   namespace: AgentNamespace<T>,
   name: string,
   options?: {
     jurisdiction?: DurableObjectJurisdiction;
     locationHint?: DurableObjectLocationHint;
+    props?: Props;
   }
 ) {
   return getServerByName<Env, T>(namespace, name, options);
@@ -1706,7 +2005,7 @@ export class StreamingResponse {
       id: this._id,
       result: chunk,
       success: true,
-      type: "rpc"
+      type: MessageType.RPC
     };
     this._connection.send(JSON.stringify(response));
   }
@@ -1725,7 +2024,7 @@ export class StreamingResponse {
       id: this._id,
       result: finalChunk,
       success: true,
-      type: "rpc"
+      type: MessageType.RPC
     };
     this._connection.send(JSON.stringify(response));
   }