dd-trace 5.110.0 → 5.112.0

package/packages/datadog-plugin-vitest/src/index.js

@@ -30,12 +30,12 @@ const {
   TEST_MANAGEMENT_IS_ATTEMPT_TO_FIX,
   TEST_MANAGEMENT_ATTEMPT_TO_FIX_PASSED,
   TEST_HAS_FAILED_ALL_RETRIES,
-  getLibraryCapabilitiesTags,
+  getLibraryCapabilitiesTags: getDefaultLibraryCapabilitiesTags,
   TEST_RETRY_REASON_TYPES,
-  isModifiedTest,
   TEST_IS_MODIFIED,
   TEST_HAS_DYNAMIC_NAME,
   TEST_FINAL_STATUS,
+  TEST_IS_TEST_FRAMEWORK_WORKER,
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 const {
@@ -67,52 +67,10 @@ class VitestPlugin extends CiPlugin {
         testCommand: this.command,
         repositoryRoot: this.repositoryRoot,
         codeOwnersEntries: this.codeOwnersEntries,
+        testEnvironmentMetadata: this.testEnvironmentMetadata,
       })
     })
 
-    this.addSub('ci:vitest:test:is-new', ({ knownTests, testSuiteAbsolutePath, testName, onDone }) => {
-      // if for whatever reason the worker does not receive valid known tests, we don't consider it as new
-      if (!knownTests.vitest) {
-        return onDone(false)
-      }
-      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
-      const testsForThisTestSuite = knownTests.vitest[testSuite] || []
-      onDone(!testsForThisTestSuite.includes(testName))
-    })
-
-    this.addSub('ci:vitest:test:is-attempt-to-fix', ({
-      testManagementTests,
-      testSuiteAbsolutePath,
-      testName,
-      onDone,
-    }) => {
-      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
-      const { isAttemptToFix } = this.getTestProperties(testManagementTests, testSuite, testName)
-
-      onDone(isAttemptToFix ?? false)
-    })
-
-    this.addSub('ci:vitest:test:is-disabled', ({ testManagementTests, testSuiteAbsolutePath, testName, onDone }) => {
-      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
-      const { isDisabled } = this.getTestProperties(testManagementTests, testSuite, testName)
-
-      onDone(isDisabled)
-    })
-
-    this.addSub('ci:vitest:test:is-quarantined', ({ testManagementTests, testSuiteAbsolutePath, testName, onDone }) => {
-      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
-      const { isQuarantined } = this.getTestProperties(testManagementTests, testSuite, testName)
-
-      onDone(isQuarantined)
-    })
-
-    this.addSub('ci:vitest:test:is-modified', ({ modifiedFiles, testSuiteAbsolutePath, onDone }) => {
-      const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
-      const isModified = isModifiedTest(testSuite, 0, 0, modifiedFiles, this.constructor.id)
-
-      onDone(isModified)
-    })
-
     this.addSub('ci:vitest:is-early-flake-detection-faulty', ({
       knownTests,
       testFilepaths,
@@ -141,12 +99,16 @@ class VitestPlugin extends CiPlugin {
         isRetryReasonAttemptToFix,
         isRetryReasonAtr,
         isModified,
+        isTestFrameworkWorker,
+        requestErrorTags,
       } = ctx
 
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
-      const store = storage('legacy').getStore()
+      const store = ctx.currentStore || storage('legacy').getStore()
+      const testSuiteSpan = store?.testSuiteSpan || this.testSuiteSpan
 
       const extraTags = {
+        ...requestErrorTags,
         [TEST_SOURCE_FILE]: testSuite,
       }
       if (isRetry) {
@@ -179,11 +141,14 @@ class VitestPlugin extends CiPlugin {
       if (isModified) {
         extraTags[TEST_IS_MODIFIED] = 'true'
       }
+      if (isTestFrameworkWorker) {
+        extraTags[TEST_IS_TEST_FRAMEWORK_WORKER] = 'true'
+      }
 
       const span = this.startTestSpan(
         testName,
         testSuite,
-        this.testSuiteSpan,
+        testSuiteSpan,
         extraTags
       )
 
@@ -204,7 +169,7 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addBind('ci:vitest:test:finish-time', (ctx) => {
-      const { status, task, attemptToFixPassed, attemptToFixFailed } = ctx
+      const { status, task, attemptToFixPassed, attemptToFixFailed, duration } = ctx
       const span = ctx.currentStore?.span
 
       // we store the finish time to finish at a later hook
@@ -218,7 +183,8 @@ class VitestPlugin extends CiPlugin {
           span.setTag(TEST_MANAGEMENT_ATTEMPT_TO_FIX_PASSED, 'false')
         }
 
-        this.taskToFinishTime.set(task, span._getTime())
+        const finishTime = typeof duration === 'number' ? span._startTime + duration : span._getTime()
+        this.taskToFinishTime.set(task, finishTime)
 
         ctx.parentStore = ctx.currentStore
         ctx.currentStore = { ...ctx.currentStore, span }
@@ -227,7 +193,7 @@ class VitestPlugin extends CiPlugin {
       return ctx.currentStore
     })
 
-    this.addSub('ci:vitest:test:pass', ({ span, task, finalStatus, earlyFlakeAbortReason }) => {
+    this.addSub('ci:vitest:test:pass', ({ span, task, finalStatus, earlyFlakeAbortReason, promises }) => {
       if (span) {
         this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', this.getTestTelemetryTags(span))
         span.setTag(TEST_STATUS, 'pass')
@@ -237,8 +203,21 @@ class VitestPlugin extends CiPlugin {
         if (earlyFlakeAbortReason) {
           span.setTag(TEST_EARLY_FLAKE_ABORT_REASON, earlyFlakeAbortReason)
         }
-        span.finish(this.taskToFinishTime.get(task))
-        finishAllTraceSpans(span)
+        const finish = () => {
+          span.finish(this.taskToFinishTime.get(task))
+          finishAllTraceSpans(span)
+        }
+
+        if (finalStatus) {
+          if (promises && this.diBreakpointHitPromise) {
+            promises.hitBreakpointPromise = this.waitForPreparedDiBreakpointHit().then(finish)
+            return
+          }
+          finish()
+          this.cancelDiBreakpointHitWait()
+          return
+        }
+        finish()
       }
     })
 
@@ -247,6 +226,7 @@ class VitestPlugin extends CiPlugin {
       duration,
       error,
       shouldSetProbe,
+      shouldWaitForHitProbe,
       promises,
       hasFailedAllRetries,
       attemptToFixFailed,
@@ -262,7 +242,8 @@ class VitestPlugin extends CiPlugin {
           const { file, line, stackIndex, setProbePromise } = probeInformation
           this.runningTestProbe = { file, line }
           this.testErrorStackIndex = stackIndex
-          promises.setProbePromise = setProbePromise
+          this.prepareDiBreakpointHitWait()
+          promises.setProbePromise = this.waitForDiOperation(setProbePromise)
         }
       }
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', this.getTestTelemetryTags(span))
@@ -283,27 +264,56 @@ class VitestPlugin extends CiPlugin {
       if (earlyFlakeAbortReason) {
         span.setTag(TEST_EARLY_FLAKE_ABORT_REASON, earlyFlakeAbortReason)
       }
-      if (duration) {
-        span.finish(span._startTime + duration - MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION) // milliseconds
-      } else {
-        span.finish() // `duration` is empty for retries, so we'll use clock time
+      const finish = () => {
+        if (duration) {
+          span.finish(span._startTime + duration - MILLISECONDS_TO_SUBTRACT_FROM_FAILED_TEST_DURATION) // milliseconds
+        } else {
+          span.finish() // `duration` is empty for retries, so we'll use clock time
+        }
+        finishAllTraceSpans(span)
+      }
+
+      if (!shouldSetProbe && finalStatus && promises && this.diBreakpointHitPromise) {
+        promises.hitBreakpointPromise = this.waitForPreparedDiBreakpointHit().then(finish)
+        return
+      }
+      finish()
+      if (shouldWaitForHitProbe) {
+        this.prepareDiBreakpointHitWait()
+      } else if (!shouldSetProbe) {
+        this.cancelDiBreakpointHitWait()
+      }
+    })
+
+    this.addSub('ci:vitest:test:di:wait', ({ promises }) => {
+      if (this.di) {
+        promises.hitBreakpointPromise = this.waitForDiBreakpointHits()
       }
-      finishAllTraceSpans(span)
     })
 
-    this.addSub('ci:vitest:test:skip', ({ testName, testSuiteAbsolutePath, isNew, isDisabled }) => {
+    this.addSub('ci:vitest:test:skip', ({
+      testName,
+      testSuiteAbsolutePath,
+      isNew,
+      isDisabled,
+      isTestFrameworkWorker,
+      requestErrorTags,
+      testSuiteSpan,
+    }) => {
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
       const testSpan = this.startTestSpan(
         testName,
         testSuite,
-        this.testSuiteSpan,
+        testSuiteSpan || this.testSuiteSpan,
         {
+          ...requestErrorTags,
           [TEST_SOURCE_FILE]: testSuite,
           [TEST_SOURCE_START]: 1, // we can't get the proper start line in vitest
           [TEST_STATUS]: 'skip',
           [TEST_FINAL_STATUS]: 'skip',
           ...(isDisabled ? { [TEST_MANAGEMENT_IS_DISABLED]: 'true' } : {}),
           ...(isNew ? { [TEST_IS_NEW]: 'true' } : {}),
+          ...(isTestFrameworkWorker ? { [TEST_IS_TEST_FRAMEWORK_WORKER]: 'true' } : {}),
         }
       )
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'test', this.getTestTelemetryTags(testSpan))
@@ -311,10 +321,25 @@ class VitestPlugin extends CiPlugin {
     })
 
     this.addBind('ci:vitest:test-suite:start', (ctx) => {
-      const { codeOwnersEntries, repositoryRoot, testSuiteAbsolutePath, frameworkVersion } = ctx
+      const {
+        codeOwnersEntries,
+        repositoryRoot,
+        testEnvironmentMetadata,
+        requestErrorTags,
+        testSuiteAbsolutePath,
+        frameworkVersion,
+        isTestFrameworkWorker,
+        isVitestNoWorkerInitActive,
+      } = ctx
 
       const testCommand = ctx.testCommand || 'vitest run'
       const { testSessionId, testModuleId } = ctx
+      if (testEnvironmentMetadata) {
+        this.testEnvironmentMetadata = {
+          ...this.testEnvironmentMetadata,
+          ...testEnvironmentMetadata,
+        }
+      }
       this._setRepositoryRoot(repositoryRoot, codeOwnersEntries)
       this.command = testCommand
       this.frameworkVersion = frameworkVersion
@@ -329,13 +354,16 @@ class VitestPlugin extends CiPlugin {
       // test suites run in a different process, so they also need to init the metadata dictionary
       const testSessionName = getTestSessionName(this.config, trimmedCommand, this.testEnvironmentMetadata)
       if (this.tracer._exporter.addMetadataTags) {
+        const libraryCapabilitiesTags = this.getLibraryCapabilitiesTags(frameworkVersion, {
+          isVitestNoWorkerInitActive,
+        })
         this.tracer._exporter.addMetadataTags({
           [TEST_LEVELS_METADATA]: {
             [TEST_COMMAND]: testCommand,
             [TEST_SESSION_NAME]: testSessionName,
             ...getTestLevelsMetadataTags(this.testEnvironmentMetadata),
           },
-          test: getLibraryCapabilitiesTags(this.constructor.id),
+          test: libraryCapabilitiesTags,
         })
       }
 
@@ -348,9 +376,13 @@ class VitestPlugin extends CiPlugin {
           testSuite,
           'vitest'
         ),
+        ...requestErrorTags,
         [TEST_SOURCE_FILE]: testSuite,
         [TEST_SOURCE_START]: 1,
       }
+      if (isTestFrameworkWorker) {
+        testSuiteMetadata[TEST_IS_TEST_FRAMEWORK_WORKER] = 'true'
+      }
 
       const codeOwners = this.getCodeOwners(testSuiteMetadata)
       if (codeOwners) {
@@ -374,13 +406,17 @@ class VitestPlugin extends CiPlugin {
       return ctx.currentStore
     })
 
-    this.addSub('ci:vitest:test-suite:finish', ({ testSuiteSpan, status, onFinish }) => {
+    this.addSub('ci:vitest:test-suite:finish', ({ testSuiteSpan, status, deferFlush, onFinish }) => {
       if (testSuiteSpan) {
         testSuiteSpan.setTag(TEST_STATUS, status)
         testSuiteSpan.finish()
         finishAllTraceSpans(testSuiteSpan)
       }
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'suite')
+      if (deferFlush) {
+        onFinish()
+        return
+      }
       this.tracer._exporter.flush(onFinish)
       if (this.runningTestProbe) {
         this.removeDiProbe(this.runningTestProbe)
@@ -409,9 +445,15 @@ class VitestPlugin extends CiPlugin {
       isEarlyFlakeDetectionEnabled,
       isEarlyFlakeDetectionFaulty,
       isTestManagementTestsEnabled,
+      requestErrorTags,
       vitestPool,
+      isVitestNoWorkerInitActive,
       onFinish,
     }) => {
+      for (const [tag, value] of Object.entries(requestErrorTags || {})) {
+        this.testSessionSpan.setTag(tag, value)
+        this.testModuleSpan.setTag(tag, value)
+      }
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)
       if (error) {
@@ -438,12 +480,12 @@ class VitestPlugin extends CiPlugin {
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'module')
       this.testSessionSpan.finish()
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'session', {
-        hasFailedTestReplay: this.libraryConfig?.isDiEnabled || undefined,
+        hasFailedTestReplay: this.libraryConfig?.isDiEnabled && !isVitestNoWorkerInitActive ? true : undefined,
       })
       finishAllTraceSpans(this.testSessionSpan)
       this.telemetry.count(TELEMETRY_TEST_SESSION, {
         provider: this.ciProviderName,
-        autoInjected: !!this._tracerConfig.DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER,
+        autoInjected: !!this._tracerConfig.testOptimization.DD_CIVISIBILITY_AUTO_INSTRUMENTATION_PROVIDER,
       })
       this.tracer._exporter.flush(onFinish)
     })
@@ -453,6 +495,19 @@ class VitestPlugin extends CiPlugin {
     })
   }
 
+  /**
+   * Returns Vitest library capability metadata tags.
+   * @param {string} frameworkVersion - The Vitest version.
+   * @param {object} [ctx] - Diagnostic channel context.
+   * @param {boolean} [ctx.isVitestNoWorkerInitActive] - Whether no-worker init is active for this run.
+   * @returns {Record<string, string|undefined>}
+   */
+  getLibraryCapabilitiesTags (frameworkVersion, ctx = {}) {
+    return getDefaultLibraryCapabilitiesTags(this.constructor.id, frameworkVersion, {
+      omitFailedTestReplay: ctx.isVitestNoWorkerInitActive,
+    })
+  }
+
   /**
    * Handles the coverage report by discovering and uploading it if enabled.
    * @param {string} rootDir - The root directory where coverage reports are located.
@@ -469,13 +524,6 @@ class VitestPlugin extends CiPlugin {
       onDone,
     })
   }
-
-  getTestProperties (testManagementTests, testSuite, testName) {
-    const { attempt_to_fix: isAttemptToFix, disabled: isDisabled, quarantined: isQuarantined } =
-      testManagementTests?.vitest?.suites?.[testSuite]?.tests?.[testName]?.properties || {}
-
-    return { isAttemptToFix, isDisabled, isQuarantined }
-  }
 }
 
 module.exports = VitestPlugin

package/packages/datadog-shimmer/src/shimmer.js

@@ -119,11 +119,14 @@ function wrapCallback (original, wrapper) {
  * object.
  * @param {string | symbol} name - The property key of the method to wrap.
  * @param {(original: Function) => (...args: unknown[]) => unknown} wrapper - The wrapper function.
- * @param {{ replaceGetter?: boolean }} [options] - If `replaceGetter` is set to
- * true, the getter is accessed and the getter is replaced with one that just
- * returns the earlier retrieved value. Use with care! This may only be done in
- * case the getter absolutely has no side effect and no setter is defined for the
- * property.
+ * @param {{ replaceGetter?: boolean }} [options] - By default the getter is
+ * wrapped in place, so each property access runs the wrapper. A getter+setter
+ * pair keeps its setter; a setter-only property throws. If `replaceGetter` is
+ * true, the getter is instead accessed once and replaced with one returning the
+ * resolved wrapped value — for a lazy getter+setter pair (e.g. Node 20's
+ * `fs.opendir`) the setter is rebuilt to materialize a writable data property on
+ * assignment, keeping the descriptor observationally identical for downstream
+ * consumers. Use with care! This may only be done when the getter has no side effect.
  * @returns {Record<string | symbol, unknown> | Function | undefined} The target object with
  * the wrapped method.
  */
@@ -151,17 +154,13 @@ function wrap (target, name, wrapper, options) {
     enumerable: false,
   }
 
-  if (descriptor.set && (!descriptor.get || options?.replaceGetter)) {
-    // It is possible to support these cases by instrumenting both the getter
-    // and setter (or only the setter, in case that is a use case).
-    // For now, this is not supported due to the complexity and the fact that
-    // this is not a common use case.
-    throw new Error(options?.replaceGetter
-      ? 'Replacing a getter/setter pair is not supported. Implement if required.'
-      : 'Replacing setters is not supported. Implement if required.')
+  // A setter-only property has nothing to wrap. Instrumenting the setter is not
+  // implemented; a getter+setter pair is handled below.
+  if (descriptor.set && !descriptor.get) {
+    throw new Error('Replacing setters is not supported. Implement if required.')
   }
 
-  const original = descriptor.value ?? options?.replaceGetter ? target[name] : descriptor.get
+  const original = (descriptor.value ?? options?.replaceGetter) ? target[name] : descriptor.get
 
   assertMethod(target, name, original)
 
@@ -176,8 +175,30 @@ function wrap (target, name, wrapper, options) {
     }
     descriptor.value = wrapped
   } else {
-    if (descriptor.get) {
-      // `replaceGetter` may only be used when the getter has no side effect.
+    if (descriptor.set && options?.replaceGetter) {
+      // A lazy accessor pair (e.g. Node 20's `fs.opendir`). `original` already
+      // resolved the value through the getter. Keep the property an accessor pair
+      // so the shape stays observationally identical — a downstream consumer may
+      // read the descriptor or assign to it on a specific Node.js version. The
+      // getter returns the wrapped value; the setter mirrors the native lazy
+      // contract (assignment self-replaces the property with a writable data
+      // property holding the assigned value), so a caller that overwrites the
+      // method gets exactly what they set, unwrapped, as before. The descriptor
+      // is the original accessor descriptor (no `value`/`writable` slots), so
+      // reassigning `get`/`set` keeps it a valid accessor descriptor.
+      descriptor.get = () => wrapped
+      descriptor.set = function (value) {
+        Object.defineProperty(this, name, {
+          configurable: descriptor.configurable,
+          enumerable: descriptor.enumerable,
+          writable: true,
+          value,
+        })
+      }
+    } else if (descriptor.get) {
+      // Wrap the getter in place; for a getter+setter pair the original setter
+      // stays untouched. `replaceGetter` (no side effect on read) instead returns
+      // the value resolved once into `wrapped`.
       descriptor.get = options?.replaceGetter ? () => wrapped : wrapped
     } else {
       descriptor.value = wrapped

package/packages/dd-trace/src/aiguard/index.js

@@ -2,15 +2,11 @@
 
 const log = require('../log')
 const { incomingHttpRequestStart } = require('./channels')
-const openaiIntegration = require('./integrations/openai')
-const vercelAiIntegration = require('./integrations/vercel-ai')
+const integrations = require('./integrations')
 const AIGuard = require('./sdk')
 
 let isEnabled = false
 let aiguard
-let block
-let disableOpenAIIntegration
-let disableVercelAiIntegration
 
 function onIncomingHttpRequestStart () {
   // No-op: subscribing ensures the HTTP plugin spreads req onto the store
@@ -21,31 +17,30 @@ function enable (tracer, config) {
 
   try {
     aiguard = new AIGuard(tracer, config)
-    block = config.experimental?.aiguard?.block !== false
+    const block = config.experimental?.aiguard?.block !== false
 
     incomingHttpRequestStart.subscribe(onIncomingHttpRequestStart)
-    disableOpenAIIntegration = openaiIntegration.enable(aiguard, block)
-    disableVercelAiIntegration = vercelAiIntegration.enable(aiguard, block)
+    integrations.enable(aiguard, block)
 
     isEnabled = true
   } catch (err) {
     log.error('AIGuard: unexpected error during initialization: %s', err.message)
-    disable()
+    reset()
   }
 }
 
 function disable () {
   if (!isEnabled) return
 
+  reset()
+}
+
+function reset () {
   incomingHttpRequestStart.unsubscribe(onIncomingHttpRequestStart)
-  disableOpenAIIntegration?.()
-  disableVercelAiIntegration?.()
+  integrations.disable()
 
   aiguard = undefined
   isEnabled = false
-  block = false
-  disableOpenAIIntegration = undefined
-  disableVercelAiIntegration = undefined
 }
 
 module.exports = { enable, disable }

package/packages/dd-trace/src/aiguard/integrations/index.js

@@ -0,0 +1,34 @@
+'use strict'
+
+const openai = require('./openai')
+const vercelAi = require('./vercel-ai')
+
+let isEnabled = false
+
+/**
+ * Enables all AI Guard provider integrations.
+ *
+ * @param {object} aiguard
+ * @param {boolean} block
+ */
+function enable (aiguard, block) {
+  if (isEnabled) return
+
+  openai.enable(aiguard, block)
+  vercelAi.enable(aiguard, block)
+
+  isEnabled = true
+}
+
+function disable () {
+  vercelAi.disable()
+  openai.disable()
+  isEnabled = false
+}
+
+module.exports = {
+  enable,
+  disable,
+  openai,
+  vercelAi,
+}

package/packages/dd-trace/src/aiguard/integrations/openai.js

@@ -16,51 +16,65 @@ const chatCompletionsAfterChannel = channel('dd-trace:openai:chat.completions:af
 const responsesBeforeChannel = channel('dd-trace:openai:responses:before')
 const responsesAfterChannel = channel('dd-trace:openai:responses:after')
 
+let isEnabled = false
+let aiguard
+let opts
+
 /**
  * Subscribes AI Guard to OpenAI lifecycle channels.
  *
- * @param {object} aiguard
+ * @param {object} aiguardInstance
  * @param {boolean} block
- * @returns {() => void}
  */
-function enable (aiguard, block) {
-  const opts = { block, source: SOURCE_AUTO, integration: 'openai' }
-
-  function onChatCompletionsBefore (ctx) {
-    pushEvaluation(ctx, aiguard, getChatCompletionsInputMessages(ctx.args?.[0]), opts)
-  }
-
-  function onChatCompletionsAfter (ctx) {
-    const inputMessages = getChatCompletionsInputMessages(ctx.args?.[0])
-    if (!inputMessages?.length) return
-    for (const message of getChatCompletionsOutputMessages(ctx.body)) {
-      pushEvaluation(ctx, aiguard, [...inputMessages, message], opts)
-    }
-  }
+function enable (aiguardInstance, block) {
+  if (isEnabled) return
 
-  function onResponsesBefore (ctx) {
-    pushEvaluation(ctx, aiguard, getResponsesInputMessages(ctx.args?.[0]), opts)
-  }
-
-  function onResponsesAfter (ctx) {
-    const inputMessages = getResponsesInputMessages(ctx.args?.[0])
-    if (!inputMessages?.length) return
-    const outputMessages = getResponsesOutputMessages(ctx.body)
-    if (!outputMessages.length) return
-    pushEvaluation(ctx, aiguard, [...inputMessages, ...outputMessages], opts)
-  }
+  aiguard = aiguardInstance
+  opts = { block, source: SOURCE_AUTO, integration: 'openai' }
 
   chatCompletionsBeforeChannel.subscribe(onChatCompletionsBefore)
   chatCompletionsAfterChannel.subscribe(onChatCompletionsAfter)
   responsesBeforeChannel.subscribe(onResponsesBefore)
   responsesAfterChannel.subscribe(onResponsesAfter)
 
-  return function disable () {
-    chatCompletionsBeforeChannel.unsubscribe(onChatCompletionsBefore)
-    chatCompletionsAfterChannel.unsubscribe(onChatCompletionsAfter)
-    responsesBeforeChannel.unsubscribe(onResponsesBefore)
-    responsesAfterChannel.unsubscribe(onResponsesAfter)
+  isEnabled = true
+}
+
+function disable () {
+  if (!isEnabled) return
+
+  chatCompletionsBeforeChannel.unsubscribe(onChatCompletionsBefore)
+  chatCompletionsAfterChannel.unsubscribe(onChatCompletionsAfter)
+  responsesBeforeChannel.unsubscribe(onResponsesBefore)
+  responsesAfterChannel.unsubscribe(onResponsesAfter)
+
+  aiguard = undefined
+  opts = undefined
+  isEnabled = false
+}
+
+function onChatCompletionsBefore (ctx) {
+  pushEvaluation(ctx, aiguard, getChatCompletionsInputMessages(ctx.args?.[0]), opts)
+}
+
+function onChatCompletionsAfter (ctx) {
+  const inputMessages = getChatCompletionsInputMessages(ctx.args?.[0])
+  if (!inputMessages?.length) return
+  for (const message of getChatCompletionsOutputMessages(ctx.body)) {
+    pushEvaluation(ctx, aiguard, [...inputMessages, message], opts)
   }
 }
 
-module.exports = { enable }
+function onResponsesBefore (ctx) {
+  pushEvaluation(ctx, aiguard, getResponsesInputMessages(ctx.args?.[0]), opts)
+}
+
+function onResponsesAfter (ctx) {
+  const inputMessages = getResponsesInputMessages(ctx.args?.[0])
+  if (!inputMessages?.length) return
+  const outputMessages = getResponsesOutputMessages(ctx.body)
+  if (!outputMessages.length) return
+  pushEvaluation(ctx, aiguard, [...inputMessages, ...outputMessages], opts)
+}
+
+module.exports = { enable, disable }