agents 0.0.0-ed3f94d → 0.0.0-edd3357

package/dist/mcp/index.js

@@ -1,21 +1,21 @@
 import {
   Agent,
-  getAgentByName,
-  getCurrentAgent
-} from "../chunk-254F4GDT.js";
+  getAgentByName
+} from "../chunk-YDUDMOL6.js";
 import {
   SSEEdgeClientTransport,
   StreamableHTTPEdgeClientTransport
-} from "../chunk-3OT2NNEW.js";
-import "../chunk-Z44WASMA.js";
+} from "../chunk-MH46VMM4.js";
+import "../chunk-LL2AFX7V.js";
 import "../chunk-QEVM4BVL.js";
-import "../chunk-BER7KXUJ.js";
-import "../chunk-PR4QN5HX.js";
+import "../chunk-AVYJQSLW.js";
 
 // src/mcp/index.ts
 import {
-  JSONRPCMessageSchema as JSONRPCMessageSchema3,
-  isJSONRPCError as isJSONRPCError2,
+  JSONRPCMessageSchema as JSONRPCMessageSchema2,
+  isJSONRPCError as isJSONRPCError3,
+  isJSONRPCNotification as isJSONRPCNotification2,
+  isJSONRPCRequest as isJSONRPCRequest2,
   isJSONRPCResponse as isJSONRPCResponse3
 } from "@modelcontextprotocol/sdk/types.js";
 
@@ -24,10 +24,12 @@ import {
   JSONRPCMessageSchema,
   InitializeRequestSchema,
   isJSONRPCResponse,
-  isJSONRPCNotification
+  isJSONRPCError,
+  isJSONRPCNotification,
+  isJSONRPCRequest
 } from "@modelcontextprotocol/sdk/types.js";
-var MCP_HTTP_METHOD_HEADER = "cf-mcp-method";
-var MCP_MESSAGE_HEADER = "cf-mcp-message";
+var STANDALONE_SSE_MARKER = "standalone-sse";
+var STANDALONE_SSE_METHOD = "cf/standalone_sse/attach";
 var MAXIMUM_MESSAGE_SIZE_BYTES = 4 * 1024 * 1024;
 var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
   let pathname = basePath;
@@ -178,8 +180,6 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
         const req = new Request(request.url, {
           headers: {
             ...existingHeaders,
-            [MCP_HTTP_METHOD_HEADER]: "POST",
-            [MCP_MESSAGE_HEADER]: JSON.stringify(messages),
             Upgrade: "websocket"
           }
         });
@@ -199,17 +199,26 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
           });
           return new Response(body2, { status: 500 });
         }
