fix(llmobs): address OpenAI Agents review feedback (#9526)

* fix(llmobs): address OpenAI Agents review feedback

* fix(llmobs): complete chat completions metadata

* fix(llmobs): avoid replacing agents model class

* fix(llmobs): address agents review gaps

Author: crysmags

packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/index.js

@@ -10,6 +10,7 @@ module.exports = [
   ...require('./langgraph'),
   ...require('./mercurius'),
   ...require('./modelcontextprotocol-sdk'),
+  ...require('./openai-agents'),
   ...require('./playwright'),
   ...require('./aws-durable-execution-sdk-js'),
 ]

packages/datadog-instrumentations/src/helpers/rewriter/instrumentations/openai-agents.js

@@ -0,0 +1,31 @@
+'use strict'
+
+const moduleName = '@openai/agents-openai'
+const versionRange = '>=0.7.0'
+const functionQuery = {
+  methodName: 'constructor',
+  className: 'OpenAIChatCompletionsModel',
+  kind: 'Sync',
+}
+const channelName = 'OpenAIChatCompletionsModel_constructor'
+
+module.exports = [
+  {
+    module: {
+      name: moduleName,
+      versionRange,
+      filePath: 'dist/openaiChatCompletionsModel.js',
+    },
+    functionQuery,
+    channelName,
+  },
+  {
+    module: {
+      name: moduleName,
+      versionRange,
+      filePath: 'dist/openaiChatCompletionsModel.mjs',
+    },
+    functionQuery,
+    channelName,
+  },
+]

packages/datadog-instrumentations/src/openai-agents.js

@@ -1,30 +1,49 @@
 'use strict'
 
-const { channel } = require('dc-polyfill')
+const { channel, tracingChannel } = require('dc-polyfill')
 const shimmer = require('../../datadog-shimmer')
-const { addHook } = require('./helpers/instrument')
+const { addHook, getHooks } = require('./helpers/instrument')
 
 // `WeakSet` keyed by module exports — replaces the underscored
 // `mod._datadogPatched` flag while keeping dedupe semantics. Mods are kept
 // alive by `require.cache` anyway, so this doesn't add lifetime to anything.
 const patchedMods = new WeakSet()
+const modelBaseURLs = new WeakMap()
+const chatCompletionsModelConstructorCh =
+  tracingChannel('orchestrion:@openai/agents-openai:OpenAIChatCompletionsModel_constructor')
+
+/**
+ * Capture the client base URL before the SDK stores the client in a #private field.
+ *
+ * @param {{ arguments?: [{ baseURL?: string }], self?: object }} ctx
+ */
+function captureChatCompletionsModelBaseURL (ctx) {
+  const baseURL = ctx.arguments?.[0]?.baseURL
+  if (ctx.self && typeof baseURL === 'string') modelBaseURLs.set(ctx.self, baseURL)
+}
+
+chatCompletionsModelConstructorCh.end.subscribe(captureChatCompletionsModelBaseURL)
+
+for (const hook of getHooks('@openai/agents-openai')) {
+  addHook(hook, moduleExports => moduleExports)
+}
 
 // Plugin subscribes to this and registers its TracingProcessor when
 // `@openai/agents` loads. Publishing from here keeps this file free of
 // any cross-package import from the plugin.
 const agentsCoreLoadedCh = channel('apm:openai-agents:agents-core:loaded')
 
-// Plugin subscribes here to keep track of the OpenAI-compatible client's
-// baseURL — used to resolve `model_provider` (openai / azure_openai /
-// deepseek / unknown).
-const responseClientCh = channel('apm:openai-agents:response:client')
-
 // Plugin uses addBind on this channel so that legacyStorage.run(store, fn) wraps
 // the model call — including async iterator advancement for streaming responses.
 // This ensures the active dd-trace span is visible to the openai plugin when it
 // creates its openai.request span, correctly parenting it under the agent span.
 const modelStartCh = channel('apm:openai-agents:model:start')
 
+// Tool.invoke runs inside agents-core's public function-span context. Bind the
+// corresponding dd-trace tool span around that public invocation boundary so
+// spans created by user tool code inherit it.
+const toolStartCh = channel('apm:openai-agents:tool:start')
+
 // Reference to the loaded @openai/agents module, captured in the first hook
 // so that wrapResponseMethod can call getCurrentSpan() without an additional
 // require (and without triggering n/no-missing-require on agents-core internals).
@@ -36,60 +55,84 @@ let agentsMod
 //   provider.getCurrentSpan()              (replaces getCurrentSpan)
 // Both APIs are tried so this file works across the full supported version range.
 // The plugin subscriber (index.js) handles processor registration via the channel.
-function getCurrentSpanId (mod) {
-  if (typeof mod?.getCurrentSpan === 'function') {
-    return mod.getCurrentSpan()?.spanId
+function getCurrentSpan () {
+  if (typeof agentsMod?.getCurrentSpan === 'function') {
+    return agentsMod.getCurrentSpan()
   }
-  if (typeof mod?.getGlobalTraceProvider === 'function') {
-    return mod.getGlobalTraceProvider().getCurrentSpan()?.spanId
+  if (typeof agentsMod?.getGlobalTraceProvider === 'function') {
+    return agentsMod.getGlobalTraceProvider().getCurrentSpan()
   }
 }
 
+function getCurrentSpanId () {
+  return getCurrentSpan()?.spanId
+}
+
 addHook({ name: '@openai/agents', versions: ['>=0.7.0'] }, (mod) => {
   if (patchedMods.has(mod)) return mod
   if (typeof mod?.addTraceProcessor !== 'function' && typeof mod?.getGlobalTraceProvider !== 'function') return mod
   patchedMods.add(mod)
   agentsMod = mod
+  if (typeof mod.tool === 'function') {
+    shimmer.wrap(mod, 'tool', wrapToolFactory, { replaceGetter: true })
+  }
   agentsCoreLoadedCh.publish({ mod })
   return mod
 })
 
+function wrapToolFactory (original) {
+  return function (...args) {
+    const tool = original.apply(this, args)
+    if (typeof tool?.invoke === 'function') {
+      shimmer.wrap(tool, 'invoke', wrapToolInvoke)
+    }
+    return tool
+  }
+}
+
+function wrapToolInvoke (original) {
+  return function (...args) {
+    const agentsCoreSpan = getCurrentSpan()
+    return toolStartCh.runStores({ agentsCoreSpan }, () => original.apply(this, args))
+  }
+}
+
 function wrapResponseMethod (original) {
   return function (...args) {
-    const agentsCoreSpanId = getCurrentSpanId(agentsMod)
-    publishClientBaseURL(this)
-    return modelStartCh.runStores({ agentsCoreSpanId }, () => original.apply(this, args))
+    const agentsCoreSpanId = getCurrentSpanId()
+    const baseURL = getClientBaseURL(this)
+    return modelStartCh.runStores({ agentsCoreSpanId, baseURL }, () => original.apply(this, args))
   }
 }
 
 function wrapStreamedResponseMethod (original) {
   return function (...args) {
-    const agentsCoreSpanId = getCurrentSpanId(agentsMod)
-    publishClientBaseURL(this)
-    const iterator = modelStartCh.runStores({ agentsCoreSpanId }, () => original.apply(this, args))
-    return wrapAsyncIterator(iterator, agentsCoreSpanId)
+    const agentsCoreSpanId = getCurrentSpanId()
+    const baseURL = getClientBaseURL(this)
+    const context = { agentsCoreSpanId, baseURL }
+    const iterator = modelStartCh.runStores(context, () => original.apply(this, args))
+    return wrapAsyncIterator(iterator, context)
   }
 }
 
-function publishClientBaseURL (model) {
-  const baseURL = model?.client?.baseURL ?? model?._client?.baseURL
-  if (baseURL) responseClientCh.publish({ baseURL })
+function getClientBaseURL (model) {
+  return model?.client?.baseURL ?? model?._client?.baseURL ?? modelBaseURLs.get(model)
 }
 
-function wrapAsyncIterator (iterator, agentsCoreSpanId) {
+function wrapAsyncIterator (iterator, context) {
   if (!iterator || typeof iterator !== 'object') return iterator
 
   return {
     next () {
-      return modelStartCh.runStores({ agentsCoreSpanId }, () => iterator.next.apply(iterator, arguments))
+      return modelStartCh.runStores(context, () => iterator.next.apply(iterator, arguments))
     },
     throw () {
       if (typeof iterator.throw !== 'function') return Promise.reject(arguments[0])
-      return modelStartCh.runStores({ agentsCoreSpanId }, () => iterator.throw.apply(iterator, arguments))
+      return modelStartCh.runStores(context, () => iterator.throw.apply(iterator, arguments))
     },
     return () {
       if (typeof iterator.return !== 'function') return Promise.resolve({ done: true, value: arguments[0] })
-      return modelStartCh.runStores({ agentsCoreSpanId }, () => iterator.return.apply(iterator, arguments))
+      return modelStartCh.runStores(context, () => iterator.return.apply(iterator, arguments))
     },
     [Symbol.asyncIterator] () {
       return this
@@ -99,11 +142,18 @@ function wrapAsyncIterator (iterator, agentsCoreSpanId) {
 
 addHook({ name: '@openai/agents-openai', versions: ['>=0.7.0'] }, (mod) => {
   if (patchedMods.has(mod)) return mod
-  const proto = mod?.OpenAIResponsesModel?.prototype
-  if (!proto) return mod
+  const responseProto = mod?.OpenAIResponsesModel?.prototype
+  const chatCompletionsProto = mod?.OpenAIChatCompletionsModel?.prototype
+  if (!responseProto && !chatCompletionsProto) return mod
 
   patchedMods.add(mod)
-  shimmer.wrap(proto, 'getResponse', wrapResponseMethod)
-  shimmer.wrap(proto, 'getStreamedResponse', wrapStreamedResponseMethod)
+  for (const proto of [responseProto, chatCompletionsProto]) {
+    if (typeof proto?.getResponse === 'function') {
+      shimmer.wrap(proto, 'getResponse', wrapResponseMethod)
+    }
+    if (typeof proto?.getStreamedResponse === 'function') {
+      shimmer.wrap(proto, 'getStreamedResponse', wrapStreamedResponseMethod)
+    }
+  }
   return mod
 })

packages/datadog-plugin-openai-agents/src/index.js

@@ -2,7 +2,7 @@
 
 const { storage } = require('../../datadog-core')
 const Plugin = require('../../dd-trace/src/plugins/plugin')
-const { OpenAIAgentsIntegration } = require('./integration')
+const { MODEL_BASE_URL_STORE_KEY, OpenAIAgentsIntegration } = require('./integration')
 const { DDOpenAIAgentsProcessor } = require('./processor')
 
 const legacyStorage = storage('legacy')
@@ -14,9 +14,6 @@ const legacyStorage = storage('legacy')
  * during its constructor (which runs synchronously between `loadChannel`'s
  * publish and the addHook callback) and registers the processor.
  *
- * The instrumentation also publishes the OpenAI-compatible client's baseURL on
- * each model response call, so the integration can resolve `model_provider`.
- *
  * The integration's `enabled` flag follows this plugin's configure()
  * lifecycle. Each loaded version of the agents package replaces all processors
  * via setTraceProcessors() on module load, so the plugin re-registers a
@@ -48,26 +45,29 @@ class OpenaiAgentsPlugin extends Plugin {
       }
     })
 
-    this.addSub('apm:openai-agents:response:client', ({ baseURL }) => {
-      if (!this.#integration.enabled) return
-      this.#integration.setClientBaseURL(baseURL)
-    })
-
     // Activate the current agent's dd-trace span in legacyStorage for the
     // duration of model response calls and stream iterator advancement. This
     // makes the openai plugin's shimmer see the correct parent when it creates its
     // openai.request span, so all spans land in the same trace.
-    this.addBind('apm:openai-agents:model:start', ({ agentsCoreSpanId }) => {
-      if (!this.#integration.enabled || !agentsCoreSpanId) return legacyStorage.getStore()
+    this.addBind('apm:openai-agents:model:start', ({ agentsCoreSpanId, baseURL }) => {
+      const store = legacyStorage.getStore()
+      if (!this.#integration.enabled || !agentsCoreSpanId) return store
       const ddSpan = this.#integration.getDDSpan(agentsCoreSpanId)
-      if (!ddSpan) return legacyStorage.getStore()
-      return { ...legacyStorage.getStore(), span: ddSpan }
+      if (!ddSpan) return store
+      return { ...store, [MODEL_BASE_URL_STORE_KEY]: baseURL, span: ddSpan }
+    })
+
+    this.addBind('apm:openai-agents:tool:start', ({ agentsCoreSpan }) => {
+      const store = legacyStorage.getStore()
+      if (!this.#integration.enabled || !agentsCoreSpan) return store
+      const ddSpan = this.#integration.getOrStartToolSpan(agentsCoreSpan)
+      return ddSpan ? { ...store, span: ddSpan } : store
     })
   }
 
   configure (config) {
     super.configure(config)
-    this.#integration.setEnabled(!!config?.enabled)
+    this.#integration.configure(config)
   }
 }