agents 0.0.0-90db5ba → 0.0.0-928211f

package/dist/mcp/index.js

@@ -1,55 +1,62 @@
 import {
   Agent
-} from "../chunk-XG52S6YY.js";
+} from "../chunk-2DJ6XAU6.js";
 import {
-  __privateAdd,
-  __privateGet,
-  __privateMethod,
-  __privateSet
-} from "../chunk-HMLY7DHA.js";
+  SSEEdgeClientTransport,
+  StreamableHTTPEdgeClientTransport
+} from "../chunk-XJR75HO7.js";
+import "../chunk-PVQZBKN7.js";
+import "../chunk-KUH345EY.js";
 
 // src/mcp/index.ts
 import { DurableObject } from "cloudflare:workers";
 import {
   InitializeRequestSchema,
+  JSONRPCMessageSchema,
   isJSONRPCError,
   isJSONRPCNotification,
   isJSONRPCRequest,
-  isJSONRPCResponse,
-  JSONRPCMessageSchema
+  isJSONRPCResponse
+} from "@modelcontextprotocol/sdk/types.js";
+import {
+  ElicitRequestSchema
 } from "@modelcontextprotocol/sdk/types.js";
 var MAXIMUM_MESSAGE_SIZE_BYTES = 4 * 1024 * 1024;
