|
| 1 | +'use strict' |
| 2 | + |
| 3 | +const { channel } = require('dc-polyfill') |
| 4 | +const shimmer = require('../../datadog-shimmer') |
| 5 | +const { addHook } = require('./helpers/instrument') |
| 6 | + |
| 7 | +// `WeakSet` keyed by module exports — replaces the underscored |
| 8 | +// `mod._datadogPatched` flag while keeping dedupe semantics. Mods are kept |
| 9 | +// alive by `require.cache` anyway, so this doesn't add lifetime to anything. |
| 10 | +const patchedMods = new WeakSet() |
| 11 | + |
| 12 | +// Plugin subscribes to this and registers its TracingProcessor when |
| 13 | +// `@openai/agents` loads. Publishing from here keeps this file free of |
| 14 | +// any cross-package import from the plugin. |
| 15 | +const agentsCoreLoadedCh = channel('apm:openai-agents:agents-core:loaded') |
| 16 | + |
| 17 | +// Plugin subscribes here to keep track of the OpenAI-compatible client's |
| 18 | +// baseURL — used to resolve `model_provider` (openai / azure_openai / |
| 19 | +// deepseek / unknown). |
| 20 | +const responseClientCh = channel('apm:openai-agents:response:client') |
| 21 | + |
| 22 | +// Plugin uses addBind on this channel so that legacyStorage.run(store, fn) wraps |
| 23 | +// the model call — including async iterator advancement for streaming responses. |
| 24 | +// This ensures the active dd-trace span is visible to the openai plugin when it |
| 25 | +// creates its openai.request span, correctly parenting it under the agent span. |
| 26 | +const modelStartCh = channel('apm:openai-agents:model:start') |
| 27 | + |
| 28 | +// Reference to the loaded @openai/agents module, captured in the first hook |
| 29 | +// so that wrapResponseMethod can call getCurrentSpan() without an additional |
| 30 | +// require (and without triggering n/no-missing-require on agents-core internals). |
| 31 | +let agentsMod |
| 32 | + |
| 33 | +// @openai/agents >=0.8.0 moved addTraceProcessor / getCurrentSpan out of the |
| 34 | +// top-level re-exports. The new public surface uses getGlobalTraceProvider(): |
| 35 | +// provider.registerProcessor(processor) (replaces addTraceProcessor) |
| 36 | +// provider.getCurrentSpan() (replaces getCurrentSpan) |
| 37 | +// Both APIs are tried so this file works across the full supported version range. |
| 38 | +// The plugin subscriber (index.js) handles processor registration via the channel. |
| 39 | +function getCurrentSpanId (mod) { |
| 40 | + if (typeof mod?.getCurrentSpan === 'function') { |
| 41 | + return mod.getCurrentSpan()?.spanId |
| 42 | + } |
| 43 | + if (typeof mod?.getGlobalTraceProvider === 'function') { |
| 44 | + return mod.getGlobalTraceProvider().getCurrentSpan()?.spanId |
| 45 | + } |
| 46 | +} |
| 47 | + |
| 48 | +addHook({ name: '@openai/agents', versions: ['>=0.7.0'] }, (mod) => { |
| 49 | + if (patchedMods.has(mod)) return mod |
| 50 | + if (typeof mod?.addTraceProcessor !== 'function' && typeof mod?.getGlobalTraceProvider !== 'function') return mod |
| 51 | + patchedMods.add(mod) |
| 52 | + agentsMod = mod |
| 53 | + agentsCoreLoadedCh.publish({ mod }) |
| 54 | + return mod |
| 55 | +}) |
| 56 | + |
| 57 | +function wrapResponseMethod (original) { |
| 58 | + return function (...args) { |
| 59 | + const agentsCoreSpanId = getCurrentSpanId(agentsMod) |
| 60 | + publishClientBaseURL(this) |
| 61 | + return modelStartCh.runStores({ agentsCoreSpanId }, () => original.apply(this, args)) |
| 62 | + } |
| 63 | +} |
| 64 | + |
| 65 | +function wrapStreamedResponseMethod (original) { |
| 66 | + return function (...args) { |
| 67 | + const agentsCoreSpanId = getCurrentSpanId(agentsMod) |
| 68 | + publishClientBaseURL(this) |
| 69 | + const iterator = modelStartCh.runStores({ agentsCoreSpanId }, () => original.apply(this, args)) |
| 70 | + return wrapAsyncIterator(iterator, agentsCoreSpanId) |
| 71 | + } |
| 72 | +} |
| 73 | + |
| 74 | +function publishClientBaseURL (model) { |
| 75 | + const baseURL = model?.client?.baseURL ?? model?._client?.baseURL |
| 76 | + if (baseURL) responseClientCh.publish({ baseURL }) |
| 77 | +} |
| 78 | + |
| 79 | +function wrapAsyncIterator (iterator, agentsCoreSpanId) { |
| 80 | + if (!iterator || typeof iterator !== 'object') return iterator |
| 81 | + |
| 82 | + return { |
| 83 | + next () { |
| 84 | + return modelStartCh.runStores({ agentsCoreSpanId }, () => iterator.next.apply(iterator, arguments)) |
| 85 | + }, |
| 86 | + throw () { |
| 87 | + if (typeof iterator.throw !== 'function') return Promise.reject(arguments[0]) |
| 88 | + return modelStartCh.runStores({ agentsCoreSpanId }, () => iterator.throw.apply(iterator, arguments)) |
| 89 | + }, |
| 90 | + return () { |
| 91 | + if (typeof iterator.return !== 'function') return Promise.resolve({ done: true, value: arguments[0] }) |
| 92 | + return modelStartCh.runStores({ agentsCoreSpanId }, () => iterator.return.apply(iterator, arguments)) |
| 93 | + }, |
| 94 | + [Symbol.asyncIterator] () { |
| 95 | + return this |
| 96 | + }, |
| 97 | + } |
| 98 | +} |
| 99 | + |
| 100 | +addHook({ name: '@openai/agents-openai', versions: ['>=0.7.0'] }, (mod) => { |
| 101 | + if (patchedMods.has(mod)) return mod |
| 102 | + const proto = mod?.OpenAIResponsesModel?.prototype |
| 103 | + if (!proto) return mod |
| 104 | + |
| 105 | + patchedMods.add(mod) |
| 106 | + shimmer.wrap(proto, 'getResponse', wrapResponseMethod) |
| 107 | + shimmer.wrap(proto, 'getStreamedResponse', wrapStreamedResponseMethod) |
| 108 | + return mod |
| 109 | +}) |
0 commit comments