dd-trace 6.2.0 → 6.4.0

package/packages/datadog-instrumentations/src/jest.js

@@ -1861,10 +1861,16 @@ function resetSuiteSkippingRunState () {
 function applySuiteSkipping (originalTests, rootDir, frameworkVersion) {
   if (!isItrEnabled || !isSuitesSkippingEnabled) return originalTests
 
-  const suitePathRoot = getRepositoryRootFromTest(originalTests[0], rootDir)
-  const jestSuitesToRun = getJestSuitesToRun(skippableSuites, originalTests, suitePathRoot)
+  const repositoryRoot = getRepositoryRootFromTest(originalTests[0], rootDir)
+  const jestSuitesToRun = getJestSuitesToRun(skippableSuites, originalTests, repositoryRoot, rootDir)
   hasFilteredSkippableSuites = true
-  log.debug('%d out of %d suites are going to run.', jestSuitesToRun.suitesToRun.length, originalTests.length)
+  log.debug(
+    'Jest discovered %d suites; skipped %d using %d skippable candidates; %d suites will run.',
+    originalTests.length,
+    jestSuitesToRun.skippedSuites.length,
+    skippableSuites.length,
+    jestSuitesToRun.suitesToRun.length
+  )
   hasUnskippableSuites = jestSuitesToRun.hasUnskippableSuites
   hasForcedToRunSuites = jestSuitesToRun.hasForcedToRunSuites
 
@@ -1877,7 +1883,7 @@ function applySuiteSkipping (originalTests, rootDir, frameworkVersion) {
     ? getTestContexts(originalTests)
     : undefined
   coverageBackfillFiles = isSuitesSkipped && isTiaCoverageBackfillEnabled() && hasSkippableSuitesCoverage()
-    ? getCoverageBackfillFiles(skippableSuitesCoverage, suitePathRoot, getTestSuitePath)
+    ? getCoverageBackfillFiles(skippableSuitesCoverage, repositoryRoot, getTestSuitePath)
     : undefined
 
   itrSkippedSuitesCh.publish({ skippedSuites: jestSuitesToRun.skippedSuites, frameworkVersion })

package/packages/datadog-instrumentations/src/modelcontextprotocol-sdk.js

@@ -1,7 +1,254 @@
 'use strict'
 
-const { addHook, getHooks } = require('./helpers/instrument')
+const { tracingChannel } = require('dc-polyfill')
+const shimmer = require('../../datadog-shimmer')
+const { addHook, channel, getHooks } = require('./helpers/instrument')
 
 for (const hook of getHooks('@modelcontextprotocol/sdk')) {
   addHook(hook, exports => exports)
 }
+
+const serverRequestCh = tracingChannel('apm:mcp:server:request')
+const clientRequestInjectCh = channel('apm:mcp:client:request:inject')
+const serverToolRegisteredCh = channel('apm:mcp:server:tool:registered')
+
+const DISTRIBUTED_TRACE_META_KEY = '_dd_trace_context'
+const TRACED_METHOD_PREFIX = /^(?:tools|resources|prompts)\//
+
+// Maps Protocol instance → Map<requestId, ctx>. Shares the ctx object between
+// _onrequest (span start, in the correct HTTP async context), the SDK request
+// handler wrapper (result/error capture), and the SDK cleanup path (span finish).
+const pendingContexts = new WeakMap()
+const wrappedAbortControllerMaps = new WeakSet()
+const wrappedRequestHandlerMaps = new WeakSet()
+const wrappedRequestHandlers = new WeakSet()
+
+function addTraceContextToRequest (request, traceContext) {
+  return {
+    ...request,
+    params: {
+      ...request.params,
+      _meta: {
+        ...request.params?._meta,
+        [DISTRIBUTED_TRACE_META_KEY]: traceContext,
+      },
+    },
+  }
+}
+
+function getPendingContext (protocol, requestId) {
+  return pendingContexts.get(protocol)?.get(requestId)
+}
+
+function finishServerRequest (protocol, requestId) {
+  const pending = pendingContexts.get(protocol)
+  const ctx = pending?.get(requestId)
+  if (!ctx) return
+
+  pending.delete(requestId)
+  serverRequestCh.asyncEnd.publish(ctx)
+}
+
+function wrapRequestHandler (protocol, handler) {
+  if (wrappedRequestHandlers.has(handler)) return handler
+
+  const wrappedHandler = function requestHandlerWithTrace (request, extra) {
+    const ctx = getPendingContext(protocol, extra?.requestId)
+    if (!ctx) return handler.apply(this, arguments)
+
+    let result
+    try {
+      result = handler.apply(this, arguments)
+    } catch (err) {
+      ctx.error = err
+      throw err
+    }
+
+    if (result && typeof result.then === 'function') {
+      result.then(result => {
+        ctx.result = result
+      }, err => {
+        ctx.error = err
+      })
+      return result
+    }
+
+    ctx.result = result
+    return result
+  }
+
+  wrappedRequestHandlers.add(wrappedHandler)
+  return wrappedHandler
+}
+
+function wrapRequestHandlers (protocol) {
+  const handlers = protocol._requestHandlers
+  if (!handlers || wrappedRequestHandlerMaps.has(handlers)) return
+
+  wrappedRequestHandlerMaps.add(handlers)
+  shimmer.wrap(handlers, 'set', function (original) {
+    return function setRequestHandlerWithTrace (method, handler) {
+      if (typeof handler === 'function') {
+        return original.call(this, method, wrapRequestHandler(protocol, handler))
+      }
+
+      return original.apply(this, arguments)
+    }
+  })
+}
+
+function wrapAbortControllers (protocol) {
+  const controllers = protocol._requestHandlerAbortControllers
+  if (!controllers || wrappedAbortControllerMaps.has(controllers)) return
+
+  wrappedAbortControllerMaps.add(controllers)
+  shimmer.wrap(controllers, 'delete', function (original) {
+    return function deleteAbortControllerWithTrace (requestId) {
+      finishServerRequest(protocol, requestId)
+      return original.apply(this, arguments)
+    }
+  })
+  shimmer.wrap(controllers, 'clear', function (original) {
+    return function clearAbortControllersWithTrace () {
+      for (const requestId of this.keys()) {
+        finishServerRequest(protocol, requestId)
+      }
+      return original.apply(this, arguments)
+    }
+  })
+}
+
+function wrapProtocol (Protocol) {
+  // Inject trace context into MCP request metadata so out-of-process MCP servers
+  // can parent server spans to the client operation span.
+  shimmer.wrap(Protocol.prototype, 'request', function (original) {
+    return function requestWithTraceContext (request, resultSchema, options) {
+      if (!clientRequestInjectCh.hasSubscribers || !TRACED_METHOD_PREFIX.test(request?.method)) {
+        return original.apply(this, arguments)
+      }
+
+      const ctx = {}
+      clientRequestInjectCh.publish(ctx)
+
+      if (!ctx.traceContext) {
+        return original.apply(this, arguments)
+      }
+
+      return original.call(this, addTraceContextToRequest(request, ctx.traceContext), resultSchema, options)
+    }
+  })
+
+  // Start spans in _onrequest — this runs inside the express POST /mcp handler's
+  // async context, so ALS correctly parents server spans under the HTTP span.
+  shimmer.wrap(Protocol.prototype, '_onrequest', function (original) {
+    return function _onrequestWithTrace (request, extra) {
+      if (!serverRequestCh.hasSubscribers || !TRACED_METHOD_PREFIX.test(request.method)) {
+        return original.call(this, request, extra)
+      }
+
+      wrapAbortControllers(this)
+      const ctx = { request, extra }
+      serverRequestCh.start.runStores(ctx, () => {
+        let pending = pendingContexts.get(this)
+        if (!pending) {
+          pending = new Map()
+          pendingContexts.set(this, pending)
+        }
+        pending.set(request.id, ctx)
+
+        try {
+          original.call(this, request, extra)
+        } catch (err) {
+          ctx.error = err
+          finishServerRequest(this, request.id)
+          throw err
+        }
+
+        // The SDK registers an AbortController only when it actually dispatches a handler.
+        // If none was registered (MethodNotFound path), no handler will run and our
+        // SDK cleanup hook will never fire — finish the span immediately.
+        if (!this._requestHandlerAbortControllers?.has(request.id)) {
+          ctx.error = { name: 'McpError', message: 'Method not found' }
+          finishServerRequest(this, request.id)
+        }
+      })
+    }
+  })
+
+  // The SDK's handler closure Zod-parses the request, stripping the JSON-RPC `id`.
+  // Use `extra.requestId` (the SDK's fullExtra field) to correlate with the pending ctx.
+  shimmer.wrap(Protocol.prototype, 'setRequestHandler', function (original) {
+    return function setRequestHandlerWithTrace () {
+      wrapRequestHandlers(this)
+      return original.apply(this, arguments)
+    }
+  })
+
+  return Protocol
+}
+
+function wrapMcpServer (McpServer) {
+  // Both public registration methods (tool/registerTool) delegate here — one hook covers both.
+  // Publishes to serverToolRegisteredCh so tracing.js can build a WeakMap of tool → name
+  // for O(1) lookup in McpServerToolCallPlugin (executeToolHandler receives the tool object,
+  // not the name string).
+  shimmer.wrap(McpServer.prototype, '_createRegisteredTool', function (original) {
+    return function (name) {
+      const result = original.apply(this, arguments)
+      if (serverToolRegisteredCh.hasSubscribers) {
+        const tool = this._registeredTools?.[name]
+        if (tool) {
+          serverToolRegisteredCh.publish({ tool, name })
+          shimmer.wrap(tool, 'update', function (update) {
+            return function updateWithTrace (updates) {
+              const result = update.apply(this, arguments)
+              if (updates && Object.hasOwn(updates, 'name')) {
+                serverToolRegisteredCh.publish({ tool, name: updates.name || undefined })
+              }
+              return result
+            }
+          })
+        }
+      }
+      return result
+    }
+  })
+
+  return McpServer
+}
+
+addHook({
+  name: '@modelcontextprotocol/sdk',
+  versions: ['>=1.27.1'],
+  file: 'dist/cjs/shared/protocol.js',
+}, exports => {
+  wrapProtocol(exports.Protocol)
+  return exports
+})
+
+addHook({
+  name: '@modelcontextprotocol/sdk',
+  versions: ['>=1.27.1'],
+  file: 'dist/esm/shared/protocol.js',
+}, exports => {
+  wrapProtocol(exports.Protocol)
+  return exports
+})
+
+addHook({
+  name: '@modelcontextprotocol/sdk',
+  versions: ['>=1.27.1'],
+  file: 'dist/cjs/server/mcp.js',
+}, exports => {
+  wrapMcpServer(exports.McpServer)
+  return exports
+})
+
+addHook({
+  name: '@modelcontextprotocol/sdk',
+  versions: ['>=1.27.1'],
+  file: 'dist/esm/server/mcp.js',
+}, exports => {
+  wrapMcpServer(exports.McpServer)
+  return exports
+})

package/packages/datadog-instrumentations/src/playwright.js

@@ -105,6 +105,7 @@ const efdStartedOriginalTestKeys = new Set()
 const efdSlowAbortedTests = new Set()
 const ddPropertiesByTestId = new Map()
 const ddPropertiesRequestsByTestId = new Map()
+const disabledTestIds = new Set()
 let rootDir = ''
 let sessionProjects = []
 
@@ -113,6 +114,8 @@ const EFD_RETRY_COUNT_REQUEST = 'ddEfdRetryCountRequest'
 const EFD_RETRY_COUNT_RESPONSE = 'ddEfdRetryCountResponse'
 const DD_PROPERTIES_REQUEST = 'ddPropertiesRequest'
 const DD_PROPERTIES_RESPONSE = 'ddProperties'
+const kDdPlaywrightDisabledTestIds = Symbol('ddPlaywrightDisabledTestIds')
+const kDdPlaywrightWorkerHostInstrumented = Symbol('ddPlaywrightWorkerHostInstrumented')
 const kDdPlaywrightWorkerInstrumented = Symbol('ddPlaywrightWorkerInstrumented')
 
 function isValidKnownTests (receivedKnownTests) {
@@ -991,14 +994,36 @@ function onDispatcherCreateWorker (dispatcher, worker) {
   const projects = getProjectsFromDispatcher(dispatcher)
   sessionProjects = projects
 
+  if (disabledTestIds.size && !worker[kDdPlaywrightWorkerHostInstrumented] &&
+      typeof worker.runTestGroup === 'function') {
+    Object.defineProperty(worker, kDdPlaywrightWorkerHostInstrumented, { value: true })
+    shimmer.wrap(worker, 'runTestGroup', runTestGroup => function (runPayload) {
+      // Serial retries can restore disabled tests that were filtered from the initial dispatcher groups.
+      let disabledIds
+      for (const { testId } of runPayload.entries) {
+        if (disabledTestIds.has(testId)) {
+          disabledIds ??= []
+          disabledIds.push(testId)
+        }
+      }
+      if (disabledIds) {
+        runPayload._ddDisabledTestIds = disabledIds
+      }
+      return runTestGroup.apply(this, arguments)
+    })
+  }
+
   worker.on('testBegin', ({ testId }) => {
     const test = getTestByTestId(dispatcher, testId)
+    if (!test) return
+
     const browser = getBrowserNameFromProjects(projects, test)
     const shouldCreateTestSpan = test.expectedStatus === 'skipped'
     testBeginHandler(test, browser, shouldCreateTestSpan)
   })
   worker.on('testEnd', ({ testId, status, errors, annotations }) => {
     const test = getTestByTestId(dispatcher, testId)
+    if (!test) return
 
     const isTimeout = status === 'timedOut'
     const testStatus = STATUS_TO_TEST_STATUS[status]
@@ -1308,6 +1333,7 @@ function runAllTestsWrapper (runAllTests, playwrightVersion) {
     efdSlowAbortedTests.clear()
     ddPropertiesByTestId.clear()
     ddPropertiesRequestsByTestId.clear()
+    disabledTestIds.clear()
 
     // TODO: we can trick playwright into thinking the session passed by returning
     // 'passed' here. We might be able to use this for both EFD and Test Management tests.
@@ -1552,6 +1578,7 @@ function processRootSuite (createRootSuiteReturnValue) {
       if (testProperties.disabled) {
         test._ddIsDisabled = true
         if (!testProperties.attemptToFix) {
+          disabledTestIds.add(test.id)
           test.expectedStatus = 'skipped'
           // setting test.expectedStatus to 'skipped' does not work for every case,
           // so we need to filter out disabled tests in dispatcherRunWrapperNew,
@@ -1795,7 +1822,17 @@ function instrumentWorkerMainMethods (workerMain) {
   let steps = []
   const stepInfoByStepId = {}
 
+  shimmer.wrap(workerMain, 'runTestGroup', runTestGroup => function (runPayload) {
+    const disabledIds = runPayload._ddDisabledTestIds
+    this[kDdPlaywrightDisabledTestIds] = disabledIds ? new Set(disabledIds) : undefined
+    return runTestGroup.apply(this, arguments)
+  })
+
   shimmer.wrap(workerMain, '_runTest', _runTest => async function (test) {
+    if (this[kDdPlaywrightDisabledTestIds]?.has(test.id)) {
+      test._ddIsDisabled = true
+      test.expectedStatus = 'skipped'
+    }
     await waitForEfdRetryCount(test)
     if (shouldSkipEfdRetry(test)) {
       test._ddShouldSkipEfdRetry = true
@@ -2042,9 +2079,9 @@ function generateSummaryWrapper (generateSummary) {
   return function (...args) {
     for (const test of this.suite.allTests()) {
       // https://github.com/microsoft/playwright/blob/bf92ffecff6f30a292b53430dbaee0207e0c61ad/packages/playwright/src/reporters/base.ts#L279
-      const didNotRun = test.outcome() === 'skipped' &&
-        (!test.results.length || test.expectedStatus !== 'skipped')
-      if (didNotRun && !testsReportedInGenerateSummary.has(test)) {
+      const shouldReportAsSkipped = test.outcome() === 'skipped' &&
+        (disabledTestIds.has(test.id) || !test.results.length || test.expectedStatus !== 'skipped')
+      if (shouldReportAsSkipped && !testsReportedInGenerateSummary.has(test)) {
         testsReportedInGenerateSummary.add(test)
         const {
           _requireFile: testSuiteAbsolutePath,