make llmobs storage accessible

DataDog · Oct 29, 2024 · 845c840 · 845c840
1 parent cc6fec5
commit 845c840
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 6 deletions.
diff --git a/packages/dd-trace/src/llmobs/index.js b/packages/dd-trace/src/llmobs/index.js
@@ -2,7 +2,7 @@
 
 const log = require('../log')
 const { PROPAGATED_PARENT_ID_KEY } = require('./constants/tags')
-const { storage } = require('../../../datadog-core')
+const { storage } = require('./storage')
 
 const LLMObsSpanProcessor = require('./span_processor')
 
@@ -63,7 +63,7 @@ function disable () {
 // since LLMObs traces can extend between services and be the same trace,
 // we need to propogate the parent id.
 function handleLLMObsParentIdInjection ({ carrier }) {
-  const parent = storage.getStore()?.llmobsSpan
+  const parent = storage.getStore()?.span
   if (!parent) return
 
   const parentId = parent?.context().toSpanId()

diff --git a/packages/dd-trace/src/llmobs/sdk.js b/packages/dd-trace/src/llmobs/sdk.js
@@ -8,8 +8,7 @@ const {
 } = require('./util')
 const { isTrue } = require('../util')
 
-const { AsyncLocalStorage } = require('async_hooks')
-const storage = new AsyncLocalStorage()
+const { storage } = require('./storage')
 
 const Span = require('../opentracing/span')
 

diff --git a/packages/dd-trace/src/llmobs/storage.js b/packages/dd-trace/src/llmobs/storage.js
@@ -0,0 +1,7 @@
+'use strict'
+
+// TODO: remove this and use namespaced storage once available
+const { AsyncLocalStorage } = require('async_hooks')
+const storage = new AsyncLocalStorage()
+
+module.exports = { storage }
diff --git a/packages/dd-trace/test/llmobs/index.spec.js b/packages/dd-trace/test/llmobs/index.spec.js
@@ -43,7 +43,7 @@ describe('module', () => {
       '../log': logger,
       './writers/spans/agentless': LLMObsAgentlessSpanWriter,
       './writers/spans/agentProxy': LLMObsAgentProxySpanWriter,
-      '../../../datadog-core': {
+      './storage': {
         storage: {
           getStore () {
             return store
@@ -73,7 +73,7 @@ describe('module', () => {
   describe('handle llmobs info injection', () => {
     it('injects LLMObs parent ID when there is a parent LLMObs span', () => {
       llmobsModule.enable(config)
-      store.llmobsSpan = {
+      store.span = {
         context () {
           return {
             toSpanId () {

diff --git a/packages/dd-trace/test/llmobs/sdk/index.spec.js b/packages/dd-trace/test/llmobs/sdk/index.spec.js
@@ -10,6 +10,9 @@ const LLMObsSpanProcessor = require('../../../src/llmobs/span_processor')
 
 const tracerVersion = require('../../../../../package.json').version
 
+const { channel } = require('dc-polyfill')
+const injectCh = channel('dd-trace:span:inject')
+
 describe('sdk', () => {
   let LLMObsSDK
   let llmobs
@@ -1005,4 +1008,20 @@ describe('sdk', () => {
       expect(() => llmobs.flush()).to.not.throw()
     })
   })
+
+  describe('distributed', () => {
+    it('adds the current llmobs span id to the injection context', () => {
+      const carrier = { 'x-datadog-tags': '' }
+      let parentId
+      llmobs.trace({ kind: 'workflow', name: 'myWorkflow' }, span => {
+        parentId = span.context().toSpanId()
+
+        // simulate injection from http integration or from tracer
+        // something that triggers the text_map injection
+        injectCh.publish({ carrier })
+      })
+
+      expect(carrier['x-datadog-tags']).to.equal(`,_dd.p.llmobs_parent_id=${parentId}`)
+    })
+  })
 })