@adcp/sdk 15.4.0 → 15.5.0

package/dist/lib/protocols/mcp-modern.js

@@ -31,6 +31,7 @@ var import_errors = require('../errors/index.js');
 var import_tracing = require('../observability/tracing.js');
 var import_mcp_scope = require('./mcp-scope.js');
 var import_session_termination = require('./session-termination.js');
+var import_mcp_call_context = require('./mcp-call-context.js');
 var import_client2 = require('../signing/client.js');
 var import_abort = require('./abort.js');
 var import_rawResponseCapture = require('./rawResponseCapture.js');
@@ -50,10 +51,8 @@ const LEGACY_CLASSIFICATION_TTL_MS = 5 * 60 * 1e3;
 const MODERN_DISCOVERY_TTL_MS = 5 * 60 * 1e3;
 const modernOAuthProviderIds = /* @__PURE__ */ new WeakMap();
 const modernFetchFnIds = /* @__PURE__ */ new WeakMap();
-const modernSignalIds = /* @__PURE__ */ new WeakMap();
 let nextModernOAuthProviderId = 0;
 let nextModernFetchFnId = 0;
-let nextModernSignalId = 0;
 let connectionGeneration = 0;
 async function closeModernClient(client, terminateSession = true) {
   const transport = modernTransports.get(client);
@@ -73,7 +72,7 @@ function buildAuthHeaders(authToken, customHeaders, authProvider) {
       return key !== "authorization" && key !== "x-adcp-auth";
     })
   ) : customHeaders;
-  return (0, import_auth.createMCPRequestHeaders)(filteredHeaders, authProvider ? void 0 : authToken);
+  return (0, import_mcp_call_context.withAmbientIdentityHeaders)((0, import_auth.createMCPRequestHeaders)(filteredHeaders, authProvider ? void 0 : authToken));
 }
 function oauthProviderCacheKey(provider) {
   if (!provider) return void 0;
@@ -93,33 +92,26 @@ function fetchFnCacheKey(fetchFn) {
   }
   return key;
 }
