fix: mega typing

Author: mattzcarey

docs/mcp-transports.md

@@ -306,14 +306,50 @@ The RPC transport is minimal by design (~350 lines) and fully supports:
 - Notifications (messages without `id` field)
 - Automatic reconnection after Durable Object hibernation
 
+### Configuring RPC Transport Server Timeout
+
+The RPC transport has a configurable timeout for waiting for tool responses. By default, the server will wait **60 seconds** for a tool handler to call `send()`. You can customize this by overriding the `getRpcTransportOptions()` method in your `McpAgent`:
+
+```typescript
+export class MyMCP extends McpAgent<Env, State, {}> {
+  server = new McpServer({
+    name: "MyMCP",
+    version: "1.0.0"
+  });
+
+  // Configure RPC transport timeout
+  protected getRpcTransportOptions() {
+    return {
+      timeout: 120000 // 2 minutes (default is 60000)
+    };
+  }
+
+  async init() {
+    this.server.tool(
+      "long-running-task",
+      "A tool that takes a while to complete",
+      { input: z.string() },
+      async ({ input }) => {
+        // This tool has up to 2 minutes to complete
+        await longRunningOperation(input);
+        return {
+          content: [{ type: "text", text: "Task completed" }]
+        };
+      }
+    );
+  }
+}
+```
+
+The timeout ensures that if a tool handler fails to respond (e.g., due to an infinite loop or forgotten `send()` call), the request will fail with a clear timeout error rather than hanging indefinitely.
+
 ### Advanced: Custom RPC function names
 
 By default, the RPC transport calls the `handleMcpMessage` function. You can customize this:
 
 ```typescript
 await this.addMcpServer("my-server", "MyMCP", {
-  transport: { type: "rpc" },
-  functionName: "customHandler"
+  transport: { type: "rpc", functionName: "customHandler" }
 });
 ```
 

packages/agents/src/index.ts

