agents 0.1.4 → 0.1.6

package/dist/mcp/index.js

@@ -1,21 +1,21 @@
 import {
   Agent,
-  getAgentByName
-} from "../chunk-YDUDMOL6.js";
+  getAgentByName,
+  getCurrentAgent
+} from "../chunk-KK7D3KRW.js";
 import {
   SSEEdgeClientTransport,
   StreamableHTTPEdgeClientTransport
-} from "../chunk-MH46VMM4.js";
+} from "../chunk-THPMWGLS.js";
 import "../chunk-LL2AFX7V.js";
 import "../chunk-QEVM4BVL.js";
-import "../chunk-AVYJQSLW.js";
+import "../chunk-BER7KXUJ.js";
+import "../chunk-PR4QN5HX.js";
 
 // src/mcp/index.ts
 import {
-  JSONRPCMessageSchema as JSONRPCMessageSchema2,
-  isJSONRPCError as isJSONRPCError3,
-  isJSONRPCNotification as isJSONRPCNotification2,
-  isJSONRPCRequest as isJSONRPCRequest2,
+  JSONRPCMessageSchema as JSONRPCMessageSchema3,
+  isJSONRPCError as isJSONRPCError2,
   isJSONRPCResponse as isJSONRPCResponse3
 } from "@modelcontextprotocol/sdk/types.js";
 
@@ -24,12 +24,10 @@ import {
   JSONRPCMessageSchema,
   InitializeRequestSchema,
   isJSONRPCResponse,
-  isJSONRPCError,
-  isJSONRPCNotification,
-  isJSONRPCRequest
+  isJSONRPCNotification
 } from "@modelcontextprotocol/sdk/types.js";
-var STANDALONE_SSE_MARKER = "standalone-sse";
-var STANDALONE_SSE_METHOD = "cf/standalone_sse/attach";
+var MCP_HTTP_METHOD_HEADER = "cf-mcp-method";
+var MCP_MESSAGE_HEADER = "cf-mcp-message";
 var MAXIMUM_MESSAGE_SIZE_BYTES = 4 * 1024 * 1024;
 var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
   let pathname = basePath;
@@ -180,6 +178,8 @@ 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,26 +199,17 @@ 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);
-              const result = JSONRPCMessageSchema.safeParse(message);
-              if (!result.success) {
+              if (message.type !== "cf_mcp_agent_event" /* CF_MCP_AGENT_EVENT */) {
                 return;
               }
-              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) {
+              await writer.write(encoder.encode(message.event));
+              if (message.close) {
                 ws?.close();
                 await writer.close().catch(() => {
                 });
@@ -247,21 +238,12 @@ data: ${JSON.stringify(result.data)}
           (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",
@@ -287,7 +269,17 @@ data: ${JSON.stringify(result.data)}
         }
         const sessionId = request.headers.get("mcp-session-id");
         if (!sessionId)
-          return new Response("Missing sessionId", { status: 400 });
+          return new Response(
+            JSON.stringify({
+              error: {
+                code: -32e3,
+                message: "Bad Request: Mcp-Session-Id header is required"
+              },
+              id: null,
+              jsonrpc: "2.0"
+            }),
+            { status: 400 }
+          );
         const { readable, writable } = new TransformStream();
         const writer = writable.getWriter();
         const encoder = new TextEncoder();
@@ -316,6 +308,7 @@ data: ${JSON.stringify(result.data)}
           new Request(request.url, {
             headers: {
               ...existingHeaders,
+              [MCP_HTTP_METHOD_HEADER]: "GET",
               Upgrade: "websocket"
             }
           })
@@ -328,24 +321,15 @@ data: ${JSON.stringify(result.data)}
           });
         }
         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 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));