-function handleCORS(request, corsOptions) {
-  const origin = request.headers.get("Origin") || "*";
-  const corsHeaders = {
-    "Access-Control-Allow-Origin": corsOptions?.origin || origin,
-    "Access-Control-Allow-Methods": corsOptions?.methods || "GET, POST, OPTIONS",
-    "Access-Control-Allow-Headers": corsOptions?.headers || "Content-Type",
-    "Access-Control-Max-Age": (corsOptions?.maxAge || 86400).toString()
+function corsHeaders(_request, corsOptions = {}) {
+  const origin = "*";
+  return {
+    "Access-Control-Allow-Headers": corsOptions.headers || "Content-Type, mcp-session-id, mcp-protocol-version",
+    "Access-Control-Allow-Methods": corsOptions.methods || "GET, POST, OPTIONS",
+    "Access-Control-Allow-Origin": corsOptions.origin || origin,
+    "Access-Control-Expose-Headers": corsOptions.exposeHeaders || "mcp-session-id",
+    "Access-Control-Max-Age": (corsOptions.maxAge || 86400).toString()
   };
+}
+function isDurableObjectNamespace(namespace) {
+  return typeof namespace === "object" && namespace !== null && "newUniqueId" in namespace && typeof namespace.newUniqueId === "function" && "idFromName" in namespace && typeof namespace.idFromName === "function";
+}
+function handleCORS(request, corsOptions) {
   if (request.method === "OPTIONS") {
-    return new Response(null, { headers: corsHeaders });
+    return new Response(null, { headers: corsHeaders(request, corsOptions) });
   }
   return null;
 }
-var _getWebSocket, _started;
 var McpSSETransport = class {
   constructor(getWebSocket) {
-    __privateAdd(this, _getWebSocket);
-    __privateAdd(this, _started, false);
-    __privateSet(this, _getWebSocket, getWebSocket);
+    this._started = false;
+    this._getWebSocket = getWebSocket;
   }
   async start() {
-    if (__privateGet(this, _started)) {
+    if (this._started) {
       throw new Error("Transport already started");
     }
-    __privateSet(this, _started, true);
+    this._started = true;
   }
   async send(message) {
-    if (!__privateGet(this, _started)) {
+    if (!this._started) {
       throw new Error("Transport not started");
     }
-    const websocket = __privateGet(this, _getWebSocket).call(this);
+    const websocket = this._getWebSocket();
     if (!websocket) {
       throw new Error("WebSocket not connected");
     }
@@ -64,52 +71,40 @@ var McpSSETransport = class {
     this.onclose?.();
   }
 };
-_getWebSocket = new WeakMap();
-_started = new WeakMap();
-var _getWebSocketForGetRequest, _getWebSocketForMessageID, _notifyResponseIdSent, _started2;
 var McpStreamableHttpTransport = class {
   constructor(getWebSocketForMessageID, notifyResponseIdSent) {
-    // TODO: If there is an open connection to send server-initiated messages
-    // back, we should use that connection
-    __privateAdd(this, _getWebSocketForGetRequest);
-    // Get the appropriate websocket connection for a given message id
-    __privateAdd(this, _getWebSocketForMessageID);
-    // Notify the server that a response has been sent for a given message id
-    // so that it may clean up it's mapping of message ids to connections
-    // once they are no longer needed
-    __privateAdd(this, _notifyResponseIdSent);
-    __privateAdd(this, _started2, false);
-    __privateSet(this, _getWebSocketForMessageID, getWebSocketForMessageID);
-    __privateSet(this, _notifyResponseIdSent, notifyResponseIdSent);
-    __privateSet(this, _getWebSocketForGetRequest, () => null);
+    this._started = false;
+    this._getWebSocketForMessageID = getWebSocketForMessageID;
+    this._notifyResponseIdSent = notifyResponseIdSent;
+    this._getWebSocketForGetRequest = () => null;
   }
   async start() {
-    if (__privateGet(this, _started2)) {
+    if (this._started) {
       throw new Error("Transport already started");
     }
-    __privateSet(this, _started2, true);
+    this._started = true;
   }
   async send(message) {
-    if (!__privateGet(this, _started2)) {
+    if (!this._started) {
       throw new Error("Transport not started");
     }
     let websocket = null;
     if (isJSONRPCResponse(message) || isJSONRPCError(message)) {
-      websocket = __privateGet(this, _getWebSocketForMessageID).call(this, message.id.toString());
+      websocket = this._getWebSocketForMessageID(message.id.toString());
       if (!websocket) {
         throw new Error(
           `Could not find WebSocket for message id: ${message.id}`
         );
       }
     } else if (isJSONRPCRequest(message)) {
-      websocket = __privateGet(this, _getWebSocketForGetRequest).call(this);
+      websocket = this._getWebSocketForGetRequest();
     } else if (isJSONRPCNotification(message)) {
       websocket = null;
     }
     try {
       websocket?.send(JSON.stringify(message));
       if (isJSONRPCResponse(message)) {
-        __privateGet(this, _notifyResponseIdSent).call(this, message.id.toString());
+        this._notifyResponseIdSent(message.id.toString());
       }
     } catch (error) {
       this.onerror?.(error);
@@ -120,28 +115,16 @@ var McpStreamableHttpTransport = class {
     this.onclose?.();
   }
 };
-_getWebSocketForGetRequest = new WeakMap();
-_getWebSocketForMessageID = new WeakMap();
-_notifyResponseIdSent = new WeakMap();
-_started2 = new WeakMap();
-var _status, _transport, _transportType, _requestIdToConnectionId, _agent, _McpAgent_instances, initialize_fn;
-var _McpAgent = class _McpAgent extends DurableObject {
+var McpAgent = class _McpAgent extends DurableObject {
   constructor(ctx, env) {
     var _a;
     super(ctx, env);
-    __privateAdd(this, _McpAgent_instances);
-    __privateAdd(this, _status, "zero");
-    __privateAdd(this, _transport);
-    __privateAdd(this, _transportType, "unset");
-    __privateAdd(this, _requestIdToConnectionId, /* @__PURE__ */ new Map());
-    /**
-     * Since McpAgent's _aren't_ yet real "Agents", let's only expose a couple of the methods
-     * to the outer class: initialState/state/setState/onStateUpdate/sql
-     */
-    __privateAdd(this, _agent);
+    this._status = "zero";
+    this._transportType = "unset";
+    this._requestIdToConnectionId = /* @__PURE__ */ new Map();
     this.initRun = false;
     const self = this;
-    __privateSet(this, _agent, new (_a = class extends Agent {
+    this._agent = new (_a = class extends Agent {
       onStateUpdate(state, source) {
         return self.onStateUpdate(state, source);
       }
@@ -150,23 +133,66 @@ var _McpAgent = class _McpAgent extends DurableObject {
       }
     }, _a.options = {
       hibernate: true
-    }, _a)(ctx, env));
+    }, _a)(ctx, env);
+  }
+  get mcp() {
+    return this._agent.mcp;
   }
   get state() {
-    return __privateGet(this, _agent).state;
+    return this._agent.state;
   }
   sql(strings, ...values) {
-    return __privateGet(this, _agent).sql(strings, ...values);
+    return this._agent.sql(strings, ...values);
   }
   setState(state) {
-    return __privateGet(this, _agent).setState(state);
+    return this._agent.setState(state);
   }
+  /**
+   * Elicit user input with a message and schema
+   */
+  async elicitInput(params) {
+    const requestId = `elicit_${Math.random().toString(36).substring(2, 11)}`;
+    await this.ctx.storage.put(`elicitation:${requestId}`, {
+      message: params.message,
+      requestedSchema: params.requestedSchema,
+      timestamp: Date.now()
+    });
+    const elicitRequest = {
+      jsonrpc: "2.0",
+      id: requestId,
+      method: "elicitation/create",
+      params: {
+        message: params.message,
+        requestedSchema: params.requestedSchema
+      }
+    };
+    if (this._transport) {
+      await this._transport.send(elicitRequest);
+    } else {
+      const connections = this._agent?.getConnections();
+      if (!connections || Array.from(connections).length === 0) {
+        await this.ctx.storage.delete(`elicitation:${requestId}`);
+        throw new Error("No active connections available for elicitation");
+      }
+      const connectionList = Array.from(connections);
+      for (const connection of connectionList) {
+        try {
+          connection.send(JSON.stringify(elicitRequest));
+        } catch (error) {
+          console.error("Failed to send elicitation request:", error);
+        }
+      }
+    }
+    return this._waitForElicitationResponse(requestId);
+  }
+  // we leave the variables as unused for autocomplete purposes
+  // biome-ignore lint/correctness/noUnusedFunctionParameters: overriden later
   onStateUpdate(state, source) {
   }
   async onStart() {
     var _a;
     const self = this;
-    __privateSet(this, _agent, new (_a = class extends Agent {
+    this._agent = new (_a = class extends Agent {
       constructor() {
         super(...arguments);
         this.initialState = self.initialState;
@@ -179,39 +205,52 @@ var _McpAgent = class _McpAgent extends DurableObject {
       }
     }, _a.options = {
       hibernate: true
-    }, _a)(this.ctx, this.env));
+    }, _a)(this.ctx, this.env);
     this.props = await this.ctx.storage.get("props");
-    __privateSet(this, _transportType, await this.ctx.storage.get(
+    this._transportType = await this.ctx.storage.get(
       "transportType"
-    ));
-    this.init?.();
-    if (__privateGet(this, _transportType) === "sse") {
-      __privateSet(this, _transport, new McpSSETransport(() => this.getWebSocket()));
-      await this.server.connect(__privateGet(this, _transport));
-    } else if (__privateGet(this, _transportType) === "streamable-http") {
-      __privateSet(this, _transport, new McpStreamableHttpTransport(
+    );
+    await this._init(this.props);
+    const server = await this.server;
+    if (this._transportType === "sse") {
+      this._transport = new McpSSETransport(() => this.getWebSocket());
+      await server.connect(this._transport);
+    } else if (this._transportType === "streamable-http") {
+      this._transport = new McpStreamableHttpTransport(
         (id) => this.getWebSocketForResponseID(id),
-        (id) => __privateGet(this, _requestIdToConnectionId).delete(id)
-      ));
-      await this.server.connect(__privateGet(this, _transport));
+        (id) => this._requestIdToConnectionId.delete(id)
+      );
+      await server.connect(this._transport);
     }
   }
   async _init(props) {
     await this.ctx.storage.put("props", props ?? {});
-    await this.ctx.storage.put("transportType", "unset");
+    if (!this.ctx.storage.get("transportType")) {
+      await this.ctx.storage.put("transportType", "unset");
+    }
     this.props = props;
     if (!this.initRun) {
       this.initRun = true;
       await this.init();
     }
   }
-  isInitialized() {
-    return this.initRun;
+  async setInitialized() {
+    await this.ctx.storage.put("initialized", true);
+  }
+  async isInitialized() {
+    return await this.ctx.storage.get("initialized") === true;
+  }
+  async _initialize() {
+    await this.ctx.blockConcurrencyWhile(async () => {
+      this._status = "starting";
+      await this.onStart();
+      this._status = "started";
+    });
   }
   // Allow the worker to fetch a websocket connection to the agent
   async fetch(request) {
-    if (__privateGet(this, _status) !== "started") {
-      await __privateMethod(this, _McpAgent_instances, initialize_fn).call(this);
+    if (this._status !== "started") {
+      await this._initialize();
     }
     if (request.headers.get("Upgrade") !== "websocket") {
       return new Response("Expected WebSocket Upgrade request", {
@@ -220,6 +259,7 @@ var _McpAgent = class _McpAgent extends DurableObject {
     }
     const url = new URL(request.url);
     const path = url.pathname;
+    const server = await this.server;
     switch (path) {
       case "/sse": {
         const websockets = this.ctx.getWebSockets();
@@ -227,24 +267,24 @@ var _McpAgent = class _McpAgent extends DurableObject {
           return new Response("Websocket already connected", { status: 400 });
         }
         await this.ctx.storage.put("transportType", "sse");
-        __privateSet(this, _transportType, "sse");
-        if (!__privateGet(this, _transport)) {
-          __privateSet(this, _transport, new McpSSETransport(() => this.getWebSocket()));
-          await this.server.connect(__privateGet(this, _transport));
+        this._transportType = "sse";
+        if (!this._transport) {
+          this._transport = new McpSSETransport(() => this.getWebSocket());
+          await server.connect(this._transport);
         }
-        return __privateGet(this, _agent).fetch(request);
+        return this._agent.fetch(request);
       }
       case "/streamable-http": {
-        if (!__privateGet(this, _transport)) {
-          __privateSet(this, _transport, new McpStreamableHttpTransport(
+        if (!this._transport) {
+          this._transport = new McpStreamableHttpTransport(
             (id) => this.getWebSocketForResponseID(id),
-            (id) => __privateGet(this, _requestIdToConnectionId).delete(id)
-          ));
-          await this.server.connect(__privateGet(this, _transport));
+            (id) => this._requestIdToConnectionId.delete(id)
+          );
+          await server.connect(this._transport);
         }
         await this.ctx.storage.put("transportType", "streamable-http");
-        __privateSet(this, _transportType, "streamable-http");
-        return __privateGet(this, _agent).fetch(request);
+        this._transportType = "streamable-http";
+        return this._agent.fetch(request);
       }
       default:
         return new Response(
@@ -263,19 +303,19 @@ var _McpAgent = class _McpAgent extends DurableObject {
     return websockets[0];
   }
   getWebSocketForResponseID(id) {
-    const connectionId = __privateGet(this, _requestIdToConnectionId).get(id);
+    const connectionId = this._requestIdToConnectionId.get(id);
     if (connectionId === void 0) {
       return null;
     }
-    return __privateGet(this, _agent).getConnection(connectionId) ?? null;
+    return this._agent.getConnection(connectionId) ?? null;
   }
   // All messages received here. This is currently never called
   async onMessage(connection, event) {
-    if (__privateGet(this, _transportType) !== "streamable-http") {
+    if (this._transportType !== "streamable-http") {
       const err = new Error(
         "Internal Server Error: Expected streamable-http protocol"
       );
-      __privateGet(this, _transport)?.onerror?.(err);
+      this._transport?.onerror?.(err);
       return;
     }
     let message;
@@ -283,21 +323,85 @@ var _McpAgent = class _McpAgent extends DurableObject {
       const data = typeof event === "string" ? event : new TextDecoder().decode(event);
       message = JSONRPCMessageSchema.parse(JSON.parse(data));
     } catch (error) {
-      __privateGet(this, _transport)?.onerror?.(error);
+      this._transport?.onerror?.(error);
+      return;
+    }
+    if (await this._handleElicitationResponse(message)) {
       return;
     }
     if (isJSONRPCRequest(message)) {
-      __privateGet(this, _requestIdToConnectionId).set(message.id.toString(), connection.id);
+      this._requestIdToConnectionId.set(message.id.toString(), connection.id);
     }
-    __privateGet(this, _transport)?.onmessage?.(message);
+    this._transport?.onmessage?.(message);
+  }
+  /**
+   * Wait for elicitation response through storage polling
+   */
+  async _waitForElicitationResponse(requestId) {
+    const startTime = Date.now();
+    const timeout = 6e4;
+    try {
+      while (Date.now() - startTime < timeout) {
+        const response = await this.ctx.storage.get(
+          `elicitation:response:${requestId}`
+        );
+        if (response) {
+          await this.ctx.storage.delete(`elicitation:${requestId}`);
+          await this.ctx.storage.delete(`elicitation:response:${requestId}`);
+          return response;
+        }
+        await new Promise((resolve) => setTimeout(resolve, 100));
+      }
+      throw new Error("Elicitation request timed out");
+    } finally {
+      await this.ctx.storage.delete(`elicitation:${requestId}`);
+      await this.ctx.storage.delete(`elicitation:response:${requestId}`);
+    }
+  }
+  /**
+   * Handle elicitation responses   */
+  async _handleElicitationResponse(message) {
+    if (isJSONRPCResponse(message) && message.result) {
+      const requestId = message.id?.toString();
+      if (!requestId || !requestId.startsWith("elicit_")) return false;
+      const pendingRequest = await this.ctx.storage.get(
+        `elicitation:${requestId}`
+      );
+      if (!pendingRequest) return false;
+      await this.ctx.storage.put(
+        `elicitation:response:${requestId}`,
+        message.result
+      );
+      return true;
+    }
+    if (isJSONRPCError(message)) {
+      const requestId = message.id?.toString();
+      if (!requestId || !requestId.startsWith("elicit_")) return false;
+      const pendingRequest = await this.ctx.storage.get(
+        `elicitation:${requestId}`
+      );
+      if (!pendingRequest) return false;
+      const errorResult = {
+        action: "cancel",
+        content: {
+          error: message.error.message || "Elicitation request failed"
+        }
+      };
+      await this.ctx.storage.put(
+        `elicitation:response:${requestId}`,
+        errorResult
+      );
+      return true;
+    }
+    return false;
   }
   // All messages received over SSE after the initial connection has been established
   // will be passed here
-  async onSSEMcpMessage(sessionId, request) {
-    if (__privateGet(this, _status) !== "started") {
-      await __privateMethod(this, _McpAgent_instances, initialize_fn).call(this);
+  async onSSEMcpMessage(_sessionId, request) {
+    if (this._status !== "started") {
+      await this._initialize();
     }
-    if (__privateGet(this, _transportType) !== "sse") {
+    if (this._transportType !== "sse") {
       return new Error("Internal Server Error: Expected SSE protocol");
     }
     try {
@@ -306,35 +410,39 @@ var _McpAgent = class _McpAgent extends DurableObject {
       try {
         parsedMessage = JSONRPCMessageSchema.parse(message);
       } catch (error) {
-        __privateGet(this, _transport)?.onerror?.(error);
+        this._transport?.onerror?.(error);
         throw error;
       }
-      __privateGet(this, _transport)?.onmessage?.(parsedMessage);
+      if (await this._handleElicitationResponse(parsedMessage)) {
+        return null;
+      }
+      this._transport?.onmessage?.(parsedMessage);
       return null;
     } catch (error) {
-      __privateGet(this, _transport)?.onerror?.(error);
+      console.error("Error forwarding message to SSE:", error);
+      this._transport?.onerror?.(error);
       return error;
     }
   }
   // Delegate all websocket events to the underlying agent
   async webSocketMessage(ws, event) {
-    if (__privateGet(this, _status) !== "started") {
-      await __privateMethod(this, _McpAgent_instances, initialize_fn).call(this);
+    if (this._status !== "started") {
+      await this._initialize();
     }
-    return await __privateGet(this, _agent).webSocketMessage(ws, event);
+    return await this._agent.webSocketMessage(ws, event);
   }
   // WebSocket event handlers for hibernation support
   async webSocketError(ws, error) {
-    if (__privateGet(this, _status) !== "started") {
-      await __privateMethod(this, _McpAgent_instances, initialize_fn).call(this);
+    if (this._status !== "started") {
+      await this._initialize();
     }
-    return await __privateGet(this, _agent).webSocketError(ws, error);
+    return await this._agent.webSocketError(ws, error);
   }
   async webSocketClose(ws, code, reason, wasClean) {
-    if (__privateGet(this, _status) !== "started") {
-      await __privateMethod(this, _McpAgent_instances, initialize_fn).call(this);
+    if (this._status !== "started") {
+      await this._initialize();
     }
-    return await __privateGet(this, _agent).webSocketClose(ws, code, reason, wasClean);
+    return await this._agent.webSocketClose(ws, code, reason, wasClean);
   }
   static mount(path, {
     binding = "MCP_OBJECT",
@@ -353,18 +461,32 @@ var _McpAgent = class _McpAgent extends DurableObject {
     const basePattern = new URLPattern({ pathname });
     const messagePattern = new URLPattern({ pathname: `${pathname}/message` });
     return {
-      fetch: async (request, env, ctx) => {
+      async fetch(request, env, ctx) {
         const corsResponse = handleCORS(request, corsOptions);
         if (corsResponse) return corsResponse;
         const url = new URL(request.url);
-        const namespace = env[binding];
+        const bindingValue = env[binding];
+        if (bindingValue == null || typeof bindingValue !== "object") {
+          console.error(
+            `Could not find McpAgent binding for ${binding}. Did you update your wrangler configuration?`
+          );
+          return new Response("Invalid binding", { status: 500 });
+        }
+        if (!isDurableObjectNamespace(bindingValue)) {
+          return new Response("Invalid binding", { status: 500 });
+        }
+        const namespace = bindingValue;
         if (request.method === "GET" && basePattern.test(url)) {
           const sessionId = url.searchParams.get("sessionId") || namespace.newUniqueId().toString();
           const { readable, writable } = new TransformStream();
           const writer = writable.getWriter();
           const encoder = new TextEncoder();
+          const endpointUrl = new URL(request.url);
+          endpointUrl.pathname = encodeURI(`${pathname}/message`);
+          endpointUrl.searchParams.set("sessionId", sessionId);
+          const relativeUrlWithSession = endpointUrl.pathname + endpointUrl.search + endpointUrl.hash;
           const endpointMessage = `event: endpoint
-data: ${encodeURI(`${pathname}/message`)}?sessionId=${sessionId}
+data: ${relativeUrlWithSession}
 
 `;
           writer.write(encoder.encode(endpointMessage));
@@ -373,9 +495,14 @@ data: ${encodeURI(`${pathname}/message`)}?sessionId=${sessionId}
           await doStub._init(ctx.props);
           const upgradeUrl = new URL(request.url);
           upgradeUrl.pathname = "/sse";
+          const existingHeaders = {};
+          request.headers.forEach((value, key) => {
+            existingHeaders[key] = value;
+          });
           const response = await doStub.fetch(
             new Request(upgradeUrl, {
               headers: {
+                ...existingHeaders,
                 Upgrade: "websocket",
                 // Required by PartyServer
                 "x-partykit-room": sessionId
@@ -392,40 +519,49 @@ data: ${encodeURI(`${pathname}/message`)}?sessionId=${sessionId}
           }
           ws.accept();
           ws.addEventListener("message", (event) => {
-            try {
-              const message = JSON.parse(event.data);
-              const result = JSONRPCMessageSchema.safeParse(message);
-              if (!result.success) {
-                return;
-              }
-              const messageText = `event: message
+            async function onMessage(event2) {
+              try {
+                const message = JSON.parse(event2.data);
+                const result = JSONRPCMessageSchema.safeParse(message);
+                if (!result.success) {
+                  return;
+                }
+                const messageText = `event: message
 data: ${JSON.stringify(result.data)}
 
 `;
-              Promise.resolve(writer.write(encoder.encode(messageText)));
-            } catch (error) {
-              console.error("Error forwarding message to SSE:", error);
+                await writer.write(encoder.encode(messageText));
+              } catch (error) {
+                console.error("Error forwarding message to SSE:", error);
+              }
             }
+            onMessage(event).catch(console.error);
           });
           ws.addEventListener("error", (error) => {
-            try {
-              Promise.resolve(writer.close());
-            } catch (e) {
+            async function onError(_error) {
+              try {
+                await writer.close();
+              } catch (_e) {
+              }
             }
+            onError(error).catch(console.error);
           });
           ws.addEventListener("close", () => {
-            try {
-              Promise.resolve(writer.close());
-            } catch (error) {
-              console.error("Error closing SSE connection:", error);
+            async function onClose() {
+              try {
+                await writer.close();
+              } catch (error) {
+                console.error("Error closing SSE connection:", error);
+              }
             }
+            onClose().catch(console.error);
           });
           return new Response(readable, {
             headers: {
-              "Content-Type": "text/event-stream",
               "Cache-Control": "no-cache",
               Connection: "keep-alive",
-              "Access-Control-Allow-Origin": corsOptions?.origin || "*"
+              "Content-Type": "text/event-stream",
+              ...corsHeaders(request, corsOptions)
             }
           });
         }
@@ -460,23 +596,23 @@ data: ${JSON.stringify(result.data)}
           const error = await doStub.onSSEMcpMessage(sessionId, request);
           if (error) {
             return new Response(error.message, {
-              status: 400,
               headers: {
-                "Content-Type": "text/event-stream",
                 "Cache-Control": "no-cache",
                 Connection: "keep-alive",
-                "Access-Control-Allow-Origin": corsOptions?.origin || "*"
-              }
+                "Content-Type": "text/event-stream",
+                ...corsHeaders(request, corsOptions)
+              },
+              status: 400
             });
           }
           return new Response("Accepted", {
-            status: 202,
             headers: {
-              "Content-Type": "text/event-stream",
               "Cache-Control": "no-cache",
               Connection: "keep-alive",
-              "Access-Control-Allow-Origin": corsOptions?.origin || "*"
-            }
+              "Content-Type": "text/event-stream",
+              ...corsHeaders(request, corsOptions)
+            },
+            status: 202
           });
         }
         return new Response("Not Found", { status: 404 });
@@ -493,35 +629,45 @@ data: ${JSON.stringify(result.data)}
     }
     const basePattern = new URLPattern({ pathname });
     return {
-      fetch: async (request, env, ctx) => {
+      async fetch(request, env, ctx) {
         const corsResponse = handleCORS(request, corsOptions);
         if (corsResponse) {
           return corsResponse;
         }
         const url = new URL(request.url);
-        const namespace = env[binding];
+        const bindingValue = env[binding];
+        if (bindingValue == null || typeof bindingValue !== "object") {
+          console.error(
+            `Could not find McpAgent binding for ${binding}. Did you update your wrangler configuration?`
+          );
+          return new Response("Invalid binding", { status: 500 });
+        }
+        if (!isDurableObjectNamespace(bindingValue)) {
+          return new Response("Invalid binding", { status: 500 });
+        }
+        const namespace = bindingValue;
         if (request.method === "POST" && basePattern.test(url)) {
           const acceptHeader = request.headers.get("accept");
           if (!acceptHeader?.includes("application/json") || !acceptHeader.includes("text/event-stream")) {
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32e3,
                 message: "Not Acceptable: Client must accept both application/json and text/event-stream"
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 406 });
           }
           const ct = request.headers.get("content-type");
           if (!ct || !ct.includes("application/json")) {
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32e3,
                 message: "Unsupported Media Type: Content-Type must be application/json"
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 415 });
           }
@@ -531,12 +677,12 @@ data: ${JSON.stringify(result.data)}
           );
           if (contentLength > MAXIMUM_MESSAGE_SIZE_BYTES) {
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32e3,
                 message: `Request body too large. Maximum size is ${MAXIMUM_MESSAGE_SIZE_BYTES} bytes`
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 413 });
           }
@@ -544,14 +690,14 @@ data: ${JSON.stringify(result.data)}
           let rawMessage;
           try {
             rawMessage = await request.json();
-          } catch (error) {
+          } catch (_error) {
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32700,
                 message: "Parse error: Invalid JSON"
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 400 });
           }
@@ -565,12 +711,12 @@ data: ${JSON.stringify(result.data)}
           for (const msg of arrayMessage) {
             if (!JSONRPCMessageSchema.safeParse(msg).success) {
               const body2 = JSON.stringify({
-                jsonrpc: "2.0",
                 error: {
                   code: -32700,
                   message: "Parse error: Invalid JSON-RPC message"
                 },
-                id: null
+                id: null,
+                jsonrpc: "2.0"
               });
               return new Response(body2, { status: 400 });
             }
@@ -581,34 +727,34 @@ data: ${JSON.stringify(result.data)}
           );
           if (isInitializationRequest && sessionId) {
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32600,
                 message: "Invalid Request: Initialization requests must not include a sessionId"
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 400 });
           }
           if (isInitializationRequest && messages.length > 1) {
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32600,
                 message: "Invalid Request: Only one initialization request is allowed"
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 400 });
           }
           if (!isInitializationRequest && !sessionId) {
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32e3,
                 message: "Bad Request: Mcp-Session-Id header is required"
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 400 });
           }
@@ -618,14 +764,15 @@ data: ${JSON.stringify(result.data)}
           const isInitialized = await doStub.isInitialized();
           if (isInitializationRequest) {
             await doStub._init(ctx.props);
+            await doStub.setInitialized();
           } else if (!isInitialized) {
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32001,
                 message: "Session not found"
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 404 });
           }
@@ -634,9 +781,14 @@ data: ${JSON.stringify(result.data)}
           const encoder = new TextEncoder();
           const upgradeUrl = new URL(request.url);
           upgradeUrl.pathname = "/streamable-http";
+          const existingHeaders = {};
+          request.headers.forEach((value, key) => {
+            existingHeaders[key] = value;
+          });
           const response = await doStub.fetch(
             new Request(upgradeUrl, {
               headers: {
+                ...existingHeaders,
                 Upgrade: "websocket",
                 // Required by PartyServer
                 "x-partykit-room": sessionId
@@ -648,52 +800,61 @@ data: ${JSON.stringify(result.data)}
             console.error("Failed to establish WebSocket connection");
             await writer.close();
             const body2 = JSON.stringify({
-              jsonrpc: "2.0",
               error: {
                 code: -32001,
                 message: "Failed to establish WebSocket connection"
               },
-              id: null
+              id: null,
+              jsonrpc: "2.0"
             });
             return new Response(body2, { status: 500 });
           }
           const requestIds = /* @__PURE__ */ new Set();
           ws.accept();
           ws.addEventListener("message", (event) => {
-            try {
-              const data = typeof event.data === "string" ? event.data : new TextDecoder().decode(event.data);
-              const message = JSON.parse(data);
-              const result = JSONRPCMessageSchema.safeParse(message);
-              if (!result.success) {
-                return;
-              }
-              if (isJSONRPCResponse(result.data) || isJSONRPCError(result.data)) {
-                requestIds.delete(result.data.id);
-              }
-              const messageText = `event: message
+            async function onMessage(event2) {
+              try {
+                const data = typeof event2.data === "string" ? event2.data : new TextDecoder().decode(event2.data);
+                const message = JSON.parse(data);
+                const result = JSONRPCMessageSchema.safeParse(message);
+                if (!result.success) {
+                  return;
+                }
+                if (isJSONRPCResponse(result.data) || isJSONRPCError(result.data)) {
+                  requestIds.delete(result.data.id);
+                }
+                const messageText = `event: message
 data: ${JSON.stringify(result.data)}
 
 `;
-              Promise.resolve(writer.write(encoder.encode(messageText)));
-              if (requestIds.size === 0) {
-                ws.close();
+                await writer.write(encoder.encode(messageText));
+                if (requestIds.size === 0) {
+                  ws.close();
+                }
+              } catch (error) {
+                console.error("Error forwarding message to SSE:", error);
               }
-            } catch (error) {
-              console.error("Error forwarding message to SSE:", error);
             }
+            onMessage(event).catch(console.error);
           });
           ws.addEventListener("error", (error) => {
-            try {
-              Promise.resolve(writer.close());
-            } catch (e) {
+            async function onError(_error) {
+              try {
+                await writer.close();
+              } catch (_e) {
+              }
             }
+            onError(error).catch(console.error);
           });
           ws.addEventListener("close", () => {
-            try {
-              Promise.resolve(writer.close());
-            } catch (error) {
-              console.error("Error closing SSE connection:", error);
+            async function onClose() {
+              try {
+                await writer.close();
+              } catch (error) {
+                console.error("Error closing SSE connection:", error);
+              }
             }
+            onClose().catch(console.error);
           });
           const hasOnlyNotificationsOrResponses = messages.every(
             (msg) => isJSONRPCNotification(msg) || isJSONRPCResponse(msg)
@@ -703,7 +864,10 @@ data: ${JSON.stringify(result.data)}
               ws.send(JSON.stringify(message));
             }
             ws.close();
-            return new Response(null, { status: 202 });
+            return new Response(null, {
+              headers: corsHeaders(request, corsOptions),
+              status: 202
+            });
           }
           for (const message of messages) {
             if (isJSONRPCRequest(message)) {
@@ -713,43 +877,32 @@ data: ${JSON.stringify(result.data)}
           }
           return new Response(readable, {
             headers: {
-              "Content-Type": "text/event-stream",
               "Cache-Control": "no-cache",
               Connection: "keep-alive",
+              "Content-Type": "text/event-stream",
               "mcp-session-id": sessionId,
-              "Access-Control-Allow-Origin": corsOptions?.origin || "*"
+              ...corsHeaders(request, corsOptions)
             },
             status: 200
           });
         }
         const body = JSON.stringify({
-          jsonrpc: "2.0",
           error: {
             code: -32e3,
             message: "Method not allowed"
           },
-          id: null
+          id: null,
+          jsonrpc: "2.0"
         });
         return new Response(body, { status: 405 });
       }
     };
   }
 };
-_status = new WeakMap();
-_transport = new WeakMap();
-_transportType = new WeakMap();
-_requestIdToConnectionId = new WeakMap();
-_agent = new WeakMap();
-_McpAgent_instances = new WeakSet();
-initialize_fn = async function() {
-  await this.ctx.blockConcurrencyWhile(async () => {
-    __privateSet(this, _status, "starting");
-    await this.onStart();
-    __privateSet(this, _status, "started");
-  });
-};
-var McpAgent = _McpAgent;
 export {
-  McpAgent
+  ElicitRequestSchema,
+  McpAgent,
+  SSEEdgeClientTransport,
+  StreamableHTTPEdgeClientTransport
 };
 //# sourceMappingURL=index.js.map