remove breaking change with async storage apis. (#676)

* refactor: simplify MCPClientManager storage implementation and fix initialization order

- Changed MCPClientManager to use DurableObjectStorage directly instead of abstraction layer
- Removed MCPClientStorage and OAuthClientStorage interfaces
- Removed AgentMCPClientStorage adapter class
- Fixed initialization order: MCPClientManager now created AFTER database tables to prevent table-not-found errors during DO restart
- Changed getMcpServers() from async to sync method
- Added defensive checks for storage

* add return types to oauth client provider

Author: mattzcarey

.changeset/plenty-friends-hope.md

@@ -1,24 +1,20 @@
 ---
-"agents": minor
+"agents": patch
 ---
 
-### Breaking Changes
-
-- **`getMcpServers()` is now async**: Changed from synchronous to asynchronous method to support storage operations
-- **`DurableObjectOAuthClientProvider` constructor**: Now accepts `OAuthClientStorage` interface instead of `DurableObjectStorage`
-
 ### New Features
 
 - **`MCPClientManager` API changes**:
   - New `registerServer()` method to register servers (replaces part of `connect()`)
   - New `connectToServer()` method to establish connection (replaces part of `connect()`)
   - `connect()` method deprecated (still works for backward compatibility)
-  - Requires `MCPClientStorage` interface implementation (provided via `AgentMCPClientStorage`)
-- **Storage abstraction layer**: New `MCPClientStorage` and `OAuthClientStorage` interfaces enable custom storage implementations beyond Durable Objects
 - **Connection state observability**: New `onServerStateChanged()` event for tracking all server state changes
 - **Improved reconnect logic**: `restoreConnectionsFromStorage()` handles failed connections
 
 ### Bug Fixes
 
 - Fixed failed connections not being recreated on restore
 - Fixed redundant storage operations during connection restoration
+- Fixed potential OAuth storage initialization issue by excluding non-serializable authProvider from stored server options
+- Added defensive checks for storage initialization in MCPClientManager and DurableObjectOAuthClientProvider
+- Fixed initialization order: MCPClientManager is now created AFTER database tables are created to prevent possible table-not-found errors during DO restart

packages/agents/src/index.ts

@@ -26,7 +26,6 @@ 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 { AgentMCPClientStorage } from "./mcp/client-storage";
 import { genericObservability, type Observability } from "./observability";
 import { DisposableStore } from "./core/events";
 import { MessageType } from "./ai-types";
@@ -404,33 +403,12 @@ export class Agent<
   constructor(ctx: AgentContext, env: Env) {
     super(ctx, env);
 
-    this.mcp = new MCPClientManager(this._ParentClass.name, "0.0.1", {
-      storage: new AgentMCPClientStorage(
-        this.sql.bind(this),
-        this.ctx.storage.kv
-      )
-    });
-
     if (!wrappedClasses.has(this.constructor)) {
       // Auto-wrap custom methods with agent context
       this._autoWrapCustomMethods();
       wrappedClasses.add(this.constructor);
     }
 
-    // Broadcast server state whenever MCP state changes (register, connect, OAuth, remove, etc.)
-    this._disposables.add(
-      this.mcp.onServerStateChanged(async () => {
-        await 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_mcp_servers (
           id TEXT PRIMARY KEY NOT NULL,
@@ -472,6 +450,25 @@ export class Agent<
       )
     `;
 
+    // Initialize MCPClientManager AFTER tables are created
+    this.mcp = new MCPClientManager(this._ParentClass.name, "0.0.1", {
+      storage: this.ctx.storage
+    });
+
+    // Broadcast server state whenever MCP state changes (register, connect, OAuth, remove, etc.)
+    this._disposables.add(
+      this.mcp.onServerStateChanged(async () => {
+        this.broadcastMcpServers();
+      })
+    );
+
+    // Emit MCP observability events
+    this._disposables.add(
+      this.mcp.onObservabilityEvent((event) => {
+        this.observability?.emit(event);
+      })
+    );
+
     const _onRequest = this.onRequest.bind(this);
     this.onRequest = (request: Request) => {
       return agentContext.run(
@@ -603,7 +600,7 @@ export class Agent<
 
           connection.send(
             JSON.stringify({
-              mcp: await this.getMcpServers(),
+              mcp: this.getMcpServers(),
               type: MessageType.CF_AGENT_MCP_SERVERS
             })
           );
@@ -637,7 +634,7 @@ export class Agent<
         async () => {
           await this._tryCatch(async () => {
             await this.mcp.restoreConnectionsFromStorage(this.name);
-            await this.broadcastMcpServers();
+            this.broadcastMcpServers();
             return _onStart(props);
           });
         }
@@ -1409,7 +1406,7 @@ export class Agent<
     const id = nanoid(8);
 
     const authProvider = new DurableObjectOAuthClientProvider(
-      this.ctx.storage.kv,
+      this.ctx.storage,
       this.name,
       callbackUrl
     );
@@ -1463,18 +1460,18 @@ export class Agent<
     if (this.mcp.mcpConnections[id]) {
       await this.mcp.closeConnection(id);
     }
-    await this.mcp.removeServer(id);
+    this.mcp.removeServer(id);
   }
 
-  async getMcpServers(): Promise<MCPServersState> {
+  getMcpServers(): MCPServersState {
     const mcpState: MCPServersState = {
       prompts: this.mcp.listPrompts(),
       resources: this.mcp.listResources(),
       servers: {},
       tools: this.mcp.listTools()
     };
 
-    const servers = await this.mcp.listServers();
+    const servers = this.mcp.listServers();
 
     if (servers && Array.isArray(servers) && servers.length > 0) {
       for (const server of servers) {
@@ -1501,10 +1498,10 @@ export class Agent<
     return mcpState;
   }
 
-  private async broadcastMcpServers() {
+  private broadcastMcpServers() {
     this.broadcast(
       JSON.stringify({
-        mcp: await this.getMcpServers(),
+        mcp: this.getMcpServers(),
         type: MessageType.CF_AGENT_MCP_SERVERS
       })
     );
@@ -1527,7 +1524,7 @@ export class Agent<
     request: Request
   ): Promise<Response | null> {
     // Check if this is an OAuth callback request
-    const isCallback = await this.mcp.isCallbackRequest(request);
+    const isCallback = this.mcp.isCallbackRequest(request);
     if (!isCallback) {
       return null;
     }

packages/agents/src/mcp/client-storage.ts

@@ -10,130 +10,3 @@ export type MCPServerRow = {
   callback_url: string;
   server_options: string | null;
 };
-
-/**
- * KV storage interface for OAuth-related data
- * Used by OAuth providers to store tokens, client info, etc.
- */
-export interface OAuthClientStorage {
-  /**
-   * Get a value from key-value storage (for OAuth data like tokens, client info, etc.)
-   */
-  get<T>(key: string): Promise<T | undefined> | undefined;
-
-  /**
-   * Put a value into key-value storage (for OAuth data like tokens, client info, etc.)
-   */
-  put(key: string, value: unknown): Promise<void> | void;
-}
-
-/**
- * Storage interface for MCP client manager
- * Abstracts storage operations to decouple from specific storage implementations
- */
-export interface MCPClientStorage extends OAuthClientStorage {
-  /**
-   * Save or update an MCP server configuration
-   */
-  saveServer(server: MCPServerRow): Promise<void>;
-
-  /**
-   * Remove an MCP server from storage
-   */
-  removeServer(serverId: string): Promise<void>;
-
-  /**
-   * List all MCP servers from storage
-   */
-  listServers(): Promise<MCPServerRow[]>;
-
-  /**
-   * Get an MCP server by its callback URL
-   * Used during OAuth callback to identify which server is being authenticated
-   */
-  getServerByCallbackUrl(callbackUrl: string): Promise<MCPServerRow | null>;
-
-  /**
-   * Clear auth_url after successful OAuth authentication
-   * This prevents the agent from continuously asking for OAuth on reconnect
-   * when stored tokens are still valid.
-   */
-  clearAuthUrl(serverId: string): Promise<void>;
-}
-
-/**
- * SQL-based storage adapter that wraps SQL operations
- * Used by Agent class to provide SQL access to MCPClientManager
- */
-export class AgentMCPClientStorage implements MCPClientStorage {
-  constructor(
-    private sql: <T extends Record<string, unknown>>(
-      strings: TemplateStringsArray,
-      ...values: (string | number | boolean | null)[]
-    ) => T[],
-    private kv: SyncKvStorage
-  ) {}
-
-  async saveServer(server: MCPServerRow) {
-    this.sql`
-      INSERT OR REPLACE INTO cf_agents_mcp_servers (
-        id,
-        name,
-        server_url,
-        client_id,
-        auth_url,
-        callback_url,
-        server_options
-      )
-      VALUES (
-        ${server.id},
-        ${server.name},
-        ${server.server_url},
-        ${server.client_id ?? null},
-        ${server.auth_url ?? null},
-        ${server.callback_url},
-        ${server.server_options ?? null}
-      )
-    `;
-  }
-
-  async removeServer(serverId: string) {
-    this.sql`
-      DELETE FROM cf_agents_mcp_servers WHERE id = ${serverId}
-    `;
-  }
-
-  async listServers() {
-    const servers = this.sql<MCPServerRow>`
-      SELECT id, name, server_url, client_id, auth_url, callback_url, server_options
-      FROM cf_agents_mcp_servers
-    `;
-    return servers;
-  }
-
-  async getServerByCallbackUrl(callbackUrl: string) {
-    const results = this.sql<MCPServerRow>`
-      SELECT id, name, server_url, client_id, auth_url, callback_url, server_options
-      FROM cf_agents_mcp_servers
-      WHERE callback_url = ${callbackUrl}
-      LIMIT 1
-    `;
-    return results.length > 0 ? results[0] : null;
-  }
-
-  async clearAuthUrl(serverId: string) {
-    this.sql`
-      UPDATE cf_agents_mcp_servers
-      SET auth_url = NULL
-      WHERE id = ${serverId}
-    `;
-  }
-
-  async get<T>(key: string) {
-    return this.kv.get<T>(key);
-  }
-
-  async put(key: string, value: unknown) {
-    return this.kv.put(key, value);
-  }
-}