feat(aiguard): evaluating anthropic calls with AI guard automatically (#9219)

* feat(aiguard): evaluating anthropic calls with AI guard automatically

Author: IlyasShabi

.github/workflows/instrumentation.yml

@@ -61,6 +61,16 @@ jobs:
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
       - uses: ./.github/actions/instrumentations/test
 
+  instrumentation-anthropic-lifecycle:
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    env:
+      PLUGINS: anthropic-lifecycle
+    steps:
+      - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+      - uses: ./.github/actions/instrumentations/test
+
   instrumentation-aws-sdk:
     runs-on: ubuntu-latest
     permissions:

packages/datadog-instrumentations/src/anthropic.js

@@ -6,6 +6,59 @@ const { addHook } = require('./helpers/instrument')
 
 const anthropicTracingChannel = tracingChannel('apm:anthropic:request')
 const onStreamedChunkCh = channel('apm:anthropic:request:chunk')
+const messagesBeforeChannel = channel('dd-trace:anthropic:messages:before')
+const messagesAfterChannel = channel('dd-trace:anthropic:messages:after')
+
+/**
+ * Publishes a provider-native lifecycle payload to a cancelable lifecycle channel.
+ *
+ * Subscribers push async work into `pending` synchronously during publication and
+ * abort `abortController` with an error before the pushed promise resolves to block.
+ *
+ * @param {object} channel
+ * @param {object} payload
+ * @returns {Promise<void>}
+ */
+function publishLifecycle (channel, payload) {
+  const abortController = new AbortController()
+  const ctx = { ...payload, abortController, pending: [] }
+
+  channel.publish(ctx)
+
+  return Promise.all(ctx.pending).then(() => {
+    if (abortController.signal.aborted) {
+      throw abortController.signal.reason
+    }
+  })
+}
+
+/**
+ * @template T
+ * @param {Promise<T>} promise
+ * @param {Promise<void>|undefined} verdict
+ * @returns {Promise<T>}
+ */
+function waitForVerdict (promise, verdict) {
+  return verdict
+    ? Promise.all([verdict, promise]).then(([, value]) => value)
+    : promise
+}
+
+/**
+ * @param {object} response
+ * @param {'json'|'text'} method
+ * @param {(body: object|string) => Promise<void>|undefined} getVerdict
+ */
+function wrapResponseReader (response, method, getVerdict) {
+  if (typeof response[method] !== 'function') return
+
+  shimmer.wrap(response, method, original => function (...args) {
+    return original.apply(this, args).then(body => {
+      const verdict = getVerdict(body)
+      return verdict ? verdict.then(() => body) : body
+    })
+  })
+}
 
 function wrapStreamIterator (iterator, ctx) {
   return function (...args) {
@@ -34,16 +87,20 @@ function wrapStreamIterator (iterator, ctx) {
 
 function wrapCreate (create) {
   return function (...args) {
-    if (!anthropicTracingChannel.start.hasSubscribers) {
+    const options = args[0]
+    const stream = options?.stream
+
+    const hasLifecycle = !stream && (messagesBeforeChannel.hasSubscribers || messagesAfterChannel.hasSubscribers)
+
+    if (!anthropicTracingChannel.start.hasSubscribers && !hasLifecycle) {
       return create.apply(this, args)
     }
 
-    const options = args[0]
-    const stream = options.stream
-
     const ctx = { options, resource: 'create', baseUrl: this._client?.baseURL }
 
     return anthropicTracingChannel.start.runStores(ctx, () => {
+      const parentSpan = hasLifecycle ? ctx.currentStore?.span : undefined
+
       let apiPromise
       try {
         apiPromise = create.apply(this, args)
@@ -52,18 +109,76 @@ function wrapCreate (create) {
         throw error
       }
 
-      shimmer.wrap(apiPromise, 'parse', parse => function (...args) {
-        return parse.apply(this, args)
+      let afterVerdict
+      let parseResult
+      let wrappedResponse
+
+      let beforeVerdict
+      function getBeforeVerdict () {
+        if (!hasLifecycle || beforeVerdict) return beforeVerdict
+        if (!messagesBeforeChannel.hasSubscribers) return
+
+        beforeVerdict = publishLifecycle(messagesBeforeChannel, { args, parentSpan })
+        return beforeVerdict
+      }
+
+      /**
+       * @param {object|string} body
+       */
+      function getAfterVerdict (body) {
+        if (!hasLifecycle || afterVerdict) return afterVerdict
+        if (!messagesAfterChannel.hasSubscribers) return
+
+        afterVerdict = publishLifecycle(messagesAfterChannel, { args, body, parentSpan })
+        return afterVerdict
+      }
+
+      shimmer.wrap(apiPromise, 'parse', parse => function (...parseArgs) {
+        if (parseResult) return parseResult
+
+        const parsed = parse.apply(this, parseArgs)
+        parseResult = waitForVerdict(parsed, getBeforeVerdict())
           .then(response => {
             if (stream) {
               shimmer.wrap(response, Symbol.asyncIterator, iterator => wrapStreamIterator(iterator, ctx))
-            } else {
+              return response
+            }
+            const verdict = getAfterVerdict(response)
+            if (!verdict) {
+              finish(ctx, response, null)
+              return response
+            }
+            // Finish after evaluation so a block propagates the error to anthropic.request
+            // and the span wraps its child instead of closing before it.
+            return verdict.then(() => {
               finish(ctx, response, null)
+              return response
+            })
+          }).catch(error => {
+            if (!ctx.finished) finish(ctx, null, error)
+            throw error
+          })
+
+        return parseResult
+      })
+
+      // Gate `.asResponse()` callers on the before verdict so raw-response paths still block,
+      // then evaluate output only if the caller consumes the JSON body.
+      shimmer.wrap(apiPromise, 'asResponse', origAsResponse => function (...asResponseArgs) {
+        return waitForVerdict(origAsResponse.apply(this, asResponseArgs), getBeforeVerdict())
+          .then(response => {
+            if (!stream && hasLifecycle && wrappedResponse !== response) {
+              wrappedResponse = response
+              wrapResponseReader(response, 'json', getAfterVerdict)
+              wrapResponseReader(response, 'text', getAfterVerdict)
             }
 
+            if (afterVerdict) return afterVerdict.then(() => response)
+            if (!stream && !ctx.finished && !parseResult) finish(ctx, null, null)
             return response
-          }).catch(error => {
-            finish(ctx, null, error)
+          })
+          .catch(error => {
+            if (!ctx.finished) finish(ctx, null, error)
             throw error
           })
       })
@@ -76,13 +191,16 @@ function wrapCreate (create) {
 }
 
 function finish (ctx, result, error) {
+  if (ctx.finished) return
+
   if (error) {
     ctx.error = error
     anthropicTracingChannel.error.publish(ctx)
   }
 
   // streamed responses are handled and set separately
   ctx.result ??= result
+  ctx.finished = true
 
   anthropicTracingChannel.asyncEnd.publish(ctx)
 }