@mastra/mcp 2.9.0-alpha.1 → 2.9.0-alpha.2

package/dist/index.js

@@ -15844,6 +15844,16 @@ function shouldDetachPersistentTransportRequest(init) {
 		return false;
 	}
 }
+/** Whether a request body is a single JSON-RPC response: the client's reply to a server request. */
+function isJsonRpcResponseBody(body) {
+	if (typeof body !== "string") return false;
+	try {
+		const message = JSON.parse(body);
+		return typeof message === "object" && message !== null && !Array.isArray(message) && !("method" in message) && "id" in message && ("result" in message || "error" in message);
+	} catch {
+		return false;
+	}
+}
 const TASK_WAIT_ERROR_IDS = {
 	timeout: "MCP_CLIENT_TOOL_TASK_TIMEOUT",
 	expired: "MCP_CLIENT_TOOL_TASK_EXPIRED",
@@ -16072,6 +16082,19 @@ var InternalMastraMCPClient = class extends MastraBase {
 	clientConnectionOnClose;
 	_authState;
 	operationContextStore = new AsyncLocalStorage();
+	/**
+	* The requestContext of the tool call whose execution opened the connection. Only the
+	* transport's fetch reads it, so the handshake and the streams the connection keeps open are
+	* sent as that call. Messages arriving on those streams are delivered outside it, so server logs
+	* are not labelled with it and requests their handlers send are not sent as that call.
+	*/
+	connectionContextStore = new AsyncLocalStorage();
+	/**
+	* The connection context of the transport a server message arrived on, as it was when that
+	* transport was created, set while the message is delivered. The transport's fetch reads it only for the client's replies to server requests,
+	* which belong to the connection rather than to the handler that produced them.
+	*/
+	replyContextStore = new AsyncLocalStorage();
 	exitHookUnsubscribe;
 	sigTermHandler;
 	sigHupHandler;
@@ -16264,7 +16287,7 @@ var InternalMastraMCPClient = class extends MastraBase {
 		if (allowedHosts !== void 0) assertHostAllowed(url, allowedHosts);
 		const policyUserFetch = userFetch && allowedHosts !== void 0 ? wrapFetchWithHostPolicy(userFetch, allowedHosts) : void 0;
 		const fetch = (requestUrl, init) => {
-			const requestContext = this.operationContextStore.getStore() ?? null;
+			const requestContext = this.fetchRequestContext(init);
 			const executeFetch = () => {
 				if (allowedHosts === void 0) return userFetch ? userFetch(requestUrl, init, requestContext) : globalThis.fetch(requestUrl, init);
 				if (policyUserFetch) return policyUserFetch(requestUrl, init, requestContext);
@@ -16279,6 +16302,7 @@ var InternalMastraMCPClient = class extends MastraBase {
 			fetch
 		});
 		observeEventStreamRequests(transport, this.noteEventStreamSent);
+		this.deliverOutsideConnectionContext(transport);
 		this.connectingTransport = transport;
 		try {
 			this.assertSameLifecycle(generation);
@@ -16300,6 +16324,59 @@ var InternalMastraMCPClient = class extends MastraBase {
 		}
 	}
 	/**
+	* Connects as a tool call: the handshake and the streams the connection keeps open reach the
+	* transport's fetch with the call's requestContext, through the connection context. Any enclosing
+	* operation scope is exited first. It would take priority in fetch, and the new streams and stdio
+	* child would deliver every later server message in it.
+	*/
+	connectAsCall(requestContext, connect) {
+		return this.operationContextStore.exit(() => this.connectionContextStore.run(requestContext, connect));
+	}
+	/**
+	* The requestContext the transport's fetch passes on for an outgoing request. An operation scope
+	* owns its requests, even when its requestContext is null. Outside one, the handshake and the
+	* streams the connection keeps open use the connection context, and a reply to a server request
+	* uses the context of the connection that delivered the request.
+	*/
+	fetchRequestContext(init) {
+		const operationContext = this.operationContextStore.getStore();
+		if (operationContext !== void 0) return operationContext;
+		const connectionContext = this.connectionContextStore.getStore();
+		if (connectionContext !== void 0) return connectionContext;
+		const replyContext = this.replyContextStore.getStore();
+		if (replyContext !== void 0 && isJsonRpcResponseBody(init?.body)) return replyContext;
+		return null;
+	}
+	/**
+	* Delivers every message the server sends outside the connection context, from the handshake on.
+	* A server log, notification or request arriving on a stream the opening call started belongs to
+	* no call, so a request its handler sends must not be sent as that call. The client's reply to a
+	* server request still is, through `replyContextStore`.
+	*
+	* The reply context is the connection context this transport was created in, captured once. Read
+	* at delivery time, it would carry over to a connection that a handler of an older connection's
+	* message opens, such as a `reconnectServer()` from a server log callback.
+	*
+	* The SDK assigns `onmessage` during `connect()` and swaps it while probing the server, so every
+	* assignment is wrapped. A handler that wraps an already wrapped one (`interceptTaskResults`, or
+	* the probe restoring its saved handler) runs with the same reply context.
+	*/
+	deliverOutsideConnectionContext(transport) {
+		const replyContext = this.connectionContextStore.getStore() ?? null;
+		let deliver;
+		Object.defineProperty(transport, "onmessage", {
+			configurable: true,
+			enumerable: true,
+			get: () => deliver,
+			set: (handler) => {
+				deliver = handler && ((message, extra) => {
+					const dispatch = () => handler.call(transport, message, extra);
+					this.connectionContextStore.exit(() => this.replyContextStore.run(replyContext, dispatch));
+				});
+			}
+		});
+	}
+	/**
 	* Detaches whatever transport is still attached to the underlying SDK Client.
 	*
 	* The SDK assigns its internal `_transport` before `transport.start()` and never
@@ -16637,7 +16714,8 @@ var InternalMastraMCPClient = class extends MastraBase {
 				this.serverImplementation = void 0;
 			}
 			this.assertSameLifecycle(generation);
-			await this.connect(generation);
+			const callContext = this.operationContextStore.getStore();
+			await (callContext === void 0 ? this.connect(generation) : this.connectAsCall(callContext, () => this.connect(generation)));
 			this.log("debug", "Successfully reconnected to MCP server");
 		})();
 		this.reconnectPromise = reconnectPromise;
@@ -17130,8 +17208,8 @@ var InternalMastraMCPClient = class extends MastraBase {
 				},
 				...tool.outputSchema ? { toModelOutput: createStructuredToolToModelOutput() } : {},
 				execute: async (input, context) => {
-					if (serverMeta.connectFirst) await this.connect();
 					const operationContext = context?.requestContext ?? null;
+					if (serverMeta.connectFirst) await this.connectAsCall(operationContext, () => this.connect());
 					return this.operationContextStore.run(operationContext, async () => {
 						const suspendHost = context?.agent ?? context?.workflow;
 						const toolCallSpan = context?.tracingContext?.currentSpan;