fix: register name on the stub in the transport

Author: mattzcarey

packages/agents/src/index.ts

@@ -1732,19 +1732,19 @@ export class Agent<
   ): Promise<MCPTransportOptions> {
     const { serverName, namespace, options } = config;
 
-    const doId = namespace.idFromName(`rpc:${serverName}`);
+    const doName = `rpc:${serverName}`;
+    const doId = namespace.idFromName(doName);
     const stub = namespace.get(doId) as unknown as DurableObjectStub<T>;
 
-    await stub.setName(`rpc:${serverName}`);
-
     if (options?.transport?.props) {
       await stub.updateProps(options.transport.props);
     }
 
     return {
       type: "rpc",
       stub,
-      functionName: options?.transport?.functionName
+      functionName: options?.transport?.functionName,
+      doName
     };
   }
 

packages/agents/src/mcp/index.ts

@@ -388,9 +388,21 @@ export abstract class McpAgent<
     message: JSONRPCMessage
   ): Promise<JSONRPCMessage | JSONRPCMessage[] | undefined> {
     if (!this._transport) {
+      this.props = await this.ctx.storage.get("props");
+
+      // Re-run init() to register tools on the server
+      await this.init();
       const server = await this.server;
-      this._transport = new RPCServerTransport(this.getRpcTransportOptions());
+
+      this._transport = this.initTransport();
+
+      if (!this._transport) {
+        throw new Error("Failed to initialize transport");
+      }
       await server.connect(this._transport);
+
+      // Reinitialize the server with any stored initialize request
+      await this.reinitializeServer();
     }
 
     if (!(this._transport instanceof RPCServerTransport)) {

packages/agents/src/mcp/rpc.ts

@@ -171,16 +171,19 @@ export interface MCPStub {
   handleMcpMessage(
     message: JSONRPCMessage | JSONRPCMessage[]
   ): Promise<JSONRPCMessage | JSONRPCMessage[] | undefined>;
+  setName(name: string): Promise<void>;
 }
 
 export interface RPCClientTransportOptions {
   stub: MCPStub;
   functionName?: string;
+  doName?: string;
 }
 
 export class RPCClientTransport implements Transport {
   private _stub: MCPStub;
   private _functionName: string;
+  private _doName?: string;
   private _started = false;
   private _protocolVersion?: string;
 
@@ -192,6 +195,7 @@ export class RPCClientTransport implements Transport {
   constructor(options: RPCClientTransportOptions) {
     this._stub = options.stub;
     this._functionName = options.functionName ?? "handleMcpMessage";
+    this._doName = options.doName;
   }
 
   setProtocolVersion(version: string): void {
@@ -229,6 +233,11 @@ export class RPCClientTransport implements Transport {
       validateJSONRPCMessage(message);
     }
 
+    // Set the name if the stub is a DO
+    if (this._doName && "setName" in this._stub) {
+      await this._stub.setName(this._doName);
+    }
+
     try {
       const result =
         await this._stub[this._functionName as keyof MCPStub](message);
@@ -425,11 +434,6 @@ export class RPCServerTransport implements Transport {
       // Use queueMicrotask to allow additional send() calls to accumulate
       // Resolver is reused for concurrent sends within the same tick
       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}`
-      );
     }
   }