+        const requestIds = /* @__PURE__ */ new Set();
         ws.accept();
         ws.addEventListener("message", (event) => {
           async function onMessage(event2) {
             try {
               const data = typeof event2.data === "string" ? event2.data : new TextDecoder().decode(event2.data);
               const message = JSON.parse(data);
-              if (message.type !== "cf_mcp_agent_event" /* CF_MCP_AGENT_EVENT */) {
+              const result = JSONRPCMessageSchema.safeParse(message);
+              if (!result.success) {
                 return;
               }
-              await writer.write(encoder.encode(message.event));
-              if (message.close) {
+              if (isJSONRPCResponse(result.data) || isJSONRPCError(result.data)) {
+                requestIds.delete(result.data.id);
+              }
+              const messageText = `event: message
+data: ${JSON.stringify(result.data)}
+
+`;
+              await writer.write(encoder.encode(messageText));
+              if (requestIds.size === 0) {
                 ws?.close();
                 await writer.close().catch(() => {
                 });
@@ -238,12 +247,21 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
           (msg) => isJSONRPCNotification(msg) || isJSONRPCResponse(msg)
         );
         if (hasOnlyNotificationsOrResponses) {
+          for (const message of messages) {
+            ws.send(JSON.stringify(message));
+          }
           ws.close();
           return new Response(null, {
             headers: corsHeaders(request, corsOptions),
             status: 202
           });
         }
+        for (const message of messages) {
+          if (isJSONRPCRequest(message)) {
+            requestIds.add(message.id);
+          }
+          ws.send(JSON.stringify(message));
+        }
         return new Response(readable, {
           headers: {
             "Cache-Control": "no-cache",
@@ -269,17 +287,7 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
         }
         const sessionId = request.headers.get("mcp-session-id");
         if (!sessionId)
-          return new Response(
-            JSON.stringify({
-              error: {
-                code: -32e3,
-                message: "Bad Request: Mcp-Session-Id header is required"
-              },
-              id: null,
-              jsonrpc: "2.0"
-            }),
-            { status: 400 }
-          );
+          return new Response("Missing sessionId", { status: 400 });
         const { readable, writable } = new TransformStream();
         const writer = writable.getWriter();
         const encoder = new TextEncoder();
@@ -308,7 +316,6 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
           new Request(request.url, {
             headers: {
               ...existingHeaders,
-              [MCP_HTTP_METHOD_HEADER]: "GET",
               Upgrade: "websocket"
             }
           })
@@ -321,15 +328,24 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
           });
         }
         ws.accept();
+        ws.send(
+          JSON.stringify({
+            jsonrpc: "2.0",
+            method: STANDALONE_SSE_METHOD,
+            params: {}
+          })
+        );
         ws.addEventListener("message", (event) => {
           try {
             async function onMessage(ev) {
               const data = typeof ev.data === "string" ? ev.data : new TextDecoder().decode(ev.data);
-              const message = JSON.parse(data);
-              if (message.type !== "cf_mcp_agent_event" /* CF_MCP_AGENT_EVENT */) {
-                return;
-              }
-              await writer.write(encoder.encode(message.event));
+              const parsed = JSONRPCMessageSchema.safeParse(JSON.parse(data));
+              if (!parsed.success) return;
+              const frame = `event: message
+data: ${JSON.stringify(parsed.data)}
+
+`;
+              await writer.write(encoder.encode(frame));
             }
             onMessage(event).catch(console.error);
           } catch (e) {
@@ -572,10 +588,8 @@ function isDurableObjectNamespace(namespace) {
 
 // src/mcp/transport.ts
 import {
-  isJSONRPCError,
-  isJSONRPCRequest,
-  isJSONRPCResponse as isJSONRPCResponse2,
-  JSONRPCMessageSchema as JSONRPCMessageSchema2
+  isJSONRPCError as isJSONRPCError2,
+  isJSONRPCResponse as isJSONRPCResponse2
 } from "@modelcontextprotocol/sdk/types.js";
 var McpSSETransport = class {
   constructor(getWebSocket) {
@@ -606,190 +620,59 @@ var McpSSETransport = class {
     this.onclose?.();
   }
 };
-var StreamableHTTPServerTransport = class {
-  constructor(options) {
+var McpStreamableHttpTransport = class {
+  constructor(getWebSocketForMessageID, notifyResponseIdSent, getWebSocketForStandaloneSse) {
     this._started = false;
-    // This is to keep track whether all messages from a single POST request have been answered.
-    // I's fine that we don't persist this since it's only for backwards compatibility as clients
-    // should no longer batch requests, per the spec.
-    this._requestResponseMap = /* @__PURE__ */ new Map();
-    const { agent } = getCurrentAgent();
-    if (!agent)
-      throw new Error("McpAgent was not found in Transport constructor");
-    this.sessionId = agent.getSessionId();
-    this._eventStore = options.eventStore;
+    this._getWebSocketForMessageID = getWebSocketForMessageID;
+    this._notifyResponseIdSent = notifyResponseIdSent;
+    this._getWebSocketForStandaloneSse = getWebSocketForStandaloneSse;
   }
-  /**
-   * Starts the transport. This is required by the Transport interface but is a no-op
-   * for the Streamable HTTP transport as connections are managed per-request.
-   */
   async start() {
     if (this._started) {
       throw new Error("Transport already started");
     }
     this._started = true;
   }
-  /**
-   * Handles GET requests for SSE stream
-   */
-  async handleGetRequest(req) {
-    const { connection } = getCurrentAgent();
-    if (!connection)
-      throw new Error("Connection was not found in handleGetRequest");
-    if (this._eventStore) {
-      const lastEventId = req.headers.get("last-event-id");
-      if (lastEventId) {
-        await this.replayEvents(lastEventId);
-        return;
-      }
-    }
-    connection.setState({
-      _standaloneSse: true
-    });
-  }
-  /**
-   * Replays events that would have been sent after the specified event ID
-   * Only used when resumability is enabled
-   */
-  async replayEvents(lastEventId) {
-    if (!this._eventStore) {
-      return;
-    }
-    const { connection } = getCurrentAgent();
-    if (!connection)
-      throw new Error("Connection was not available in replayEvents");
-    try {
-      await this._eventStore?.replayEventsAfter(lastEventId, {
-        send: async (eventId, message) => {
-          try {
-            this.writeSSEEvent(connection, message, eventId);
-          } catch (error) {
-            this.onerror?.(error);
-          }
-        }
-      });
-    } catch (error) {
-      this.onerror?.(error);
-    }
-  }
-  /**
-   * Writes an event to the SSE stream with proper formatting
-   */
-  writeSSEEvent(connection, message, eventId, close) {
-    let eventData = "event: message\n";
-    if (eventId) {
-      eventData += `id: ${eventId}
-`;
-    }
-    eventData += `data: ${JSON.stringify(message)}
-
-`;
-    return connection.send(
-      JSON.stringify({
-        type: "cf_mcp_agent_event" /* CF_MCP_AGENT_EVENT */,
-        event: eventData,
-        close
-      })
-    );
-  }
-  /**
-   * Handles POST requests containing JSON-RPC messages
-   */
-  async handlePostRequest(req, parsedBody) {
-    const authInfo = req.auth;
-    const requestInfo = {
-      headers: Object.fromEntries(req.headers.entries())
-    };
-    delete requestInfo.headers[MCP_HTTP_METHOD_HEADER];
-    delete requestInfo.headers[MCP_MESSAGE_HEADER];
-    delete requestInfo.headers.upgrade;
-    const rawMessage = parsedBody;
-    let messages;
-    if (Array.isArray(rawMessage)) {
-      messages = rawMessage.map((msg) => JSONRPCMessageSchema2.parse(msg));
-    } else {
-      messages = [JSONRPCMessageSchema2.parse(rawMessage)];
-    }
-    const hasRequests = messages.some(isJSONRPCRequest);
-    if (!hasRequests) {
-      for (const message of messages) {
-        this.onmessage?.(message, { authInfo, requestInfo });
-      }
-    } else if (hasRequests) {
-      const { connection } = getCurrentAgent();
-      if (!connection)
-        throw new Error("Connection was not found in handlePostRequest");
-      const requestIds = messages.filter(isJSONRPCRequest).map((message) => message.id);
-      connection.setState({
-        requestIds
-      });
-      for (const message of messages) {
-        this.onmessage?.(message, { authInfo, requestInfo });
-      }
-    }
-  }
-  async close() {
-    const { agent } = getCurrentAgent();
-    if (!agent) throw new Error("Agent was not found in close");
-    for (const conn of agent.getConnections()) {
-      conn.close(1e3, "Session closed");
-    }
-    this.onclose?.();
-  }
   async send(message, options) {
-    const { agent } = getCurrentAgent();
-    if (!agent) throw new Error("Agent was not found in send");
+    if (!this._started) {
+      throw new Error("Transport not started");
+    }
     let requestId = options?.relatedRequestId;
-    if (isJSONRPCResponse2(message) || isJSONRPCError(message)) {
+    if (isJSONRPCResponse2(message) || isJSONRPCError2(message)) {
       requestId = message.id;
     }
     if (requestId === void 0) {
-      if (isJSONRPCResponse2(message) || isJSONRPCError(message)) {
+      if (isJSONRPCResponse2(message) || isJSONRPCError2(message)) {
         throw new Error(
           "Cannot send a response on a standalone SSE stream unless resuming a previous client request"
         );
       }
-      let standaloneConnection;
-      for (const conn of agent.getConnections()) {
-        if (conn.state?._standaloneSse) standaloneConnection = conn;
-      }
-      if (standaloneConnection === void 0) {
+      const standaloneSseSocket = this._getWebSocketForStandaloneSse();
+      if (!standaloneSseSocket) {
         return;
       }
-      let eventId2;
-      if (this._eventStore) {
-        eventId2 = await this._eventStore.storeEvent(
-          standaloneConnection.id,
-          message
-        );
+      try {
+        standaloneSseSocket?.send(JSON.stringify(message));
+      } catch (error) {
+        this.onerror?.(error);
       }
-      this.writeSSEEvent(standaloneConnection, message, eventId2);
       return;
     }
-    const connection = Array.from(
-      agent.getConnections()
-    ).find((conn) => conn.state?.requestIds?.includes(requestId));
-    if (!connection) {
-      throw new Error(
-        `No connection established for request ID: ${String(requestId)}`
-      );
-    }
-    let eventId;
-    if (this._eventStore) {
-      eventId = await this._eventStore.storeEvent(connection.id, message);
+    const websocket = this._getWebSocketForMessageID(requestId.toString());
+    if (!websocket) {
+      throw new Error(`Could not find WebSocket for message id: ${requestId}`);
     }
-    let shouldClose = false;
-    if (isJSONRPCResponse2(message) || isJSONRPCError(message)) {
-      this._requestResponseMap.set(requestId, message);
-      const relatedIds = connection.state?.requestIds ?? [];
-      shouldClose = relatedIds.every((id) => this._requestResponseMap.has(id));
-      if (shouldClose) {
-        for (const id of relatedIds) {
-          this._requestResponseMap.delete(id);
-        }
+    try {
+      websocket?.send(JSON.stringify(message));
+      if (isJSONRPCResponse2(message) || isJSONRPCError2(message)) {
+        this._notifyResponseIdSent(message.id.toString());
       }
+    } catch (error) {
+      this.onerror?.(error);
     }
-    this.writeSSEEvent(connection, message, eventId, shouldClose);
+  }
+  async close() {
+    this.onclose?.();
   }
 };
 
@@ -798,6 +681,17 @@ import {
   ElicitRequestSchema
 } from "@modelcontextprotocol/sdk/types.js";
 var McpAgent = class _McpAgent extends Agent {
+  constructor(ctx, env) {
+    super(ctx, env);
+    this._requestIdToConnectionId = /* @__PURE__ */ new Map();
+    for (const connection of this.getConnections()) {
+      const meta = connection.state;
+      if (meta?.role === STANDALONE_SSE_MARKER) {
+        this._standaloneSseConnectionId = connection.id;
+        return;
+      }
+    }
+  }
   /*
    * Helpers
    */
@@ -824,18 +718,10 @@ var McpAgent = class _McpAgent extends Agent {
         );
     }
   }
-  /** Read the sessionId for this agent.
-   * This relies on the naming scheme being `sse:${sessionId}`
-   * or `streamable-http:${sessionId}`.
-   */
-  getSessionId() {
-    const [_, sessionId] = this.name.split(":");
-    if (!sessionId) {
-      throw new Error(
-        "Invalid session id. McpAgent must be addressed with a valid session id."
-      );
-    }
-    return sessionId;
+  /** Get the WebSocket for the standalone SSE if any. Streamable HTTP only. */
+  getWebSocketForStandaloneSse() {
+    if (!this._standaloneSseConnectionId) return null;
+    return this.getConnection(this._standaloneSseConnectionId) ?? null;
   }
   /** Get the unique WebSocket. SSE transport only. */
   getWebSocket() {
@@ -845,6 +731,14 @@ var McpAgent = class _McpAgent extends Agent {
     }
     return websockets[0];
   }
+  /** Get the corresponding WebSocket for a responseId. Streamable HTTP only. */
+  getWebSocketForResponseID(id) {
+    const connectionId = this._requestIdToConnectionId.get(id);
+    if (connectionId === void 0) {
+      return null;
+    }
+    return this.getConnection(connectionId) ?? null;
+  }
   /** Returns a new transport matching the type of the Agent. */
   initTransport() {
     switch (this.getTransportType()) {
@@ -852,7 +746,11 @@ var McpAgent = class _McpAgent extends Agent {
         return new McpSSETransport(() => this.getWebSocket());
       }
       case "streamable-http": {
-        return new StreamableHTTPServerTransport({});
+        return new McpStreamableHttpTransport(
+          (id) => this.getWebSocketForResponseID(id),
+          (id) => this._requestIdToConnectionId.delete(id),
+          () => this.getWebSocketForStandaloneSse()
+        );
       }
     }
   }
@@ -861,12 +759,6 @@ var McpAgent = class _McpAgent extends Agent {
     await this.ctx.storage.put("props", props ?? {});
     this.props = props;
   }
-  async reinitializeServer() {
-    const initializeRequest = await this.getInitializeRequest();
-    if (initializeRequest) {
-      this._transport?.onmessage?.(initializeRequest);
-    }
-  }
   /*
    * Base Agent / Parykit Server overrides
    */
@@ -878,10 +770,13 @@ var McpAgent = class _McpAgent extends Agent {
     const server = await this.server;
     this._transport = this.initTransport();
     await server.connect(this._transport);
-    await this.reinitializeServer();
+    const initializeRequest = await this.getInitializeRequest();
+    if (initializeRequest) {
+      this._transport.onmessage?.(initializeRequest);
+    }
   }
   /** Validates new WebSocket connections. */
-  async onConnect(conn, { request: req }) {
+  async onConnect(conn, _) {
     switch (this.getTransportType()) {
       case "sse": {
         const websockets = Array.from(this.getConnections());
@@ -892,19 +787,54 @@ var McpAgent = class _McpAgent extends Agent {
         break;
       }
       case "streamable-http":
-        if (this._transport instanceof StreamableHTTPServerTransport) {
-          switch (req.headers.get(MCP_HTTP_METHOD_HEADER)) {
-            case "POST": {
-              const payloadHeader = req.headers.get(MCP_MESSAGE_HEADER);
-              const parsedBody = await JSON.parse(payloadHeader ?? "{}");
-              this._transport?.handlePostRequest(req, parsedBody);
-              break;
-            }
-            case "GET":
-              this._transport?.handleGetRequest(req);
-              break;
-          }
-        }
+        break;
+    }
+  }
+  /** Handles MCP Messages for Streamable HTTP. */
+  async onMessage(connection, event) {
+    if (this.getTransportType() !== "streamable-http") {
+      const err = new Error(
+        "Internal Server Error: Expected streamable-http protocol"
+      );
+      this._transport?.onerror?.(err);
+      return;
+    }
+    let message;
+    try {
+      const data = typeof event === "string" ? event : new TextDecoder().decode(event);
+      message = JSONRPCMessageSchema2.parse(JSON.parse(data));
+    } catch (error) {
+      this._transport?.onerror?.(error);
+      return;
+    }
+    if (isJSONRPCNotification2(message) && message.method === STANDALONE_SSE_METHOD) {
+      if (this._standaloneSseConnectionId && this._standaloneSseConnectionId !== connection.id) {
+        const standaloneSseSocket = this.getConnection(
+          this._standaloneSseConnectionId
+        );
+        standaloneSseSocket?.close(1e3, "replaced");
+      }
+      connection.setState({
+        role: STANDALONE_SSE_MARKER
+      });
+      this._standaloneSseConnectionId = connection.id;
+      return;
+    }
+    if (await this._handleElicitationResponse(message)) {
+      return;
+    }
+    if (isJSONRPCRequest2(message)) {
+      this._requestIdToConnectionId.set(message.id.toString(), connection.id);
+    }
+    this._transport?.onmessage?.(message);
+  }
+  /** Remove clients from our cache when they disconnect */
+  async onClose(conn, _code, _reason, _wasClean) {
+    for (const [reqId, connId] of this._requestIdToConnectionId) {
+      if (connId === conn.id) this._requestIdToConnectionId.delete(reqId);
+    }
+    if (this._standaloneSseConnectionId === conn.id) {
+      this._standaloneSseConnectionId = void 0;
     }
   }
   /*
@@ -918,7 +848,7 @@ var McpAgent = class _McpAgent extends Agent {
     try {
       let parsedMessage;
       try {
-        parsedMessage = JSONRPCMessageSchema3.parse(messageBody);
+        parsedMessage = JSONRPCMessageSchema2.parse(messageBody);
       } catch (error) {
         this._transport?.onerror?.(error);
         throw error;
@@ -1007,7 +937,7 @@ var McpAgent = class _McpAgent extends Agent {
       );
       return true;
     }
-    if (isJSONRPCError2(message)) {
+    if (isJSONRPCError3(message)) {
       const requestId = message.id?.toString();
       if (!requestId || !requestId.startsWith("elicit_")) return false;
       const pendingRequest = await this.ctx.storage.get(