@@ -37,7 +37,11 @@ import type {
 import { genericObservability, type Observability } from "./observability";
 import { DisposableStore } from "./core/events";
 import { MessageType } from "./ai-types";
-import type { McpAgent } from "./mcp";
+import type {
+  McpAgent,
+  RpcConnectionOptions,
+  HttpConnectionOptions
+} from "./mcp";
 
 export type { Connection, ConnectionContext, WSMessage } from "partyserver";
 
@@ -1470,12 +1474,7 @@ export class Agent<
   >(
     serverName: string,
     binding: DurableObjectNamespace<T> | string,
-    options: {
-      transport: { type: "rpc" };
-      functionName?: string;
-      client?: ConstructorParameters<typeof Client>[1];
-      props?: T extends McpAgent<unknown, unknown, infer Props> ? Props : never;
-    }
+    options: RpcConnectionOptions<T>
   ): Promise<{ id: string; authUrl: undefined }>;
 
   // Overload for HTTP/SSE transport
@@ -1484,38 +1483,17 @@ export class Agent<
     url: string,
     callbackHost?: string,
     agentsPrefix?: string,
-    options?: {
-      client?: ConstructorParameters<typeof Client>[1];
-      transport?: {
-        headers?: HeadersInit;
-        type?: Exclude<TransportType, "rpc">;
-      };
-    }
+    options?: HttpConnectionOptions
   ): Promise<{ id: string; authUrl: string | undefined }>;
 
   async addMcpServer<
     T extends McpAgent<unknown, unknown, Record<string, unknown>> = McpAgent
   >(
     serverName: string,
     urlOrBinding: string | DurableObjectNamespace<T>,
-    callbackHostOrOptions?:
-      | string
-      | {
-          transport: { type: "rpc" };
-          functionName?: string;
-          client?: ConstructorParameters<typeof Client>[1];
-          props?: T extends McpAgent<unknown, unknown, infer Props>
-            ? Props
-            : never;
-        },
+    callbackHostOrOptions?: string | RpcConnectionOptions<T>,
     agentsPrefix = "agents",
-    options?: {
-      client?: ConstructorParameters<typeof Client>[1];
-      transport?: {
-        headers?: HeadersInit;
-        type?: TransportType;
-      };
-    }
+    options?: HttpConnectionOptions
   ): Promise<{ id: string; authUrl: string | undefined }> {
     // Determine if this is RPC or HTTP based on parameters
     const isRpc =
@@ -1524,14 +1502,7 @@ export class Agent<
 
     if (isRpc) {
       // RPC transport path
-      const rpcOptions = callbackHostOrOptions as {
-        transport: { type: "rpc" };
-        functionName?: string;
-        client?: ConstructorParameters<typeof Client>[1];
-        props?: T extends McpAgent<unknown, unknown, infer Props>
-          ? Props
-          : never;
-      };
+      const rpcOptions = callbackHostOrOptions as RpcConnectionOptions<T>;
 
       const namespace = this._resolveRpcBinding<T>(
         urlOrBinding as DurableObjectNamespace<T> | string
@@ -1551,11 +1522,7 @@ export class Agent<
         serverName,
         namespace,
         url,
-        options: {
-          functionName: rpcOptions.functionName,
-          client: rpcOptions.client,
-          props: rpcOptions.props
-        },
+        options: rpcOptions,
         reconnect
       });
 
@@ -1770,14 +1737,14 @@ export class Agent<
 
     await stub.setName(`rpc:${serverName}`);
 
-    if (options?.props) {
-      await stub.updateProps(options.props);
+    if (options?.transport?.props) {
+      await stub.updateProps(options.transport.props);
     }
 
     return {
       type: "rpc",
       stub,
-      functionName: options?.functionName
+      functionName: options?.transport?.functionName
     };
   }
 

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

@@ -34,7 +34,12 @@ import {
 import { SSEEdgeClientTransport } from "./sse-edge";
 import { StreamableHTTPEdgeClientTransport } from "./streamable-http-edge";
 import { RPCClientTransport, type RPCClientTransportOptions } from "./rpc";
-import type { BaseTransportType, TransportType } from "./types";
+import type {
+  BaseTransportType,
+  TransportType,
+  McpClientOptions,
+  HttpTransportType
+} from "./types";
 
 /**
  * Connection state for MCP client connections
@@ -46,12 +51,18 @@ export type MCPConnectionState =
   | "discovering"
   | "failed";
 
+/**
+ * Transport options for MCP client connections.
+ * Combines transport-specific options with auth provider and type selection.
+ */
 export type MCPTransportOptions = (
   | SSEClientTransportOptions
   | StreamableHTTPClientTransportOptions
   | RPCClientTransportOptions
 ) & {
+  /** Optional OAuth provider for authenticating with the MCP server */
   authProvider?: AgentsOAuthProvider;
+  /** The transport type to use. "auto" will try streamable-http, then fall back to SSE */
   type?: TransportType;
 };
 
@@ -75,7 +86,7 @@ export class MCPClientConnection {
     info: ConstructorParameters<typeof Client>[0],
     public options: {
       transport: MCPTransportOptions;
-      client: ConstructorParameters<typeof Client>[1];
+      client: McpClientOptions;
     } = { client: {}, transport: {} }
   ) {
     const clientOptions = {
@@ -144,10 +155,7 @@ export class MCPClientConnection {
       throw new Error("Transport type must be specified");
     }
 
-    const finishAuth = async (base: BaseTransportType) => {
-      if (base === "rpc") {
-        throw new Error("RPC transport does not support authentication");
-      }
+    const finishAuth = async (base: HttpTransportType) => {
       const transport = this.getTransport(base);
       if (
         "finishAuth" in transport &&

packages/agents/src/mcp/index.ts

@@ -20,7 +20,7 @@ import {
   MCP_MESSAGE_HEADER
 } from "./utils";
 import { McpSSETransport, StreamableHTTPServerTransport } from "./transport";
-import { RPCServerTransport } from "./rpc";
+import { RPCServerTransport, type RPCServerTransportOptions } from "./rpc";
 
 export abstract class McpAgent<
   Env = unknown,
@@ -88,6 +88,23 @@ export abstract class McpAgent<
     return websockets[0];
   }
 
+  /**
+   * Returns options for configuring the RPC server transport.
+   * Override this method to customize RPC transport behavior (e.g., timeout).
+   *
+   * @example
+   * ```typescript
+   * class MyMCP extends McpAgent {
+   *   protected getRpcTransportOptions() {
+   *     return { timeout: 120000 }; // 2 minutes
+   *   }
+   * }
+   * ```
+   */
+  protected getRpcTransportOptions(): RPCServerTransportOptions {
+    return {};
+  }
+
   /** Returns a new transport matching the type of the Agent. */
   private initTransport() {
     switch (this.getTransportType()) {
@@ -98,7 +115,7 @@ export abstract class McpAgent<
         return new StreamableHTTPServerTransport({});
       }
       case "rpc": {
-        return new RPCServerTransport();
+        return new RPCServerTransport(this.getRpcTransportOptions());
       }
     }
   }
@@ -372,7 +389,7 @@ export abstract class McpAgent<
   ): Promise<JSONRPCMessage | JSONRPCMessage[] | undefined> {
     if (!this._transport) {
       const server = await this.server;
-      this._transport = new RPCServerTransport();
+      this._transport = new RPCServerTransport(this.getRpcTransportOptions());
       await server.connect(this._transport);
     }
 
@@ -489,3 +506,12 @@ export type {
   MCPClientOAuthResult,
   MCPClientOAuthCallbackConfig
 } from "./client";
+
+// Export connection configuration types
+export type {
+  RpcConnectionOptions,
+  HttpConnectionOptions,
+  McpClientOptions,
+  RpcTransportOptions,
+  HttpTransportOptions
+} from "./types";

packages/agents/src/mcp/rpc.ts

@@ -319,6 +319,21 @@ export interface RPCServerTransportOptions {
    * ```
    */
   onsessionclosed?: (sessionId: string) => void | Promise<void>;
+
+  /**
+   * Timeout in milliseconds for waiting for a response from the onmessage handler.
+   * If the handler doesn't call send() within this time, the request will fail with a timeout error.
+   *
+   * @default 60000 (60 seconds)
+   *
+   * @example
+   * ```typescript
+   * const transport = new RPCServerTransport({
+   *   timeout: 30000 // 30 seconds
+   * });
+   * ```
+   */
+  timeout?: number;
 }
 
 export class RPCServerTransport implements Transport {
@@ -331,6 +346,7 @@ export class RPCServerTransport implements Transport {
   private _initialized = false;
   private _onsessioninitialized?: (sessionId: string) => void | Promise<void>;
   private _onsessionclosed?: (sessionId: string) => void | Promise<void>;
+  private _timeout: number;
 
   sessionId?: string;
   onclose?: () => void;
@@ -341,6 +357,7 @@ export class RPCServerTransport implements Transport {
     this._sessionIdGenerator = options?.sessionIdGenerator;
     this._onsessioninitialized = options?.onsessioninitialized;
     this._onsessionclosed = options?.onsessionclosed;
+    this._timeout = options?.timeout ?? 60000; // Default 60 seconds
   }
 
   setProtocolVersion(version: string): void {
@@ -402,12 +419,19 @@ export class RPCServerTransport implements Transport {
       this._pendingResponse = [this._pendingResponse, message];
     }
 
-    // Resolve the promise on the next tick to allow multiple sends to accumulate
+    // Resolve the promise on the next tick to allow multiple send() calls to accumulate
+    // Note: Only the first send() triggers resolution; subsequent sends just accumulate
+    // until the microtask executes and handle() returns all responses
     if (this._responseResolver) {
       const resolver = this._responseResolver;
       this._responseResolver = null;
       // Use queueMicrotask to allow additional send() calls to accumulate
       queueMicrotask(() => resolver());
+    } else if (this._currentRequestId !== null) {
+      // This shouldn't happen - send() called after promise already resolved
+      console.warn(
+        `send() called after response already resolved for request ${this._currentRequestId}`
+      );
     }
   }
 
@@ -516,14 +540,40 @@ export class RPCServerTransport implements Transport {
     this._currentRequestId = (message as { id: string | number | null }).id;
 
     // Set up the promise before calling onmessage to handle race conditions
-    const responsePromise = new Promise<void>((resolve) => {
-      this._responseResolver = resolve;
+    let timeoutId: ReturnType<typeof setTimeout> | null = null;
+    const responsePromise = new Promise<void>((resolve, reject) => {
+      // Set up timeout
+      timeoutId = setTimeout(() => {
+        this._responseResolver = null;
+        reject(
+          new Error(
+            `Request timeout: No response received within ${this._timeout}ms for request ID ${this._currentRequestId}`
+          )
+        );
+      }, this._timeout);
+
+      // Wrap the resolver to clear timeout when response is received
+      this._responseResolver = () => {
+        if (timeoutId) {
+          clearTimeout(timeoutId);
+          timeoutId = null;
+        }
+        resolve();
+      };
     });
 
     this.onmessage?.(message);
 
     // Wait for a response using a promise that resolves when send() is called
-    await responsePromise;
+    try {
+      await responsePromise;
+    } catch (error) {
+      // Clean up on timeout
+      this._pendingResponse = null;
+      this._currentRequestId = null;
+      this._responseResolver = null;
+      throw error;
+    }
 
     const response = this._pendingResponse;
     this._pendingResponse = null;