agents 0.1.3 → 0.1.5

package/dist/mcp/index.js

@@ -1,21 +1,21 @@
 import {
   Agent,
-  getAgentByName
-} from "../chunk-YDUDMOL6.js";
+  getAgentByName,
+  getCurrentAgent
+} from "../chunk-HS7VEROK.js";
 import {
   SSEEdgeClientTransport,
   StreamableHTTPEdgeClientTransport
-} from "../chunk-MH46VMM4.js";
+} from "../chunk-SKACXF37.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;
@@ -114,10 +112,10 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
           }
         }
         messages = arrayMessage.map((msg) => JSONRPCMessageSchema.parse(msg));
-        const isInitializationRequest = messages.some(
+        const maybeInitializeRequest = messages.find(
           (msg) => InitializeRequestSchema.safeParse(msg).success
         );
-        if (isInitializationRequest && sessionId) {
+        if (!!maybeInitializeRequest && sessionId) {
           const body2 = JSON.stringify({
             error: {
               code: -32600,
@@ -128,7 +126,7 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
           });
           return new Response(body2, { status: 400 });
         }
-        if (isInitializationRequest && messages.length > 1) {
+        if (!!maybeInitializeRequest && messages.length > 1) {
           const body2 = JSON.stringify({
             error: {
               code: -32600,
@@ -139,7 +137,7 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
           });
           return new Response(body2, { status: 400 });
         }
-        if (!isInitializationRequest && !sessionId) {
+        if (!maybeInitializeRequest && !sessionId) {
           const body2 = JSON.stringify({
             error: {
               code: -32e3,
@@ -156,9 +154,9 @@ var createStreamingHttpHandler = (basePath, namespace, corsOptions) => {
           `streamable-http:${sessionId}`,
           { props: ctx.props }
         );
-        const isInitialized = await agent.isInitialized();
-        if (isInitializationRequest) {
-          await agent.setInitialized();
+        const isInitialized = await agent.getInitializeRequest();
+        if (maybeInitializeRequest) {
+          await agent.setInitializeRequest(maybeInitializeRequest);
         } else if (!isInitialized) {
           const body2 = JSON.stringify({
             error: {
@@ -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,28 +199,20 @@ 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();
+                await writer.close().catch(() => {
+                });
               }
             } catch (error) {
               console.error("Error forwarding message to SSE:", error);
@@ -230,20 +222,15 @@ data: ${JSON.stringify(result.data)}
         });
         ws.addEventListener("error", (error) => {
           async function onError(_error) {
-            try {
-              await writer.close();
-            } catch (_e) {
-            }
+            await writer.close().catch(() => {
+            });
           }
           onError(error).catch(console.error);
         });
         ws.addEventListener("close", () => {
           async function onClose() {
-            try {
-              await writer.close();
-            } catch (error) {
-              console.error("Error closing SSE connection:", error);
-            }
+            await writer.close().catch(() => {
+            });
           }
           onClose().catch(console.error);
         });
@@ -251,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",
@@ -289,9 +267,19 @@ data: ${JSON.stringify(result.data)}
           });
           return new Response(body2, { status: 406 });
         }
-        const sessionId = url.searchParams.get("sessionId");
+        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();
@@ -300,7 +288,7 @@ data: ${JSON.stringify(result.data)}
           `streamable-http:${sessionId}`,
           { props: ctx.props }
         );
-        const isInitialized = await agent.isInitialized();
+        const isInitialized = await agent.getInitializeRequest();
         if (!isInitialized) {
           return new Response(
             JSON.stringify({
@@ -320,6 +308,7 @@ data: ${JSON.stringify(result.data)}
           new Request(request.url, {
             headers: {
               ...existingHeaders,
+              [MCP_HTTP_METHOD_HEADER]: "GET",
               Upgrade: "websocket"
             }
           })
@@ -332,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) {
@@ -393,7 +373,7 @@ data: ${JSON.stringify(parsed.data)}
           namespace,
           `streamable-http:${sessionId}`
         );
-        const isInitialized = await agent.isInitialized();
+        const isInitialized = await agent.getInitializeRequest();
         if (!isInitialized) {
           return new Response(
             JSON.stringify({
@@ -592,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) {
@@ -624,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);
   }
 };
 
@@ -685,25 +798,14 @@ 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
    */
-  async setInitialized() {
-    await this.ctx.storage.put("initialized", true);
+  async setInitializeRequest(initializeRequest) {
+    await this.ctx.storage.put("initializeRequest", initializeRequest);
   }
-  async isInitialized() {
-    return await this.ctx.storage.get("initialized") === true;
+  async getInitializeRequest() {
+    return this.ctx.storage.get("initializeRequest");
   }
   /** Read the transport type for this agent.
    * This relies on the naming scheme being `sse:${sessionId}`
@@ -722,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() {
@@ -735,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()) {
@@ -750,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({});
       }
     }
   }
@@ -763,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
    */
@@ -774,9 +878,10 @@ var McpAgent = class _McpAgent extends Agent {
     const server = await this.server;
     this._transport = this.initTransport();
     await server.connect(this._transport);
+    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(