-function signalCacheKey(signal) {
-  if (!signal) return void 0;
-  let key = modernSignalIds.get(signal);
-  if (!key) {
-    key = `signal:${++nextModernSignalId}`;
-    modernSignalIds.set(signal, key);
-  }
-  return key;
-}
-function connectionCacheKey(agentUrl, headers, signingCacheKey, authProvider, fetchFn, signal, requestTimeoutMs, handleLegacy, allowPrivateIp) {
-  const normalizedHeaders = Object.entries(headers).map(([key, value]) => [key.toLowerCase(), value]).sort(([left], [right]) => left.localeCompare(right));
+function connectionCacheKey(agentUrl, headers, signingCacheKey, authProvider, fetchFn, handleLegacy, allowPrivateIp) {
+  const normalizedHeaders = Object.entries((0, import_mcp_call_context.splitConnectionHeaders)(headers).identity).map(([key, value]) => [key.toLowerCase(), value]).sort(([left], [right]) => left.localeCompare(right));
   const parts = [agentUrl, `headers:${cacheDisambiguator(JSON.stringify(normalizedHeaders))}`];
   if (signingCacheKey) parts.push(signingCacheKey);
   const providerKey = oauthProviderCacheKey(authProvider);
   if (providerKey) parts.push(providerKey);
   const fetchKey = fetchFnCacheKey(fetchFn);
   if (fetchKey) parts.push(fetchKey);
-  const signalKey = signalCacheKey(signal);
-  if (signalKey) parts.push(signalKey);
-  if (requestTimeoutMs !== void 0) parts.push(`timeout:${requestTimeoutMs}`);
   if (handleLegacy !== void 0) parts.push(`handle-legacy:${handleLegacy}`);
   if (allowPrivateIp !== void 0) parts.push(`allow-private-ip:${allowPrivateIp}`);
   const scopeKey = (0, import_mcp_scope.currentMCPConnectionScopeKey)();
   if (scopeKey) parts.push(scopeKey);
   return parts.join("::");
 }
-function isKnownLegacy(cacheKey) {
+function classificationKey(cacheKey) {
+  const marker = cacheKey.lastIndexOf("::scope:");
+  return marker >= 0 ? cacheKey.slice(0, marker) : cacheKey;
+}
+function isKnownLegacy(connectionKey) {
+  const cacheKey = classificationKey(connectionKey);
   const classifiedAt = knownLegacyConnections.get(cacheKey);
   if (classifiedAt === void 0) return false;
   if (Date.now() - classifiedAt > LEGACY_CLASSIFICATION_TTL_MS) {
@@ -130,7 +122,9 @@ function isKnownLegacy(cacheKey) {
   knownLegacyConnections.set(cacheKey, classifiedAt);
   return true;
 }
-function markKnownLegacy(cacheKey) {
+function markKnownLegacy(connectionKey) {
+  const cacheKey = classificationKey(connectionKey);
+  modernDiscoveries.delete(connectionKey);
   modernDiscoveries.delete(cacheKey);
   knownLegacyConnections.delete(cacheKey);
   knownLegacyConnections.set(cacheKey, Date.now());
@@ -152,6 +146,7 @@ function getCachedModernDiscovery(cacheKey) {
   return { kind: "modern", discover: cached.discover };
 }
 function cacheModernDiscovery(cacheKey, discover) {
+  knownLegacyConnections.delete(classificationKey(cacheKey));
   modernDiscoveries.delete(cacheKey);
   modernDiscoveries.set(cacheKey, {
     discover,
@@ -178,8 +173,12 @@ function withPerRequestTraceHeaders(fetchImpl) {
   return (input, init) => {
     const headers = new Headers(input instanceof Request ? input.headers : void 0);
     new Headers(init?.headers).forEach((value, key) => headers.set(key, value));
-    for (const [key, value] of Object.entries((0, import_tracing.injectTraceHeaders)())) headers.set(key, value);
-    return fetchImpl(input, { ...init, headers });
+    (0, import_mcp_call_context.applyMCPCallHeaders)(headers);
+    if (headers.has("traceparent")) return fetchImpl(input, { ...init, headers, ...(0, import_mcp_call_context.linkedCallSignal)(init) });
+    for (const [key, value] of Object.entries((0, import_tracing.injectTraceHeaders)())) {
+      if ((0, import_mcp_call_context.isCorrelationHeader)(key) && !headers.has(key)) headers.set(key, value);
+    }
+    return fetchImpl(input, { ...init, headers, ...(0, import_mcp_call_context.linkedCallSignal)(init) });
   };
 }
 function getCachedConnection(cacheKey) {
@@ -189,7 +188,7 @@ function getCachedConnection(cacheKey) {
     if (legacyExpiry !== void 0 && legacyExpiry <= Date.now()) {
       modernConnections.delete(cacheKey);
       legacyConnectionExpiresAt.delete(cacheKey);
-      void closeModernClient(client).catch(() => {
+      void (0, import_mcp_call_context.closeWhenIdle)(client, () => closeModernClient(client)).catch(() => {
       });
       return void 0;
     }
@@ -205,7 +204,7 @@ function evictLeastRecentlyUsed() {
   const client = modernConnections.get(oldestKey);
   modernConnections.delete(oldestKey);
   legacyConnectionExpiresAt.delete(oldestKey);
-  if (client) void closeModernClient(client).catch(() => {
+  if (client) void (0, import_mcp_call_context.closeWhenIdle)(client, () => closeModernClient(client)).catch(() => {
   });
 }
 function isNegotiationAuthFailure(error) {
@@ -220,13 +219,24 @@ function presentedStaticCredential(authHeaders, authProvider) {
 }
 async function createNegotiatedClient(cacheKey, options, authHeaders, useCachedDiscovery = true, skipProbe = false) {
   const generation = connectionGeneration;
-  const requestTimeoutMs = (0, import_abort.resolveRequestTimeoutMs)(options.requestTimeoutMs);
   const clientRequestTimeoutMs = (0, import_abort.resolveClientRequestTimeoutMs)(options.requestTimeoutMs);
   const rawNetworkFetch = (0, import_agent_transport_fetch.createAgentTransportFetch)(options.agentUrl, {
     trustedFetchFn: options.fetchFn,
     allowPrivateIp: options.allowPrivateIp
   });
-  const networkFetch = (input, init) => (0, import_abort.withAbortSignal)([init?.signal], requestTimeoutMs, (signal) => rawNetworkFetch(input, { ...init, signal }));
+  const networkFetch = (input, init) => (
+    // `init.signal` already carries the active call's cancellation (linked at the
+    // outermost transport wrapper), including on legacy-era sessions where the
+    // v2 client cancels with a notification and leaves the HTTP stream open.
+    (0, import_abort.withAbortSignal)(
+      [init?.signal],
+      (0, import_mcp_call_context.requestTimeoutFor)(init),
+      (deadline) => rawNetworkFetch(input, {
+        ...init,
+        signal: init?.signal && deadline ? AbortSignal.any([init.signal, deadline]) : deadline ?? init?.signal
+      })
+    )
+  );
   const diagnosticFetch = (0, import_transportDiagnostics.wrapFetchWithTransportDiagnostics)((0, import_responseSizeLimit.wrapFetchWithSizeLimit)(networkFetch));
   const signedFetch = options.signingContext ? (0, import_client2.buildAgentSigningFetch)({
     upstream: (0, import_signedRequestRejection.wrapFetchWithSignedRequestRejection)(diagnosticFetch, options.agentUrl),
@@ -235,7 +245,7 @@ async function createNegotiatedClient(cacheKey, options, authHeaders, useCachedD
     adcpVersion: options.signingContext.adcpVersion
   }) : diagnosticFetch;
   const transport = new import_client.StreamableHTTPClientTransport(new URL(options.agentUrl), {
-    requestInit: { headers: authHeaders, redirect: "manual" },
+    requestInit: { headers: (0, import_mcp_call_context.splitConnectionHeaders)(authHeaders).identity, redirect: "manual" },
     fetch: (0, import_rawResponseCapture.wrapFetchWithCapture)(withPerRequestTraceHeaders(signedFetch)),
     ...options.authProvider && {
       authProvider: options.authProvider
@@ -283,10 +293,19 @@ async function createNegotiatedClient(cacheKey, options, authHeaders, useCachedD
   }
 }
 async function getOrCreateModernConnection(cacheKey, options, authHeaders) {
-  const cached = getCachedConnection(cacheKey);
-  if (cached) return cached;
-  const pending = pendingModernConnections.get(cacheKey);
-  if (pending) return pending;
+  for (let attempt = 0; ; attempt++) {
+    const cached = getCachedConnection(cacheKey);
+    if (cached) return cached;
+    const pending = pendingModernConnections.get(cacheKey);
+    if (pending) {
+      const joined = await (0, import_mcp_call_context.joinPendingConnection)(pending, options, attempt);
+      if (joined === import_mcp_call_context.RETRY_CONNECT) continue;
+      return joined;
+    }
+    return createSharedModernConnection(cacheKey, options, authHeaders);
+  }
+}
+function createSharedModernConnection(cacheKey, options, authHeaders) {
   const generation = connectionGeneration;
   const promise = createNegotiatedClient(cacheKey, options, authHeaders).then(async (client) => {
     if ((client.getProtocolEra() === "modern" || options.handleLegacy === true) && generation === connectionGeneration) {
@@ -317,6 +336,16 @@ async function getOrCreateModernConnection(cacheKey, options, authHeaders) {
   (0, import_mcp_scope.registerMCPConnectionScopePending)(promise);
   return promise;
 }
+async function retireFailedConnection(cacheKey, client, error) {
+  if ((0, import_abort.isAbortOrTimeoutError)(error)) return;
+  if (modernConnections.get(cacheKey) === client) {
+    modernConnections.delete(cacheKey);
+    legacyConnectionExpiresAt.delete(cacheKey);
+  }
+  modernDiscoveries.delete(cacheKey);
+  await (0, import_mcp_call_context.closeWhenIdle)(client, () => closeModernClient(client, true)).catch(() => {
+  });
+}
 async function callOnModernClient(client, toolName, args, signal, requestTimeoutMs) {
   const resolvedRequestTimeoutMs = (0, import_abort.resolveClientRequestTimeoutMs)(requestTimeoutMs);
   return await client.callTool(
@@ -335,14 +364,11 @@ async function attemptModernCall(options, toolName, args) {
     options.signingContext?.cacheKey,
     options.authProvider,
     options.fetchFn,
-    options.signal,
-    options.requestTimeoutMs,
     options.handleLegacy,
     options.allowPrivateIp
   );
   if (isKnownLegacy(cacheKey)) return { handled: false };
-  const guardedConnection = options.signal !== void 0 || options.requestTimeoutMs !== void 0 || options.fetchFn !== void 0;
-  const oneShot = guardedConnection && (0, import_mcp_scope.currentMCPConnectionScopeKey)() === void 0;
+  const oneShot = options.fetchFn !== void 0 && (0, import_mcp_scope.currentMCPConnectionScopeKey)() === void 0;
   let client;
   let callSucceeded = false;
   try {
@@ -396,16 +422,21 @@ async function attemptModernCall(options, toolName, args) {
     timestamp: (/* @__PURE__ */ new Date()).toISOString()
   });
   try {
-    const response = await callOnModernClient(client, toolName, args, options.signal, options.requestTimeoutMs);
+    const response = await (0, import_mcp_call_context.trackConnectionUse)(
+      client,
+      () => (0, import_mcp_call_context.runCallPhase)(
+        options.signal,
+        () => callOnModernClient(client, toolName, args, options.signal, options.requestTimeoutMs)
+      )
+    );
     callSucceeded = true;
     return { handled: true, response };
   } catch (error) {
-    modernConnections.delete(cacheKey);
-    legacyConnectionExpiresAt.delete(cacheKey);
-    modernDiscoveries.delete(cacheKey);
-    try {
-      await closeModernClient(client, !(0, import_abort.isAbortOrTimeoutError)(error));
-    } catch {
+    if (oneShot) {
+      await closeModernClient(client, !(0, import_abort.isAbortOrTimeoutError)(error)).catch(() => {
+      });
+    } else {
+      await retireFailedConnection(cacheKey, client, error);
     }
     throw error;
   } finally {
@@ -413,28 +444,39 @@ async function attemptModernCall(options, toolName, args) {
     });
   }
 }
+function callContextFor(options) {
+  const { perRequest } = (0, import_mcp_call_context.splitConnectionHeaders)(options.customHeaders);
+  const requestTimeoutMs = (0, import_abort.resolveRequestTimeoutMs)(options.requestTimeoutMs);
+  return {
+    ...Object.keys(perRequest).length > 0 && { headers: perRequest },
+    ...requestTimeoutMs !== void 0 && { requestTimeoutMs }
+  };
+}
 async function tryCallModernMCPTool(agentUrl, toolName, args, authToken, debugLogs = [], customHeaders, options = {}) {
   return (0, import_tracing.withSpan)(
     "adcp.mcp.negotiate",
     { "adcp.tool": toolName, "http.url": agentUrl },
     () => import_client2.signingContextStorage.run(
       options.signingContext,
-      () => attemptModernCall(
-        {
-          agentUrl,
-          authToken,
-          customHeaders,
-          debugLogs,
-          signingContext: options.signingContext,
-          authProvider: options.authProvider,
-          signal: options.signal,
-          requestTimeoutMs: options.requestTimeoutMs,
-          fetchFn: options.fetchFn,
-          allowPrivateIp: options.allowPrivateIp,
-          handleLegacy: options.handleLegacy
-        },
-        toolName,
-        args
+      () => (0, import_mcp_call_context.runWithMCPCallContext)(
+        callContextFor({ customHeaders, requestTimeoutMs: options.requestTimeoutMs }),
+        () => attemptModernCall(
+          {
+            agentUrl,
+            authToken,
+            customHeaders,
+            debugLogs,
+            signingContext: options.signingContext,
+            authProvider: options.authProvider,
+            signal: options.signal,
+            requestTimeoutMs: options.requestTimeoutMs,
+            fetchFn: options.fetchFn,
+            allowPrivateIp: options.allowPrivateIp,
+            handleLegacy: options.handleLegacy
+          },
+          toolName,
+          args
+        )
       )
     )
   );
@@ -460,26 +502,34 @@ async function probeModernMCPConnection(agentUrl, authToken, customHeaders, opti
     options.signingContext?.cacheKey,
     options.authProvider,
     options.fetchFn,
-    options.signal,
-    options.requestTimeoutMs,
     options.handleLegacy,
     options.allowPrivateIp
   );
+  const scoped = (0, import_mcp_scope.currentMCPConnectionScopeKey)() !== void 0;
   let client;
-  try {
-    client = await createNegotiatedClient(cacheKey, connectionOptions, authHeaders, false);
-    return { connected: true, era: client.getProtocolEra() };
-  } catch (error) {
-    modernDiscoveries.delete(cacheKey);
-    if ((0, import_errors.is401Error)(error) || (0, import_abort.isAbortOrTimeoutError)(error)) throw error;
-    const status = httpStatusOf(error);
-    if (status === 404 || status === 405) return { connected: false };
-    if (isLegacyEraNegotiationFailure(error, status)) return { connected: false };
-    throw error;
-  } finally {
-    if (client) await closeModernClient(client).catch(() => {
-    });
-  }
+  return (0, import_mcp_call_context.runWithMCPCallContext)(
+    callContextFor({ customHeaders, requestTimeoutMs: options.requestTimeoutMs }),
+    async () => {
+      try {
+        client = scoped ? await getOrCreateModernConnection(cacheKey, connectionOptions, authHeaders) : await createNegotiatedClient(cacheKey, connectionOptions, authHeaders, false);
+        const era = client.getProtocolEra();
+        if (scoped && era === "legacy" && connectionOptions.handleLegacy !== true) markKnownLegacy(cacheKey);
+        return { connected: true, era };
+      } catch (error) {
+        modernDiscoveries.delete(cacheKey);
+        if ((0, import_errors.is401Error)(error) || (0, import_abort.isAbortOrTimeoutError)(error)) throw error;
+        const status = httpStatusOf(error);
+        if (status === 404 || status === 405) return { connected: false };
+        if (isLegacyEraNegotiationFailure(error, status)) return { connected: false };
+        throw error;
+      } finally {
+        if (client && !(scoped && modernConnections.get(cacheKey) === client)) {
+          await closeModernClient(client).catch(() => {
+          });
+        }
+      }
+    }
+  );
 }
 async function tryListModernMCPTools(agentUrl, authToken, customHeaders, options = {}) {
   const connectionOptions = {
@@ -492,7 +542,8 @@ async function tryListModernMCPTools(agentUrl, authToken, customHeaders, options
     signal: options.signal,
     requestTimeoutMs: options.requestTimeoutMs,
     fetchFn: options.fetchFn,
-    allowPrivateIp: options.allowPrivateIp
+    allowPrivateIp: options.allowPrivateIp,
+    handleLegacy: options.handleLegacy
   };
   const authHeaders = buildAuthHeaders(authToken, customHeaders, options.authProvider);
   const cacheKey = connectionCacheKey(
@@ -501,50 +552,75 @@ async function tryListModernMCPTools(agentUrl, authToken, customHeaders, options
     options.signingContext?.cacheKey,
     options.authProvider,
     options.fetchFn,
-    options.signal,
-    options.requestTimeoutMs,
     options.handleLegacy,
     options.allowPrivateIp
   );
+  const scoped = (0, import_mcp_scope.currentMCPConnectionScopeKey)() !== void 0;
+  if (scoped && isKnownLegacy(cacheKey)) return { handled: false };
   let client;
+  const sharedClients = /* @__PURE__ */ new WeakSet();
+  const acquire = async (create) => {
+    const acquired = await create();
+    if (scoped && modernConnections.get(cacheKey) === acquired) sharedClients.add(acquired);
+    return acquired;
+  };
   const listTools = async (connectedClient) => {
-    if (connectedClient.getProtocolEra() !== "modern") return { handled: false };
+    if (connectedClient.getProtocolEra() !== "modern" && connectionOptions.handleLegacy !== true) {
+      if (scoped) markKnownLegacy(cacheKey);
+      return { handled: false };
+    }
     const resolvedRequestTimeoutMs = (0, import_abort.resolveClientRequestTimeoutMs)(options.requestTimeoutMs);
-    const result = await connectedClient.listTools(
-      options.adcpVersion === void 0 ? void 0 : { _meta: { adcp_version: options.adcpVersion } },
-      {
-        ...options.signal && { signal: options.signal },
-        ...resolvedRequestTimeoutMs !== void 0 && { timeout: resolvedRequestTimeoutMs }
-      }
+    const result = await (0, import_mcp_call_context.trackConnectionUse)(
+      connectedClient,
+      () => (0, import_mcp_call_context.runCallPhase)(
+        options.signal,
+        () => connectedClient.listTools(
+          options.adcpVersion === void 0 ? void 0 : { _meta: { adcp_version: options.adcpVersion } },
+          {
+            ...options.signal && { signal: options.signal },
+            ...resolvedRequestTimeoutMs !== void 0 && { timeout: resolvedRequestTimeoutMs }
+          }
+        )
+      )
     );
     return { handled: true, tools: result.tools };
   };
-  try {
-    client = await createNegotiatedClient(cacheKey, connectionOptions, authHeaders);
-    return await listTools(client);
-  } catch (error) {
-    let failure = error;
-    modernDiscoveries.delete(cacheKey);
-    if (!(0, import_errors.is401Error)(failure) && !(0, import_abort.isAbortOrTimeoutError)(failure) && client && clientsUsingCachedDiscovery.has(client)) {
-      await closeModernClient(client, false).catch(() => {
-      });
+  return (0, import_mcp_call_context.runWithMCPCallContext)(
+    callContextFor({ customHeaders, requestTimeoutMs: options.requestTimeoutMs }),
+    async () => {
       try {
-        client = await createNegotiatedClient(cacheKey, connectionOptions, authHeaders, false);
+        client = await acquire(
+          () => scoped ? getOrCreateModernConnection(cacheKey, connectionOptions, authHeaders) : createNegotiatedClient(cacheKey, connectionOptions, authHeaders)
+        );
         return await listTools(client);
-      } catch (retryError) {
-        failure = retryError;
+      } catch (error) {
+        let failure = error;
         modernDiscoveries.delete(cacheKey);
+        if (client && sharedClients.has(client)) {
+          await retireFailedConnection(cacheKey, client, failure);
+        }
+        if (!(0, import_errors.is401Error)(failure) && !(0, import_abort.isAbortOrTimeoutError)(failure) && client && clientsUsingCachedDiscovery.has(client)) {
+          if (!sharedClients.has(client)) await closeModernClient(client, false).catch(() => {
+          });
+          try {
+            client = await acquire(() => createNegotiatedClient(cacheKey, connectionOptions, authHeaders, false));
+            return await listTools(client);
+          } catch (retryError) {
+            failure = retryError;
+            modernDiscoveries.delete(cacheKey);
+          }
+        }
+        if ((0, import_errors.is401Error)(failure) || (0, import_abort.isAbortOrTimeoutError)(failure)) throw failure;
+        const status = httpStatusOf(failure);
+        if (status === 404 || status === 405) return { handled: false };
+        if (isLegacyEraNegotiationFailure(failure, status)) return { handled: false };
+        throw failure;
+      } finally {
+        if (client && !sharedClients.has(client)) await closeModernClient(client).catch(() => {
+        });
       }
     }
-    if ((0, import_errors.is401Error)(failure) || (0, import_abort.isAbortOrTimeoutError)(failure)) throw failure;
-    const status = httpStatusOf(failure);
-    if (status === 404 || status === 405) return { handled: false };
-    if (isLegacyEraNegotiationFailure(failure, status)) return { handled: false };
-    throw failure;
-  } finally {
-    if (client) await closeModernClient(client).catch(() => {
-    });
-  }
+  );
 }
 async function closeModernMCPConnections() {
   connectionGeneration++;