fix(llmobs): tagger reads propogated mlApp and sessionId from registry tags (#5102)

sabrenner · web-flow · commit af641d60efd1 · 2025-01-13T14:28:43.000-05:00
* wip

* get rid of unecessary test
diff --git a/packages/dd-trace/src/llmobs/tagger.js b/packages/dd-trace/src/llmobs/tagger.js
@@ -60,10 +60,10 @@ class LLMObsTagger {
     if (modelName) this._setTag(span, MODEL_NAME, modelName)
     if (modelProvider) this._setTag(span, MODEL_PROVIDER, modelProvider)
 
-    sessionId = sessionId || parent?.context()._tags[SESSION_ID]
+    sessionId = sessionId || registry.get(parent)?.[SESSION_ID]
     if (sessionId) this._setTag(span, SESSION_ID, sessionId)
 
-    if (!mlApp) mlApp = parent?.context()._tags[ML_APP] || this._config.llmobs.mlApp
+    if (!mlApp) mlApp = registry.get(parent)?.[ML_APP] || this._config.llmobs.mlApp
     this._setTag(span, ML_APP, mlApp)
 
     const parentId =
diff --git a/packages/dd-trace/test/llmobs/sdk/integration.spec.js b/packages/dd-trace/test/llmobs/sdk/integration.spec.js
@@ -177,45 +177,6 @@ describe('end to end sdk integration tests', () => {
     check(expected, llmobsSpans)
   })
 
-  it('instruments and uninstruments as needed', () => {
-    payloadGenerator = function () {
-      llmobs.disable()
-      llmobs.trace({ kind: 'agent', name: 'llmobsParent' }, () => {
-        llmobs.annotate({ inputData: 'hello', outputData: 'world' })
-        llmobs.enable({ mlApp: 'test1' })
-        llmobs.trace({ kind: 'workflow', name: 'child1' }, () => {
-          llmobs.disable()
-          llmobs.trace({ kind: 'workflow', name: 'child2' }, () => {
-            llmobs.enable({ mlApp: 'test2' })
-            llmobs.trace({ kind: 'workflow', name: 'child3' }, () => {})
-          })
-        })
-      })
-    }
-
-    const { spans, llmobsSpans } = run(payloadGenerator)
-    expect(spans).to.have.lengthOf(4)
-    expect(llmobsSpans).to.have.lengthOf(2)
-
-    const expected = [
-      expectedLLMObsNonLLMSpanEvent({
-        span: spans[1],
-        spanKind: 'workflow',
-        tags: { ...tags, ml_app: 'test1' },
-        name: 'child1'
-      }),
-      expectedLLMObsNonLLMSpanEvent({
-        span: spans[3],
-        spanKind: 'workflow',
-        tags: { ...tags, ml_app: 'test2' },
-        name: 'child3',
-        parentId: spans[1].context().toSpanId()
-      })
-    ]
-
-    check(expected, llmobsSpans)
-  })
-
   it('submits evaluations', () => {
     sinon.stub(Date, 'now').returns(1234567890)
     payloadGenerator = function () {
diff --git a/packages/dd-trace/test/llmobs/tagger.spec.js b/packages/dd-trace/test/llmobs/tagger.spec.js
@@ -111,14 +111,16 @@ describe('tagger', () => {
         const parentSpan = {
           context () {
             return {
-              _tags: {
-                '_ml_obs.meta.ml_app': 'my-ml-app',
-                '_ml_obs.session_id': 'my-session'
-              },
               toSpanId () { return '5678' }
             }
           }
         }
+
+        Tagger.tagMap.set(parentSpan, {
+          '_ml_obs.meta.ml_app': 'my-ml-app',
+          '_ml_obs.session_id': 'my-session'
+        })
+
         tagger.registerLLMObsSpan(span, { kind: 'llm', parent: parentSpan })
 
         expect(Tagger.tagMap.get(span)).to.deep.equal({
