address comments

Author: jj22ee

aws-distro-opentelemetry-node-autoinstrumentation/src/llo-handler.ts

@@ -1,7 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { diag, Attributes, HrTime, ROOT_CONTEXT, createContextKey } from '@opentelemetry/api';
+import { Attributes, HrTime, ROOT_CONTEXT, createContextKey } from '@opentelemetry/api';
 import { LoggerProvider } from '@opentelemetry/sdk-logs';
 import { EventLoggerProvider } from '@opentelemetry/sdk-events';
 import { Event } from '@opentelemetry/api-events';
@@ -11,6 +11,7 @@ import { AnyValue } from '@opentelemetry/api-logs';
 const ROLE_SYSTEM = 'system';
 const ROLE_USER = 'user';
 const ROLE_ASSISTANT = 'assistant';
+const SESSION_ID = 'session.id';
 
 // Types of LLO attribute patterns
 export enum PatternType {
@@ -381,14 +382,14 @@ export class LLOHandler {
     const eventLogger = this.eventLoggerProvider.getEventLogger(span.instrumentationLibrary.name);
 
     // Hack - Workaround to add a custom-made Context to an Event so that the emitted event log
-    // has the correct associated traceId and spanId. This is needed because a ReadableSpan only
+    // has the correct associated traceId, spanId, flag. This is needed because a ReadableSpan only
     // provides its SpanContext, but does not provide access to its associated Context. We can only
     // provide a Context to an Event, but not a SpanContext. Here we attach a custom Context that
     // is associated to the ReadableSpan to mimic the ReadableSpan's actual Context.
     //
     // When a log record instance is created from this event, it will use the "custom Context" to
     // extract the ReadableSpan from the "custom Context", then extract the SpanContext from the
-    // RedableSpan. This way, the emitted log event has the correct SpanContext (traceId, spanId).
+    // RedableSpan. This way, the emitted log event has the correct SpanContext (traceId, spanId, flag).
     // - https://github.com/open-telemetry/opentelemetry-js/blob/experimental/v0.57.1/experimental/packages/sdk-logs/src/LogRecord.ts#L101
     // - https://github.com/open-telemetry/opentelemetry-js/blob/experimental/v0.57.1/api/src/trace/context-utils.ts#L78-L85
     //
@@ -404,8 +405,13 @@ export class LLOHandler {
       context: customContext,
     };
 
+    if (span.attributes[SESSION_ID]) {
+      event.attributes = {
+        [SESSION_ID]: span.attributes[SESSION_ID],
+      };
+    }
+
     eventLogger.emit(event);
-    diag.debug('Emitted Gen AI Event with input/output message format');
   }
 
   /**

aws-distro-opentelemetry-node-autoinstrumentation/src/otlp-aws-span-exporter.ts

@@ -8,7 +8,6 @@ import { ExportResult } from '@opentelemetry/core';
 import { LLOHandler } from './llo-handler';
 import { LoggerProvider as APILoggerProvider, logs } from '@opentelemetry/api-logs';
 import { LoggerProvider } from '@opentelemetry/sdk-logs';
-import { LogsAPI } from '@opentelemetry/api-logs/build/src/api/logs';
 import { getNodeVersion, isAgentObservabilityEnabled } from './utils';
 import { diag } from '@opentelemetry/api';
 
@@ -36,7 +35,6 @@ export class OTLPAwsSpanExporter extends OTLPProtoTraceExporter {
 
   private lloHandler: LLOHandler | undefined;
   private loggerProvider: APILoggerProvider | undefined;
-  private logsApi: LogsAPI = logs;
 
   constructor(endpoint: string, config?: OTLPExporterNodeConfigBase, loggerProvider?: APILoggerProvider) {
     super(OTLPAwsSpanExporter.changeUrlConfig(endpoint, config));
@@ -48,14 +46,15 @@ export class OTLPAwsSpanExporter extends OTLPProtoTraceExporter {
     this.loggerProvider = loggerProvider;
   }
 
-  // Lazily initialize LLO handler when needed to avoid initialization order issues"""
+  // Lazily initialize LLO handler when needed to avoid initialization order issues
   private ensureLloHandler(): boolean {
     if (!this.lloHandler && isAgentObservabilityEnabled()) {
       // If loggerProvider wasn't provided, try to get the current one
       if (!this.loggerProvider) {
         try {
-          this.loggerProvider = this.logsApi.getLoggerProvider();
+          this.loggerProvider = logs.getLoggerProvider();
         } catch (e: unknown) {
+          diag.debug('Failed to get logger provider', e);
           return false;
         }
       }

aws-distro-opentelemetry-node-autoinstrumentation/test/llo-handler.events.test.ts

@@ -552,4 +552,87 @@ describe('TestLLOHandlerEvents', () => {
     expect(outputMessages[0].content).toBe('Quantum computing is...');
     expect(outputMessages[0].role).toBe('assistant');
   });
+
+  /**
+   * Verify session.id attribute from span is copied to event attributes when present.
+   */
+  it('emitLloAttributes should copy session.id to event attributes when present', () => {
+    const attributes = {
+      'session.id': 'test-session-123',
+      'gen_ai.prompt': 'Hello, AI',
+      'gen_ai.completion': 'Hello! How can I help you?',
+    };
+
+    const span = testBase.createMockSpan(attributes);
+    testBase.updateMockSpanKey(span, 'endTime', [1234567899, 0]);
+    testBase.updateMockSpanKey(span, 'instrumentationLibrary', { name: 'test.scope', version: '1.0.0' });
+
+    testBase.lloHandler['emitLloAttributes'](span, attributes);
+
+    sinon.assert.calledOnce(testBase.eventLoggerMock.emit as any);
+    const emittedEvent = (testBase.eventLoggerMock.emit as any).getCall(0).args[0] as Event;
+
+    // Verify session.id was copied to event attributes
+    expect(emittedEvent.attributes).toBeDefined();
+    expect(emittedEvent.attributes?.['session.id']).toBe('test-session-123');
+
+    // Verify event body still contains LLO data
+    const eventBody = emittedEvent.data as any;
+    expect(eventBody.input).toBeDefined();
+    expect(eventBody.output).toBeDefined();
+  });
+
+  /**
+   * Verify event attributes do not contain session.id when not present in span attributes.
+   */
+  it('emitLloAttributes should not include session.id in event attributes when not present', () => {
+    const attributes = {
+      'gen_ai.prompt': 'Hello, AI',
+      'gen_ai.completion': 'Hello! How can I help you?',
+    };
+
+    const span = testBase.createMockSpan(attributes);
+    testBase.updateMockSpanKey(span, 'endTime', [1234567899, 0]);
+    testBase.updateMockSpanKey(span, 'instrumentationLibrary', { name: 'test.scope', version: '1.0.0' });
+
+    testBase.lloHandler['emitLloAttributes'](span, attributes);
+
+    sinon.assert.calledOnce(testBase.eventLoggerMock.emit as any);
+    const emittedEvent = (testBase.eventLoggerMock.emit as any).getCall(0).args[0] as Event;
+
+    // Verify session.id is not in event attributes (because the event doesn't have attributes)
+    expect(emittedEvent.attributes).toBeUndefined();
+    expect(emittedEvent).not.toHaveProperty('attributes');
+  });
+
+  /**
+   * Verify only session.id is copied from span attributes when mixed with other attributes.
+   */
+  it('emitLloAttributes should only copy session.id when mixed with other attributes', () => {
+    const attributes = {
+      'session.id': 'session-456',
+      'user.id': 'user-789',
+      'gen_ai.prompt': "What's the weather?",
+      'gen_ai.completion': "I can't check the weather.",
+      'other.attribute': 'some-value',
+    };
+
+    const span = testBase.createMockSpan(attributes);
+    testBase.updateMockSpanKey(span, 'endTime', [1234567899, 0]);
+    testBase.updateMockSpanKey(span, 'instrumentationLibrary', { name: 'test.scope', version: '1.0.0' });
+
+    testBase.lloHandler['emitLloAttributes'](span, attributes);
+
+    sinon.assert.calledOnce(testBase.eventLoggerMock.emit as any);
+    const emittedEvent = (testBase.eventLoggerMock.emit as any).getCall(0).args[0] as Event;
+
+    // Verify only session.id was copied to event attributes (plus event.name from Event class)
+    expect(emittedEvent.attributes).toBeDefined();
+    expect(emittedEvent.attributes?.['session.id']).toBe('session-456');
+    // Verify other span attributes were not copied
+    expect(emittedEvent.attributes).not.toHaveProperty('user.id');
+    expect(emittedEvent.attributes).not.toHaveProperty('other.attribute');
+    expect(emittedEvent.attributes).not.toHaveProperty('gen_ai.prompt');
+    expect(emittedEvent.attributes).not.toHaveProperty('gen_ai.completion');
+  });
 });