+              const message = JSON.parse(data);
+              if (message.type !== "cf_mcp_agent_event" /* CF_MCP_AGENT_EVENT */) {
+                return;
+              }
+              await writer.write(encoder.encode(message.event));
             }
             onMessage(event).catch(console.error);
           } catch (e) {
@@ -588,8 +572,10 @@ function isDurableObjectNamespace(namespace) {
 
 // src/mcp/transport.ts
 import {
-  isJSONRPCError as isJSONRPCError2,
-  isJSONRPCResponse as isJSONRPCResponse2
+  isJSONRPCError,
+  isJSONRPCRequest,
+  isJSONRPCResponse as isJSONRPCResponse2,
+  JSONRPCMessageSchema as JSONRPCMessageSchema2
 } from "@modelcontextprotocol/sdk/types.js";
 var McpSSETransport = class {
   constructor(getWebSocket) {
@@ -620,59 +606,190 @@ var McpSSETransport = class {
     this.onclose?.();
   }
 };
-var McpStreamableHttpTransport = class {
-  constructor(getWebSocketForMessageID, notifyResponseIdSent, getWebSocketForStandaloneSse) {
+var StreamableHTTPServerTransport = class {
+  constructor(options) {
     this._started = false;
-    this._getWebSocketForMessageID = getWebSocketForMessageID;
-    this._notifyResponseIdSent = notifyResponseIdSent;
-    this._getWebSocketForStandaloneSse = getWebSocketForStandaloneSse;
+    // 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;
   }
+  /**
+   * 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;
   }
-  async send(message, options) {
-    if (!this._started) {
-      throw new Error("Transport not started");
+  /**
+   * 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");
     let requestId = options?.relatedRequestId;
-    if (isJSONRPCResponse2(message) || isJSONRPCError2(message)) {
+    if (isJSONRPCResponse2(message) || isJSONRPCError(message)) {
       requestId = message.id;
     }
     if (requestId === void 0) {
-      if (isJSONRPCResponse2(message) || isJSONRPCError2(message)) {
+      if (isJSONRPCResponse2(message) || isJSONRPCError(message)) {
         throw new Error(
           "Cannot send a response on a standalone SSE stream unless resuming a previous client request"
         );
       }
-      const standaloneSseSocket = this._getWebSocketForStandaloneSse();
-      if (!standaloneSseSocket) {
+      let standaloneConnection;
+      for (const conn of agent.getConnections()) {
+        if (conn.state?._standaloneSse) standaloneConnection = conn;
+      }
+      if (standaloneConnection === void 0) {
         return;
       }
-      try {
-        standaloneSseSocket?.send(JSON.stringify(message));
-      } catch (error) {
-        this.onerror?.(error);
+      let eventId2;
+      if (this._eventStore) {
+        eventId2 = await this._eventStore.storeEvent(
+          standaloneConnection.id,
+          message
+        );
       }
+      this.writeSSEEvent(standaloneConnection, message, eventId2);
       return;
     }
-    const websocket = this._getWebSocketForMessageID(requestId.toString());
-    if (!websocket) {
-      throw new Error(`Could not find WebSocket for message id: ${requestId}`);
+    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)}`
+      );
     }
-    try {
-      websocket?.send(JSON.stringify(message));
-      if (isJSONRPCResponse2(message) || isJSONRPCError2(message)) {
-        this._notifyResponseIdSent(message.id.toString());
+    let eventId;
+    if (this._eventStore) {
+      eventId = await this._eventStore.storeEvent(connection.id, message);
+    }
+    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);
+        }
       }
-    } catch (error) {
-      this.onerror?.(error);
     }
-  }
-  async close() {
-    this.onclose?.();
+    this.writeSSEEvent(connection, message, eventId, shouldClose);
   }
 };
 
@@ -681,17 +798,6 @@ 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
    */
@@ -718,10 +824,18 @@ var McpAgent = class _McpAgent extends Agent {
         );
     }
   }
-  /** Get the WebSocket for the standalone SSE if any. Streamable HTTP only. */
-  getWebSocketForStandaloneSse() {
-    if (!this._standaloneSseConnectionId) return null;
-    return this.getConnection(this._standaloneSseConnectionId) ?? null;
+  /** 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 unique WebSocket. SSE transport only. */
   getWebSocket() {
@@ -731,14 +845,6 @@ 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()) {
@@ -746,11 +852,7 @@ var McpAgent = class _McpAgent extends Agent {
         return new McpSSETransport(() => this.getWebSocket());
       }
       case "streamable-http": {
-        return new McpStreamableHttpTransport(
-          (id) => this.getWebSocketForResponseID(id),
-          (id) => this._requestIdToConnectionId.delete(id),
-          () => this.getWebSocketForStandaloneSse()
-        );
+        return new StreamableHTTPServerTransport({});
       }
     }
   }
@@ -759,6 +861,12 @@ 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
    */
@@ -770,13 +878,10 @@ var McpAgent = class _McpAgent extends Agent {
     const server = await this.server;
     this._transport = this.initTransport();
     await server.connect(this._transport);
-    const initializeRequest = await this.getInitializeRequest();
-    if (initializeRequest) {
-      this._transport.onmessage?.(initializeRequest);
-    }
+    await this.reinitializeServer();
   }
   /** Validates new WebSocket connections. */
-  async onConnect(conn, _) {
+  async onConnect(conn, { request: req }) {
     switch (this.getTransportType()) {
       case "sse": {
         const websockets = Array.from(this.getConnections());
@@ -787,54 +892,19 @@ var McpAgent = class _McpAgent extends Agent {
         break;
       }
       case "streamable-http":
-        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;
+        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;
+          }
+        }
     }
   }
   /*
@@ -848,7 +918,7 @@ var McpAgent = class _McpAgent extends Agent {
     try {
       let parsedMessage;
       try {
-        parsedMessage = JSONRPCMessageSchema2.parse(messageBody);
+        parsedMessage = JSONRPCMessageSchema3.parse(messageBody);
       } catch (error) {
         this._transport?.onerror?.(error);
         throw error;
@@ -937,7 +1007,7 @@ var McpAgent = class _McpAgent extends Agent {
       );
       return true;
     }
-    if (isJSONRPCError3(message)) {
+    if (isJSONRPCError2(message)) {
       const requestId = message.id?.toString();
       if (!requestId || !requestId.startsWith("elicit_")) return false;
       const pendingRequest = await this.ctx.storage.get(