@illuma-ai/agents 1.0.90 → 1.0.94

package/src/specs/observability.integration.test.ts

@@ -0,0 +1,2223 @@
+/**
+ * Integration tests for Illuma Observability SDK integration in agents.
+ *
+ * These tests verify that the ObservabilityCallbackHandler and IllumaSpanProcessor
+ * correctly route LangChain/OTel events to trace payloads with proper structure,
+ * parent-child relationships, and token usage metrics.
+ *
+ * Uses a mock client to capture enqueued events without hitting a real server.
+ */
+/* eslint-disable @typescript-eslint/no-explicit-any, @typescript-eslint/explicit-function-return-type, @typescript-eslint/no-unused-vars, @typescript-eslint/no-require-imports, import/no-unresolved */
+
+import { ObservabilityCallbackHandler } from '@illuma-ai/observability-langchain';
+import type { Serialized } from '@langchain/core/load/serializable';
+import type { LLMResult } from '@langchain/core/outputs';
+import type { Document } from '@langchain/core/documents';
+
+/** Minimal IngestionEvent type for our mock (avoids importing core) */
+interface IngestionEvent {
+  id: string;
+  type: string;
+  timestamp: string;
+  body: Record<string, unknown>;
+}
+
+// ---------------------------------------------------------------------------
+// Mock Client - captures enqueued events for assertions
+// ---------------------------------------------------------------------------
+
+class MockObservabilityClient {
+  events: IngestionEvent[] = [];
+  flushed = false;
+  shutdownCalled = false;
+
+  enqueue(event: IngestionEvent): void {
+    this.events.push(event);
+  }
+
+  async flush(): Promise<void> {
+    this.flushed = true;
+  }
+
+  async shutdown(): Promise<void> {
+    this.shutdownCalled = true;
+    await this.flush();
+  }
+
+  /** Get events filtered by type */
+  getByType(type: string): IngestionEvent[] {
+    return this.events.filter((e) => e.type === type);
+  }
+
+  /** Get all event types in order */
+  getEventTypes(): string[] {
+    return this.events.map((e) => e.type);
+  }
+
+  /** Pretty-print all events for debugging */
+  printEvents(): void {
+    for (const event of this.events) {
+      const body = event.body as Record<string, unknown>;
+      const usage = body.usage as Record<string, unknown> | undefined;
+      console.log(
+        `  ${event.type} | name=${body.name ?? '—'} | id=${body.id ?? '—'} | traceId=${body.traceId ?? '—'} | parentObsId=${body.parentObservationId ?? '—'}` +
+          (usage
+            ? ` | usage={prompt:${usage.promptTokens ?? '—'},completion:${usage.completionTokens ?? '—'},total:${usage.totalTokens ?? '—'}}`
+            : '') +
+          (body.model ? ` | model=${body.model}` : '') +
+          (body.provider ? ` | provider=${body.provider}` : '') +
+          (body.level === 'ERROR' ? ` | ERROR: ${body.statusMessage}` : '')
+      );
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+const serialized = (name: string): Serialized => ({
+  lc: 1,
+  type: 'not_implemented',
+  id: ['langchain', name],
+});
+
+function makeUUID(): string {
+  return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
+    const r = (Math.random() * 16) | 0;
+    const v = c === 'x' ? r : (r & 0x3) | 0x8;
+    return v.toString(16);
+  });
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe('ObservabilityCallbackHandler integration', () => {
+  let client: MockObservabilityClient;
+  let handler: ObservabilityCallbackHandler;
+
+  beforeEach(() => {
+    client = new MockObservabilityClient();
+    handler = new ObservabilityCallbackHandler({
+      client,
+      traceName: 'test-agent-run',
+      userId: 'user-42',
+      sessionId: 'session-abc',
+      tags: ['integration-test'],
+      metadata: { messageId: 'msg-001' },
+      environment: 'test',
+      debug: false,
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 1. Root trace creation
+  // -----------------------------------------------------------------------
+
+  describe('root trace creation', () => {
+    it('should create a trace on the first chain start', async () => {
+      const chainRunId = makeUUID();
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'Hello agent' },
+        chainRunId
+      );
+
+      const traceEvents = client.getByType('trace-create');
+      expect(traceEvents).toHaveLength(1);
+
+      const trace = traceEvents[0].body as Record<string, unknown>;
+      expect(trace.name).toBe('test-agent-run');
+      expect(trace.userId).toBe('user-42');
+      expect(trace.sessionId).toBe('session-abc');
+      expect(trace.tags).toEqual(['integration-test']);
+      expect(trace.metadata).toEqual({ messageId: 'msg-001' });
+      expect(trace.environment).toBe('test');
+      expect(trace.input).toEqual({ input: 'Hello agent' });
+    });
+
+    it('should not create a second trace for nested chains', async () => {
+      const outerRunId = makeUUID();
+      const innerRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'outer' },
+        outerRunId
+      );
+      await handler.handleChainStart(
+        serialized('ChatPromptTemplate'),
+        { input: 'inner' },
+        innerRunId,
+        outerRunId
+      );
+
+      const traceEvents = client.getByType('trace-create');
+      expect(traceEvents).toHaveLength(1); // Only one trace, not two
+    });
+
+    it('should return the trace ID after creation', async () => {
+      expect(handler.getTraceId()).toBeNull();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        makeUUID()
+      );
+
+      expect(handler.getTraceId()).toBeTruthy();
+      expect(typeof handler.getTraceId()).toBe('string');
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 2. LLM generation events
+  // -----------------------------------------------------------------------
+
+  describe('LLM generation tracing', () => {
+    it('should create generation-create with input prompts', async () => {
+      const chainRunId = makeUUID();
+      const llmRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+
+      await handler.handleLLMStart(
+        serialized('ChatAnthropic'),
+        ['You are a helpful assistant.\n\nHuman: Hello'],
+        llmRunId,
+        chainRunId,
+        undefined,
+        undefined,
+        { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' }
+      );
+
+      const genEvents = client.getByType('generation-create');
+      expect(genEvents).toHaveLength(1);
+
+      const gen = genEvents[0].body as Record<string, unknown>;
+      expect(gen.name).toBe('ChatAnthropic');
+      expect(gen.model).toBe('claude-sonnet-4-20250514');
+      expect(gen.provider).toBe('anthropic');
+      expect(gen.input).toEqual([
+        'You are a helpful assistant.\n\nHuman: Hello',
+      ]);
+      expect(gen.traceId).toBe(handler.getTraceId());
+
+      // Parent should be the chain
+      const chainEvents = client.getByType('chain-create');
+      const chainObsId = (chainEvents[0].body as Record<string, unknown>).id;
+      expect(gen.parentObservationId).toBe(chainObsId);
+    });
+
+    it('should create generation-update with token usage on LLM end', async () => {
+      const chainRunId = makeUUID();
+      const llmRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleLLMStart(
+        serialized('ChatOpenAI'),
+        ['Hello'],
+        llmRunId,
+        chainRunId
+      );
+
+      const llmResult: LLMResult = {
+        generations: [
+          [{ text: 'Hi there! How can I help?', generationInfo: {} }],
+        ],
+        llmOutput: {
+          tokenUsage: {
+            promptTokens: 150,
+            completionTokens: 42,
+            totalTokens: 192,
+          },
+          modelName: 'gpt-4o',
+          provider: 'openai',
+        },
+      };
+
+      await handler.handleLLMEnd(llmResult, llmRunId);
+
+      const genUpdates = client.getByType('generation-update');
+      expect(genUpdates).toHaveLength(1);
+
+      const update = genUpdates[0].body as Record<string, unknown>;
+      expect(update.output).toBe('Hi there! How can I help?');
+      expect(update.model).toBe('gpt-4o');
+      expect(update.provider).toBe('openai');
+      expect(update.endTime).toBeDefined();
+
+      const usage = update.usage as Record<string, unknown>;
+      expect(usage.promptTokens).toBe(150);
+      expect(usage.completionTokens).toBe(42);
+      expect(usage.totalTokens).toBe(192);
+    });
+
+    it('should handle LLM errors with ERROR level', async () => {
+      const chainRunId = makeUUID();
+      const llmRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleLLMStart(
+        serialized('ChatAnthropic'),
+        ['Hello'],
+        llmRunId,
+        chainRunId
+      );
+
+      await handler.handleLLMError(new Error('Rate limit exceeded'), llmRunId);
+
+      const genUpdates = client.getByType('generation-update');
+      expect(genUpdates).toHaveLength(1);
+
+      const update = genUpdates[0].body as Record<string, unknown>;
+      expect(update.level).toBe('ERROR');
+      expect(update.statusMessage).toBe('Rate limit exceeded');
+      expect(update.endTime).toBeDefined();
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 3. Chain events
+  // -----------------------------------------------------------------------
+
+  describe('chain tracing', () => {
+    it('should create chain-create events with input', async () => {
+      const chainRunId = makeUUID();
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { query: 'What is AI?' },
+        chainRunId
+      );
+
+      const chainEvents = client.getByType('chain-create');
+      expect(chainEvents).toHaveLength(1);
+
+      const chain = chainEvents[0].body as Record<string, unknown>;
+      expect(chain.name).toBe('RunnableSequence');
+      expect(chain.input).toEqual({ query: 'What is AI?' });
+      expect(chain.traceId).toBe(handler.getTraceId());
+    });
+
+    it('should create span-update on chain end with output', async () => {
+      const chainRunId = makeUUID();
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+
+      await handler.handleChainEnd(
+        { output: 'AI is artificial intelligence.' },
+        chainRunId
+      );
+
+      const spanUpdates = client.getByType('span-update');
+      expect(spanUpdates).toHaveLength(1);
+
+      const update = spanUpdates[0].body as Record<string, unknown>;
+      expect(update.output).toEqual({
+        output: 'AI is artificial intelligence.',
+      });
+      expect(update.endTime).toBeDefined();
+    });
+
+    it('should handle chain error with ERROR level', async () => {
+      const chainRunId = makeUUID();
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+
+      await handler.handleChainError(
+        new Error('Chain failed: missing tool'),
+        chainRunId
+      );
+
+      const spanUpdates = client.getByType('span-update');
+      const errorUpdate = spanUpdates.find(
+        (e) => (e.body as Record<string, unknown>).level === 'ERROR'
+      );
+      expect(errorUpdate).toBeDefined();
+      expect((errorUpdate!.body as Record<string, unknown>).statusMessage).toBe(
+        'Chain failed: missing tool'
+      );
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 4. Tool events
+  // -----------------------------------------------------------------------
+
+  describe('tool tracing', () => {
+    it('should create tool-create events', async () => {
+      const chainRunId = makeUUID();
+      const toolRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleToolStart(
+        serialized('web_search'),
+        '{"query": "latest AI news"}',
+        toolRunId,
+        chainRunId
+      );
+
+      const toolEvents = client.getByType('tool-create');
+      expect(toolEvents).toHaveLength(1);
+
+      const tool = toolEvents[0].body as Record<string, unknown>;
+      expect(tool.name).toBe('Tool: web_search');
+      expect(tool.input).toBe('{"query": "latest AI news"}');
+      expect(tool.traceId).toBe(handler.getTraceId());
+    });
+
+    it('should create span-update on tool end', async () => {
+      const chainRunId = makeUUID();
+      const toolRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleToolStart(
+        serialized('calculator'),
+        '{"expression": "2+2"}',
+        toolRunId,
+        chainRunId
+      );
+
+      await handler.handleToolEnd('4', toolRunId);
+
+      const spanUpdates = client.getByType('span-update');
+      expect(spanUpdates).toHaveLength(1);
+      expect((spanUpdates[0].body as Record<string, unknown>).output).toBe('4');
+    });
+
+    it('should handle tool error', async () => {
+      const chainRunId = makeUUID();
+      const toolRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleToolStart(
+        serialized('api_call'),
+        '{"url": "https://example.com"}',
+        toolRunId,
+        chainRunId
+      );
+
+      await handler.handleToolError(new Error('Connection timeout'), toolRunId);
+
+      const spanUpdates = client.getByType('span-update');
+      const errorUpdate = spanUpdates.find(
+        (e) => (e.body as Record<string, unknown>).level === 'ERROR'
+      );
+      expect(errorUpdate).toBeDefined();
+      expect((errorUpdate!.body as Record<string, unknown>).statusMessage).toBe(
+        'Connection timeout'
+      );
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 5. Retriever events
+  // -----------------------------------------------------------------------
+
+  describe('retriever tracing', () => {
+    it('should create retriever-create events', async () => {
+      const chainRunId = makeUUID();
+      const retrieverRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleRetrieverStart(
+        serialized('VectorStoreRetriever'),
+        'How does photosynthesis work?',
+        retrieverRunId,
+        chainRunId
+      );
+
+      const retrieverEvents = client.getByType('retriever-create');
+      expect(retrieverEvents).toHaveLength(1);
+
+      const retriever = retrieverEvents[0].body as Record<string, unknown>;
+      expect(retriever.name).toBe('Retriever: VectorStoreRetriever');
+      expect(retriever.input).toBe('How does photosynthesis work?');
+    });
+
+    it('should produce span-update with documents on retriever end', async () => {
+      const chainRunId = makeUUID();
+      const retrieverRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleRetrieverStart(
+        serialized('VectorStoreRetriever'),
+        'photosynthesis',
+        retrieverRunId,
+        chainRunId
+      );
+
+      const documents: Document[] = [
+        {
+          pageContent: 'Photosynthesis is the process...',
+          metadata: { source: 'wiki', page: 1 },
+        },
+        {
+          pageContent: 'Plants use chlorophyll...',
+          metadata: { source: 'textbook', page: 42 },
+        },
+      ];
+
+      await handler.handleRetrieverEnd(documents, retrieverRunId);
+
+      const spanUpdates = client.getByType('span-update');
+      expect(spanUpdates).toHaveLength(1);
+
+      const update = spanUpdates[0].body as Record<string, unknown>;
+      const output = update.output as Array<Record<string, unknown>>;
+      expect(output).toHaveLength(2);
+      expect(output[0].pageContent).toBe('Photosynthesis is the process...');
+      expect(output[0].metadata).toEqual({ source: 'wiki', page: 1 });
+
+      const meta = update.metadata as Record<string, unknown>;
+      expect(meta.documentCount).toBe(2);
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 6. Parent-child hierarchy
+  // -----------------------------------------------------------------------
+
+  describe('parent-child hierarchy', () => {
+    it('should correctly link nested observations to parent', async () => {
+      const outerChainId = makeUUID();
+      const innerChainId = makeUUID();
+      const llmRunId = makeUUID();
+      const toolRunId = makeUUID();
+
+      // Outer chain starts
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'orchestrate' },
+        outerChainId
+      );
+
+      // Inner chain nested under outer
+      await handler.handleChainStart(
+        serialized('ChatPromptTemplate'),
+        { input: 'format prompt' },
+        innerChainId,
+        outerChainId
+      );
+
+      // LLM nested under inner chain
+      await handler.handleLLMStart(
+        serialized('ChatOpenAI'),
+        ['formatted prompt'],
+        llmRunId,
+        innerChainId
+      );
+
+      // Tool nested under outer chain
+      await handler.handleToolStart(
+        serialized('calculator'),
+        '2+2',
+        toolRunId,
+        outerChainId
+      );
+
+      // Verify hierarchy
+      const chainCreates = client.getByType('chain-create');
+      const genCreates = client.getByType('generation-create');
+      const toolCreates = client.getByType('tool-create');
+
+      // All share the same traceId
+      const traceId = handler.getTraceId();
+      expect((chainCreates[0].body as Record<string, unknown>).traceId).toBe(
+        traceId
+      );
+      expect((chainCreates[1].body as Record<string, unknown>).traceId).toBe(
+        traceId
+      );
+      expect((genCreates[0].body as Record<string, unknown>).traceId).toBe(
+        traceId
+      );
+      expect((toolCreates[0].body as Record<string, unknown>).traceId).toBe(
+        traceId
+      );
+
+      // Inner chain -> parent is outer chain's observationId
+      const outerChainObsId = (chainCreates[0].body as Record<string, unknown>)
+        .id;
+      const innerChainObsId = (chainCreates[1].body as Record<string, unknown>)
+        .id;
+      expect(
+        (chainCreates[1].body as Record<string, unknown>).parentObservationId
+      ).toBe(outerChainObsId);
+
+      // LLM -> parent is inner chain
+      expect(
+        (genCreates[0].body as Record<string, unknown>).parentObservationId
+      ).toBe(innerChainObsId);
+
+      // Tool -> parent is outer chain
+      expect(
+        (toolCreates[0].body as Record<string, unknown>).parentObservationId
+      ).toBe(outerChainObsId);
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 7. Full agent pipeline simulation
+  // -----------------------------------------------------------------------
+
+  describe('full agent pipeline simulation', () => {
+    it('should produce correct event sequence for a RAG agent with tools', async () => {
+      const chainRunId = makeUUID();
+      const agentChainId = makeUUID();
+      const retrieverRunId = makeUUID();
+      const llmRunId1 = makeUUID();
+      const toolRunId = makeUUID();
+      const llmRunId2 = makeUUID();
+
+      // 1. Root chain starts (the main graph)
+      await handler.handleChainStart(
+        serialized('CompiledStateGraph'),
+        {
+          messages: [
+            {
+              role: 'user',
+              content: 'What papers did OpenAI publish in 2024?',
+            },
+          ],
+        },
+        chainRunId
+      );
+
+      // 2. Agent chain starts (nested)
+      await handler.handleChainStart(
+        serialized('AgentExecutor'),
+        { input: 'What papers did OpenAI publish in 2024?' },
+        agentChainId,
+        chainRunId
+      );
+
+      // 3. First LLM call (decides to use tool)
+      await handler.handleLLMStart(
+        serialized('ChatAnthropic'),
+        ['System: You are a research assistant.\nHuman: What papers...'],
+        llmRunId1,
+        agentChainId,
+        undefined,
+        undefined,
+        { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' }
+      );
+
+      await handler.handleLLMEnd(
+        {
+          generations: [
+            [
+              {
+                text: "I'll search for OpenAI papers from 2024.",
+                generationInfo: {},
+              },
+            ],
+          ],
+          llmOutput: {
+            tokenUsage: {
+              promptTokens: 200,
+              completionTokens: 30,
+              totalTokens: 230,
+            },
+            modelName: 'claude-sonnet-4-20250514',
+            provider: 'anthropic',
+          },
+        },
+        llmRunId1
+      );
+
+      // 4. Retriever runs
+      await handler.handleRetrieverStart(
+        serialized('VectorStoreRetriever'),
+        'OpenAI papers 2024',
+        retrieverRunId,
+        agentChainId
+      );
+
+      await handler.handleRetrieverEnd(
+        [
+          {
+            pageContent: 'GPT-4o: omni-model for text, vision, audio...',
+            metadata: { source: 'arxiv', year: 2024 },
+          },
+          {
+            pageContent: 'Sora: text-to-video generation model...',
+            metadata: { source: 'blog', year: 2024 },
+          },
+        ] as Document[],
+        retrieverRunId
+      );
+
+      // 5. Tool call (web search)
+      await handler.handleToolStart(
+        serialized('web_search'),
+        '{"query": "OpenAI 2024 research papers list"}',
+        toolRunId,
+        agentChainId
+      );
+
+      await handler.handleToolEnd(
+        JSON.stringify([
+          { title: 'GPT-4o Technical Report', url: 'https://arxiv.org/...' },
+          { title: 'Sora', url: 'https://openai.com/sora' },
+        ]),
+        toolRunId
+      );
+
+      // 6. Second LLM call (final answer)
+      await handler.handleLLMStart(
+        serialized('ChatAnthropic'),
+        ['System: ...\nContext: ...\nHuman: What papers...'],
+        llmRunId2,
+        agentChainId,
+        undefined,
+        undefined,
+        { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' }
+      );
+
+      await handler.handleLLMEnd(
+        {
+          generations: [
+            [
+              {
+                text: 'OpenAI published several notable papers in 2024 including GPT-4o and Sora...',
+                generationInfo: {},
+              },
+            ],
+          ],
+          llmOutput: {
+            tokenUsage: {
+              promptTokens: 800,
+              completionTokens: 150,
+              totalTokens: 950,
+            },
+            modelName: 'claude-sonnet-4-20250514',
+            provider: 'anthropic',
+          },
+        },
+        llmRunId2
+      );
+
+      // 7. Chains end
+      await handler.handleChainEnd(
+        { output: 'OpenAI published several notable papers...' },
+        agentChainId
+      );
+      await handler.handleChainEnd(
+        {
+          messages: [
+            {
+              role: 'assistant',
+              content: 'OpenAI published several notable papers...',
+            },
+          ],
+        },
+        chainRunId
+      );
+
+      // Flush
+      await handler.flushAsync();
+      expect(client.flushed).toBe(true);
+
+      // --- Assertions ---
+
+      // Uncomment to debug: client.printEvents();
+
+      // Verify event counts
+      const eventTypes = client.getEventTypes();
+      const typeCounts = eventTypes.reduce(
+        (acc, t) => {
+          acc[t] = (acc[t] || 0) + 1;
+          return acc;
+        },
+        {} as Record<string, number>
+      );
+
+      expect(typeCounts['trace-create']).toBe(1); // Root trace
+      expect(typeCounts['chain-create']).toBe(1); // Outer chain only
+      expect(typeCounts['agent-create']).toBe(1); // AgentExecutor detected as agent!
+      expect(typeCounts['generation-create']).toBe(2); // Two LLM calls
+      expect(typeCounts['generation-update']).toBe(2); // Two LLM ends
+      expect(typeCounts['retriever-create']).toBe(1); // One retriever
+      expect(typeCounts['tool-create']).toBe(1); // One tool
+      expect(typeCounts['span-update']).toBe(4); // retriever end + tool end + 2 chain ends
+
+      // Verify all events share the same traceId
+      const traceId = handler.getTraceId()!;
+      for (const event of client.events) {
+        const body = event.body as Record<string, unknown>;
+        if (body.traceId) {
+          expect(body.traceId).toBe(traceId);
+        }
+      }
+
+      // Verify token usage in generation-update events
+      const genUpdates = client.getByType('generation-update');
+
+      // First LLM call usage
+      const firstLLMUpdate = genUpdates[0].body as Record<string, unknown>;
+      const firstUsage = firstLLMUpdate.usage as Record<string, unknown>;
+      expect(firstUsage.promptTokens).toBe(200);
+      expect(firstUsage.completionTokens).toBe(30);
+      expect(firstUsage.totalTokens).toBe(230);
+      expect(firstLLMUpdate.model).toBe('claude-sonnet-4-20250514');
+      expect(firstLLMUpdate.provider).toBe('anthropic');
+
+      // Second LLM call usage
+      const secondLLMUpdate = genUpdates[1].body as Record<string, unknown>;
+      const secondUsage = secondLLMUpdate.usage as Record<string, unknown>;
+      expect(secondUsage.promptTokens).toBe(800);
+      expect(secondUsage.completionTokens).toBe(150);
+      expect(secondUsage.totalTokens).toBe(950);
+
+      // Verify total token computation
+      const totalPrompt = 200 + 800;
+      const totalCompletion = 30 + 150;
+      const totalTokens = totalPrompt + totalCompletion;
+      expect(totalPrompt).toBe(1000);
+      expect(totalCompletion).toBe(180);
+      expect(totalTokens).toBe(1180);
+
+      // Total: 1180 tokens (1000 prompt + 180 completion) across 2 LLM calls
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 8. Token usage edge cases
+  // -----------------------------------------------------------------------
+
+  describe('token usage edge cases', () => {
+    it('should handle different token usage field names (prompt_tokens vs promptTokens)', async () => {
+      const chainRunId = makeUUID();
+      const llmRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('Chain'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleLLMStart(
+        serialized('ChatOpenAI'),
+        ['test'],
+        llmRunId,
+        chainRunId
+      );
+
+      // OpenAI style: uses prompt_tokens (snake_case)
+      await handler.handleLLMEnd(
+        {
+          generations: [[{ text: 'response', generationInfo: {} }]],
+          llmOutput: {
+            tokenUsage: {
+              prompt_tokens: 100,
+              completion_tokens: 50,
+              total_tokens: 150,
+            },
+          },
+        },
+        llmRunId
+      );
+
+      const updates = client.getByType('generation-update');
+      const usage = (updates[0].body as Record<string, unknown>)
+        .usage as Record<string, unknown>;
+
+      // The handler should normalize both field naming conventions
+      expect(usage.promptTokens ?? usage.prompt_tokens).toBeTruthy();
+    });
+
+    it('should handle missing token usage gracefully', async () => {
+      const chainRunId = makeUUID();
+      const llmRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('Chain'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleLLMStart(
+        serialized('ChatOpenAI'),
+        ['test'],
+        llmRunId,
+        chainRunId
+      );
+
+      // No llmOutput at all
+      await handler.handleLLMEnd(
+        { generations: [[{ text: 'response', generationInfo: {} }]] },
+        llmRunId
+      );
+
+      const updates = client.getByType('generation-update');
+      expect(updates).toHaveLength(1);
+      // Should not crash, usage should be present but with undefined values
+      const usage = (updates[0].body as Record<string, unknown>)
+        .usage as Record<string, unknown>;
+      expect(usage).toBeDefined();
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 9. Flush and shutdown
+  // -----------------------------------------------------------------------
+
+  describe('flush and shutdown', () => {
+    it('flushAsync should call client.flush()', async () => {
+      await handler.handleChainStart(
+        serialized('Chain'),
+        { input: 'test' },
+        makeUUID()
+      );
+      await handler.flushAsync();
+      expect(client.flushed).toBe(true);
+    });
+
+    it('shutdown should call client.flush() (not shutdown) for external client', async () => {
+      await handler.handleChainStart(
+        serialized('Chain'),
+        { input: 'test' },
+        makeUUID()
+      );
+      await handler.shutdown();
+      // Since we provided an external client, it should only flush, not shutdown
+      expect(client.flushed).toBe(true);
+      expect(client.shutdownCalled).toBe(false);
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 10. Agent detection
+  // -----------------------------------------------------------------------
+
+  describe('agent detection', () => {
+    it('should emit agent-create for chains with langgraph_node agent= metadata', async () => {
+      const outerChainId = makeUUID();
+      const agentChainId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('CompiledStateGraph'),
+        { messages: [{ role: 'user', content: 'Hello' }] },
+        outerChainId
+      );
+
+      // Agent chain with langgraph_node metadata
+      // Cast to any to pass tags/metadata params (ts-jest resolves different .d.ts for BaseCallbackHandler)
+      await (handler as any).handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'agent task' },
+        agentChainId,
+        outerChainId,
+        undefined, // tags
+        { langgraph_node: 'agent=research-agent' } // metadata
+      );
+
+      const agentEvents = client.getByType('agent-create');
+      expect(agentEvents).toHaveLength(1);
+
+      const agentBody = agentEvents[0].body as Record<string, unknown>;
+      expect(agentBody.name).toBe('Agent: research-agent');
+      expect(agentBody.traceId).toBe(handler.getTraceId());
+
+      const meta = agentBody.metadata as Record<string, unknown>;
+      expect(meta.agentId).toBe('research-agent');
+      expect(meta.langchainMetadata).toEqual({
+        langgraph_node: 'agent=research-agent',
+      });
+    });
+
+    it('should emit agent-create for AgentExecutor chain name', async () => {
+      const agentRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('AgentExecutor'),
+        { input: 'do something' },
+        agentRunId
+      );
+
+      const agentEvents = client.getByType('agent-create');
+      expect(agentEvents).toHaveLength(1);
+      expect((agentEvents[0].body as Record<string, unknown>).name).toBe(
+        'AgentExecutor'
+      );
+
+      // Should NOT produce a chain-create
+      const chainEvents = client.getByType('chain-create');
+      expect(chainEvents).toHaveLength(0);
+    });
+
+    it('should emit agent-create for chain names ending in Agent', async () => {
+      const agentRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('ReactAgent'),
+        { input: 'think and act' },
+        agentRunId
+      );
+
+      const agentEvents = client.getByType('agent-create');
+      expect(agentEvents).toHaveLength(1);
+      expect((agentEvents[0].body as Record<string, unknown>).name).toBe(
+        'ReactAgent'
+      );
+    });
+
+    it('should NOT emit agent-create for regular chains', async () => {
+      const chainRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'regular chain' },
+        chainRunId
+      );
+
+      const agentEvents = client.getByType('agent-create');
+      expect(agentEvents).toHaveLength(0);
+
+      const chainEvents = client.getByType('chain-create');
+      expect(chainEvents).toHaveLength(1);
+    });
+
+    it('should correctly end agent observations with span-update', async () => {
+      const agentRunId = makeUUID();
+
+      await (handler as any).handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'agent task' },
+        agentRunId,
+        undefined, // parentRunId
+        undefined, // tags
+        { langgraph_node: 'agent=planner' } // metadata
+      );
+
+      await handler.handleChainEnd({ output: 'agent completed' }, agentRunId);
+
+      const agentEvents = client.getByType('agent-create');
+      expect(agentEvents).toHaveLength(1);
+
+      const spanUpdates = client.getByType('span-update');
+      expect(spanUpdates).toHaveLength(1);
+
+      const update = spanUpdates[0].body as Record<string, unknown>;
+      expect(update.output).toEqual({ output: 'agent completed' });
+      expect(update.endTime).toBeDefined();
+      // The span-update should reference the same observation ID as agent-create
+      expect(update.id).toBe(
+        (agentEvents[0].body as Record<string, unknown>).id
+      );
+    });
+
+    it('should handle agent error with ERROR level', async () => {
+      const agentRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('AgentExecutor'),
+        { input: 'fail task' },
+        agentRunId
+      );
+
+      await handler.handleChainError(
+        new Error('Agent failed: tool not found'),
+        agentRunId
+      );
+
+      const agentEvents = client.getByType('agent-create');
+      expect(agentEvents).toHaveLength(1);
+
+      const spanUpdates = client.getByType('span-update');
+      const errorUpdate = spanUpdates.find(
+        (e) => (e.body as Record<string, unknown>).level === 'ERROR'
+      );
+      expect(errorUpdate).toBeDefined();
+      expect((errorUpdate!.body as Record<string, unknown>).statusMessage).toBe(
+        'Agent failed: tool not found'
+      );
+    });
+
+    it('should integrate agent-create in full pipeline with correct counts', async () => {
+      const rootChainId = makeUUID();
+      const agentChainId = makeUUID();
+      const llmRunId = makeUUID();
+
+      // Root chain
+      await handler.handleChainStart(
+        serialized('CompiledStateGraph'),
+        { messages: [{ role: 'user', content: 'Hello' }] },
+        rootChainId
+      );
+
+      // Agent chain (detected via metadata)
+      await (handler as any).handleChainStart(
+        serialized('RunnableSequence'),
+        { input: 'agent work' },
+        agentChainId,
+        rootChainId,
+        undefined, // tags
+        { langgraph_node: 'agent=qa-agent' } // metadata
+      );
+
+      // LLM inside agent
+      await handler.handleLLMStart(
+        serialized('ChatAnthropic'),
+        ['prompt'],
+        llmRunId,
+        agentChainId,
+        undefined,
+        undefined,
+        { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' }
+      );
+
+      await handler.handleLLMEnd(
+        {
+          generations: [[{ text: 'answer', generationInfo: {} }]],
+          llmOutput: {
+            tokenUsage: {
+              promptTokens: 100,
+              completionTokens: 20,
+              totalTokens: 120,
+            },
+            modelName: 'claude-sonnet-4-20250514',
+          },
+        },
+        llmRunId
+      );
+
+      await handler.handleChainEnd({ output: 'done' }, agentChainId);
+      await handler.handleChainEnd({ output: 'done' }, rootChainId);
+
+      const typeCounts = client.getEventTypes().reduce(
+        (acc, t) => {
+          acc[t] = (acc[t] || 0) + 1;
+          return acc;
+        },
+        {} as Record<string, number>
+      );
+
+      expect(typeCounts['trace-create']).toBe(1);
+      expect(typeCounts['chain-create']).toBe(1); // Only root chain
+      expect(typeCounts['agent-create']).toBe(1); // Agent detected!
+      expect(typeCounts['generation-create']).toBe(1);
+      expect(typeCounts['generation-update']).toBe(1);
+      expect(typeCounts['span-update']).toBe(2); // agent end + root chain end
+
+      // LLM parent should be the agent observation
+      const agentObsId = (
+        client.getByType('agent-create')[0].body as Record<string, unknown>
+      ).id;
+      const genCreate = client.getByType('generation-create')[0].body as Record<
+        string,
+        unknown
+      >;
+      expect(genCreate.parentObservationId).toBe(agentObsId);
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 11. Event structure validation
+  // -----------------------------------------------------------------------
+
+  describe('event structure validation', () => {
+    it('every event should have id, type, timestamp, and body', async () => {
+      const chainRunId = makeUUID();
+      const llmRunId = makeUUID();
+
+      await handler.handleChainStart(
+        serialized('Chain'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleLLMStart(
+        serialized('LLM'),
+        ['test'],
+        llmRunId,
+        chainRunId
+      );
+      await handler.handleLLMEnd(
+        {
+          generations: [[{ text: 'ok', generationInfo: {} }]],
+          llmOutput: { tokenUsage: {} },
+        },
+        llmRunId
+      );
+      await handler.handleChainEnd({ output: 'done' }, chainRunId);
+
+      for (const event of client.events) {
+        expect(event.id).toBeDefined();
+        expect(typeof event.id).toBe('string');
+        expect(event.type).toBeDefined();
+        expect(event.timestamp).toMatch(/^\d{4}-\d{2}-\d{2}T/);
+        expect(event.body).toBeDefined();
+        expect(typeof event.body).toBe('object');
+      }
+    });
+
+    it('observation events should have traceId and observationId', async () => {
+      const chainRunId = makeUUID();
+      await handler.handleChainStart(
+        serialized('Chain'),
+        { input: 'test' },
+        chainRunId
+      );
+      await handler.handleChainEnd({ output: 'done' }, chainRunId);
+
+      const nonTraceEvents = client.events.filter(
+        (e) => e.type !== 'trace-create'
+      );
+      for (const event of nonTraceEvents) {
+        const body = event.body as Record<string, unknown>;
+        expect(body.traceId).toBe(handler.getTraceId());
+        expect(body.id).toBeDefined();
+      }
+    });
+  });
+});
+
+// ---------------------------------------------------------------------------
+// IllumaSpanProcessor unit tests (OTel integration)
+// ---------------------------------------------------------------------------
+
+describe('IllumaSpanProcessor integration', () => {
+  it('should be importable from @illuma-ai/observability-otel', async () => {
+    const { IllumaSpanProcessor } = await import(
+      '@illuma-ai/observability-otel'
+    );
+    expect(IllumaSpanProcessor).toBeDefined();
+    expect(typeof IllumaSpanProcessor).toBe('function');
+  });
+});
+
+// ---------------------------------------------------------------------------
+// Core SDK features: sampling, PII masking, observation types, prompt/dataset
+// ---------------------------------------------------------------------------
+
+describe('Core SDK features', () => {
+  // Use require to bypass ts-jest module resolution issues with dynamic imports
+
+  const {
+    ObservabilityCoreClient,
+    TraceClient,
+    SpanClient,
+    IngestionEventType,
+  } = require('@illuma-ai/observability-core');
+
+  // Create a concrete subclass at runtime to avoid TS abstract class issues
+  const TestClientClass = class extends (ObservabilityCoreClient as any) {
+    protected async fetchWithRetry(): Promise<any> {
+      return {
+        status: 200,
+        statusText: 'OK',
+        ok: true,
+        json: async () => ({}),
+        text: async () => '',
+      };
+    }
+  };
+
+  function createTestClient(config: Record<string, unknown> = {}): {
+    client: any;
+    events: any[];
+  } {
+    const events: any[] = [];
+
+    const client = new (TestClientClass as any)({
+      publicKey: 'pk-test',
+      secretKey: 'sk-test',
+      baseUrl: 'http://localhost:9999',
+      flushInterval: 0,
+      ...config,
+    });
+
+    // Intercept enqueue to capture events
+    const origEnqueue = client.enqueue.bind(client);
+    client.enqueue = (event: any) => {
+      origEnqueue(event);
+      events.push(event);
+    };
+
+    return { client, events };
+  }
+
+  // -----------------------------------------------------------------------
+  // 12. Sampling rate
+  // -----------------------------------------------------------------------
+
+  describe('sampling rate', () => {
+    it('should trace everything when sampleRate=1.0 (default)', () => {
+      const { client, events } = createTestClient({ sampleRate: 1.0 });
+
+      for (let i = 0; i < 20; i++) {
+        client.trace({ name: `trace-${i}` });
+      }
+
+      // All 20 traces should be created (each produces a trace-create event queued)
+      const traceEvents = events.filter((e: any) => e.type === 'trace-create');
+      expect(traceEvents.length).toBe(20);
+    });
+
+    it('should trace nothing when sampleRate=0.0', () => {
+      const { client, events } = createTestClient({ sampleRate: 0.0 });
+
+      for (let i = 0; i < 20; i++) {
+        const trace = client.trace({ name: `trace-${i}` });
+        // Even child events should be silently dropped
+        trace.generation({ name: 'gen', model: 'gpt-4o' });
+        trace.span({ name: 'span' });
+      }
+
+      // No events should be enqueued at all
+      expect(events.length).toBe(0);
+    });
+
+    it('should sample approximately the right percentage', () => {
+      // Use a fixed seed approach: run many traces with 50% sample rate
+      const { client, events } = createTestClient({ sampleRate: 0.5 });
+
+      const iterations = 1000;
+      for (let i = 0; i < iterations; i++) {
+        client.trace({ name: `trace-${i}` });
+      }
+
+      const traceEvents = events.filter((e: any) => e.type === 'trace-create');
+      // With 1000 iterations at 50%, expect roughly 400-600 traces
+      expect(traceEvents.length).toBeGreaterThan(300);
+      expect(traceEvents.length).toBeLessThan(700);
+    });
+
+    it('should include all children once a trace is sampled in', () => {
+      const { client, events } = createTestClient({ sampleRate: 1.0 });
+
+      const trace = client.trace({ name: 'sampled-trace' });
+      trace.span({ name: 'child-span' });
+      trace.generation({ name: 'child-gen', model: 'gpt-4o' });
+      trace.agent({ name: 'child-agent' });
+      trace.tool({ name: 'child-tool' });
+
+      // 1 trace + 4 children = 5 events
+      expect(events.length).toBe(5);
+    });
+
+    it('should clamp sampleRate to [0, 1] range', () => {
+      // sampleRate > 1 should be clamped to 1 (trace everything)
+      const { events: events1 } = createTestClient({ sampleRate: 5.0 });
+      // sampleRate < 0 should be clamped to 0 (trace nothing)
+      const { client: client2, events: events2 } = createTestClient({
+        sampleRate: -1,
+      });
+
+      client2.trace({ name: 'should-not-appear' });
+      expect(events2.length).toBe(0);
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 13. PII masking
+  // -----------------------------------------------------------------------
+
+  describe('PII masking', () => {
+    it('should apply mask function to event bodies', () => {
+      const maskFunction = (body: Record<string, unknown>) => {
+        const masked = { ...body };
+        if (typeof masked.input === 'string') {
+          masked.input = masked.input.replace(
+            /[\w.-]+@[\w.-]+\.\w+/g,
+            '[EMAIL]'
+          );
+        }
+        return masked;
+      };
+
+      const { client } = createTestClient({ maskFunction });
+
+      const trace = client.trace({ name: 'pii-test' });
+      trace.span({
+        name: 'user-query',
+        input: 'My email is john@example.com and I need help' as any,
+      });
+
+      // Check the queued events - the span should have masked input
+      const queue = client.getQueue();
+      const spanEvent = [...queue].find((e: any) => e.type === 'span-create');
+      if (spanEvent) {
+        expect((spanEvent.body as any).input).toBe(
+          'My email is [EMAIL] and I need help'
+        );
+        expect((spanEvent.body as any).input).not.toContain('john@example.com');
+      }
+    });
+
+    it('should mask multiple PII patterns', () => {
+      const maskFunction = (body: Record<string, unknown>) => {
+        const masked = { ...body };
+        const str = JSON.stringify(masked);
+        const redacted = str
+          .replace(/[\w.-]+@[\w.-]+\.\w+/g, '[EMAIL]')
+          .replace(/\b\d{3}[-.]?\d{3}[-.]?\d{4}\b/g, '[PHONE]')
+          .replace(/\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b/g, '[CARD]');
+        return JSON.parse(redacted);
+      };
+
+      const { client } = createTestClient({ maskFunction });
+
+      const trace = client.trace({
+        name: 'multi-pii',
+        metadata: { userEmail: 'test@corp.com', phone: '555-123-4567' },
+      });
+
+      const queue = client.getQueue();
+      const traceEvent = [...queue].find((e: any) => e.type === 'trace-create');
+      if (traceEvent) {
+        const meta = (traceEvent.body as any).metadata;
+        expect(meta.userEmail).toBe('[EMAIL]');
+        expect(meta.phone).toBe('[PHONE]');
+      }
+    });
+
+    it('should gracefully handle mask function errors', () => {
+      const maskFunction = () => {
+        throw new Error('Mask function exploded');
+      };
+
+      const { client } = createTestClient({ maskFunction });
+
+      // Should not throw — falls back to unmasked
+      expect(() => {
+        client.trace({ name: 'should-not-crash' });
+      }).not.toThrow();
+
+      const queue = client.getQueue();
+      expect(queue.length).toBeGreaterThan(0);
+    });
+
+    it('should not modify events when no mask function is set', () => {
+      const { client } = createTestClient();
+
+      client.trace({
+        name: 'no-mask',
+        metadata: { email: 'visible@example.com' },
+      });
+
+      const queue = client.getQueue();
+      const traceEvent = [...queue].find((e: any) => e.type === 'trace-create');
+      expect((traceEvent!.body as any).metadata.email).toBe(
+        'visible@example.com'
+      );
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 14. All observation types (TraceClient + SpanClient)
+  // -----------------------------------------------------------------------
+
+  describe('observation types completeness', () => {
+    it('TraceClient should have all observation type methods', () => {
+      const enqueue = jest.fn();
+      const trace = new TraceClient('trace-123', enqueue);
+
+      // All typed observation methods should exist and return SpanClient or GenerationClient
+      expect(typeof trace.span).toBe('function');
+      expect(typeof trace.agent).toBe('function');
+      expect(typeof trace.tool).toBe('function');
+      expect(typeof trace.chain).toBe('function');
+      expect(typeof trace.retriever).toBe('function');
+      expect(typeof trace.guardrail).toBe('function');
+      expect(typeof trace.evaluator).toBe('function');
+      expect(typeof trace.embedding).toBe('function');
+      expect(typeof trace.generation).toBe('function');
+      expect(typeof trace.event).toBe('function');
+      expect(typeof trace.score).toBe('function');
+    });
+
+    it('SpanClient should have all observation type methods', () => {
+      const enqueue = jest.fn();
+      const span = new SpanClient('span-123', 'trace-123', enqueue);
+
+      expect(typeof span.span).toBe('function');
+      expect(typeof span.agent).toBe('function');
+      expect(typeof span.tool).toBe('function');
+      expect(typeof span.chain).toBe('function');
+      expect(typeof span.retriever).toBe('function');
+      expect(typeof span.guardrail).toBe('function');
+      expect(typeof span.evaluator).toBe('function');
+      expect(typeof span.embedding).toBe('function');
+      expect(typeof span.generation).toBe('function');
+      expect(typeof span.event).toBe('function');
+      expect(typeof span.score).toBe('function');
+    });
+
+    it('should emit correct event types for each observation method on TraceClient', () => {
+      const events: any[] = [];
+      const enqueue = (event: any) => events.push(event);
+      const trace = new TraceClient('trace-456', enqueue);
+
+      trace.span({ name: 'span' });
+      trace.agent({ name: 'agent' });
+      trace.tool({ name: 'tool' });
+      trace.chain({ name: 'chain' });
+      trace.retriever({ name: 'retriever' });
+      trace.guardrail({ name: 'guardrail' });
+      trace.evaluator({ name: 'evaluator' });
+      trace.embedding({ name: 'embedding' });
+      trace.generation({ name: 'generation' });
+      trace.event({ name: 'event' });
+
+      const types = events.map((e) => e.type);
+      expect(types).toContain('span-create');
+      expect(types).toContain('agent-create');
+      expect(types).toContain('tool-create');
+      expect(types).toContain('chain-create');
+      expect(types).toContain('retriever-create');
+      expect(types).toContain('guardrail-create');
+      expect(types).toContain('evaluator-create');
+      expect(types).toContain('embedding-create');
+      expect(types).toContain('generation-create');
+      expect(types).toContain('event-create');
+    });
+
+    it('should emit correct event types for each observation method on SpanClient', () => {
+      const events: any[] = [];
+      const enqueue = (event: any) => events.push(event);
+      const span = new SpanClient('span-789', 'trace-789', enqueue);
+
+      span.span({ name: 'child-span' });
+      span.agent({ name: 'child-agent' });
+      span.tool({ name: 'child-tool' });
+      span.chain({ name: 'child-chain' });
+      span.retriever({ name: 'child-retriever' });
+      span.guardrail({ name: 'child-guardrail' });
+      span.evaluator({ name: 'child-evaluator' });
+      span.embedding({ name: 'child-embedding' });
+      span.generation({ name: 'child-generation' });
+      span.event({ name: 'child-event' });
+
+      const types = events.map((e) => e.type);
+      expect(types).toContain('span-create');
+      expect(types).toContain('agent-create');
+      expect(types).toContain('tool-create');
+      expect(types).toContain('chain-create');
+      expect(types).toContain('retriever-create');
+      expect(types).toContain('guardrail-create');
+      expect(types).toContain('evaluator-create');
+      expect(types).toContain('embedding-create');
+      expect(types).toContain('generation-create');
+      expect(types).toContain('event-create');
+    });
+
+    it('SpanClient children should set parentObservationId automatically', () => {
+      const events: any[] = [];
+      const enqueue = (event: any) => events.push(event);
+      const span = new SpanClient('parent-span', 'trace-abc', enqueue);
+
+      span.guardrail({ name: 'guardrail-check' });
+      span.evaluator({ name: 'eval-run' });
+      span.embedding({ name: 'embed-op' });
+
+      for (const event of events) {
+        expect(event.body.parentObservationId).toBe('parent-span');
+        expect(event.body.traceId).toBe('trace-abc');
+      }
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 15. Prompt & Dataset SDK methods
+  // -----------------------------------------------------------------------
+
+  describe('prompt and dataset SDK methods', () => {
+    it('core client should expose getPrompt method', () => {
+      const { client } = createTestClient();
+      expect(typeof client.getPrompt).toBe('function');
+    });
+
+    it('core client should expose createPrompt method', () => {
+      const { client } = createTestClient();
+      expect(typeof client.createPrompt).toBe('function');
+    });
+
+    it('core client should expose getDataset method', () => {
+      const { client } = createTestClient();
+      expect(typeof client.getDataset).toBe('function');
+    });
+
+    it('core client should expose createDataset method', () => {
+      const { client } = createTestClient();
+      expect(typeof client.createDataset).toBe('function');
+    });
+
+    it('core client should expose createDatasetItem method', () => {
+      const { client } = createTestClient();
+      expect(typeof client.createDatasetItem).toBe('function');
+    });
+
+    it('getPrompt should return parsed response from server', async () => {
+      const { client } = createTestClient();
+      // Mock fetchWithRetry returns { ok: true, json: async () => ({}) }
+      const result = await client.getPrompt({ name: 'my-prompt' });
+      expect(result).toBeDefined();
+      expect(result).toEqual({});
+    });
+
+    it('getDataset should return parsed response from server', async () => {
+      const { client } = createTestClient();
+      const result = await client.getDataset({ name: 'my-dataset' });
+      expect(result).toBeDefined();
+      expect(result).toEqual({});
+    });
+
+    it('prompt/dataset methods should return null when disabled', async () => {
+      const { client } = createTestClient({ enabled: false });
+      expect(await client.getPrompt({ name: 'test' })).toBeNull();
+      expect(
+        await client.createPrompt({ name: 'test', prompt: 'hello' })
+      ).toBeNull();
+      expect(await client.getDataset({ name: 'test' })).toBeNull();
+      expect(await client.createDataset({ name: 'test' })).toBeNull();
+      expect(
+        await client.createDatasetItem({ datasetName: 'test', input: 'x' })
+      ).toBeNull();
+    });
+  });
+
+  // -----------------------------------------------------------------------
+  // 16. Node SDK: sampling + PII masking propagation
+  // -----------------------------------------------------------------------
+
+  describe('node SDK config propagation', () => {
+    const { Observability } = require('@illuma-ai/observability-node');
+
+    it('Observability class should accept sampleRate option', async () => {
+      expect(Observability).toBeDefined();
+
+      const obs = new Observability({
+        publicKey: 'pk-test',
+        secretKey: 'sk-test',
+        baseUrl: 'http://localhost:9999',
+        sampleRate: 0.5,
+        flushInterval: 0,
+      });
+
+      expect(obs).toBeDefined();
+      await obs.shutdown();
+    });
+
+    it('Observability class should accept maskFunction option', async () => {
+      const obs = new Observability({
+        publicKey: 'pk-test',
+        secretKey: 'sk-test',
+        baseUrl: 'http://localhost:9999',
+        maskFunction: (body: Record<string, unknown>) => ({
+          ...body,
+          input: '[REDACTED]',
+        }),
+        flushInterval: 0,
+      });
+
+      expect(obs).toBeDefined();
+      await obs.shutdown();
+    });
+  });
+});
+
+// ---------------------------------------------------------------------------
+// 17. Guardrail tracing via ObservabilityCallbackHandler
+// ---------------------------------------------------------------------------
+
+/** Re-use MockObservabilityClient shape for guardrail tests (same file, top-level scope) */
+class GuardrailMockClient {
+  events: {
+    id: string;
+    type: string;
+    timestamp: string;
+    body: Record<string, unknown>;
+  }[] = [];
+  flushed = false;
+  shutdownCalled = false;
+  enqueue(event: {
+    id: string;
+    type: string;
+    timestamp: string;
+    body: Record<string, unknown>;
+  }): void {
+    this.events.push(event);
+  }
+  async flush(): Promise<void> {
+    this.flushed = true;
+  }
+  async shutdown(): Promise<void> {
+    this.shutdownCalled = true;
+  }
+  getByType(type: string) {
+    return this.events.filter((e) => e.type === type);
+  }
+  getEventTypes() {
+    return this.events.map((e) => e.type);
+  }
+  printEvents(): void {
+    for (const event of this.events) {
+      const body = event.body;
+      const usage = body.usage as Record<string, unknown> | undefined;
+      console.log(
+        `  ${event.type} | name=${body.name ?? '—'} | id=${body.id ?? '—'} | traceId=${body.traceId ?? '—'} | parentObsId=${body.parentObservationId ?? '—'}` +
+          (usage
+            ? ` | usage={prompt:${usage.promptTokens ?? '—'},completion:${usage.completionTokens ?? '—'},total:${usage.totalTokens ?? '—'}}`
+            : '') +
+          (body.model ? ` | model=${body.model}` : '') +
+          (body.level === 'ERROR' || body.level === 'WARNING'
+            ? ` | ${body.level}: ${body.statusMessage}`
+            : '')
+      );
+    }
+  }
+}
+
+const grSerialized = (name: string): Serialized => ({
+  lc: 1,
+  type: 'not_implemented',
+  id: ['langchain', name],
+});
+
+describe('Guardrail tracing', () => {
+  let grClient: GuardrailMockClient;
+  let grHandler: any; // Use any to bypass ts-jest type resolution for traceGuardrail
+
+  beforeEach(() => {
+    grClient = new GuardrailMockClient();
+    grHandler = new ObservabilityCallbackHandler({
+      client: grClient as any,
+    });
+  });
+
+  it('outcome=passed should create guardrail-create + span-update with DEFAULT level', async () => {
+    const chainRunId = makeUUID();
+    await grHandler.handleChainStart(
+      grSerialized('RunnableSequence'),
+      { input: 'hello' },
+      chainRunId
+    );
+
+    const result = grHandler.traceGuardrail({
+      name: 'Output Moderation',
+      guardrailId: 'bedrock-guardrail-123',
+      guardrailVersion: '1',
+      outcome: 'passed',
+      actionApplied: false,
+      action: 'NONE',
+      reason: 'passed',
+      source: 'OUTPUT',
+      input: 'What is the weather?',
+      output: 'The weather is sunny.',
+      violations: [],
+      assessments: [{ contentPolicy: { filters: [] } }],
+    });
+
+    expect(result).not.toBeNull();
+
+    const guardrailEvents = grClient.getByType('guardrail-create');
+    expect(guardrailEvents).toHaveLength(1);
+
+    const grBody = guardrailEvents[0].body;
+    expect(grBody.name).toBe('Output Moderation');
+    expect(grBody.input).toBe('What is the weather?');
+    expect((grBody.metadata as Record<string, unknown>).guardrailId).toBe(
+      'bedrock-guardrail-123'
+    );
+    expect((grBody.metadata as Record<string, unknown>).outcome).toBe('passed');
+    expect((grBody.metadata as Record<string, unknown>).source).toBe('OUTPUT');
+
+    // span-update should show PASSED
+    const spanUpdates = grClient.getByType('span-update');
+    const guardrailUpdate = spanUpdates.find((e: any) => e.body.id === result);
+    expect(guardrailUpdate).toBeDefined();
+    expect(guardrailUpdate!.body.output).toBe('The weather is sunny.');
+    expect(guardrailUpdate!.body.level).toBe('DEFAULT');
+    expect(guardrailUpdate!.body.statusMessage).toBe('PASSED');
+  });
+
+  it('outcome=blocked should mark with ERROR level when enforced', async () => {
+    const chainRunId = makeUUID();
+    await grHandler.handleChainStart(
+      grSerialized('RunnableSequence'),
+      { input: 'test' },
+      chainRunId
+    );
+
+    grHandler.traceGuardrail({
+      name: 'Input Moderation',
+      guardrailId: 'bedrock-guardrail-456',
+      outcome: 'blocked',
+      actionApplied: true,
+      action: 'GUARDRAIL_INTERVENED',
+      reason: 'policy_violation',
+      source: 'INPUT',
+      input: 'How to hack a system?',
+      output: 'Sorry, I cannot help with that.',
+      violations: [
+        { type: 'CONTENT_POLICY', category: 'VIOLENCE', action: 'BLOCKED' },
+      ],
+    });
+
+    const spanUpdates = grClient.getByType('span-update');
+    const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
+    expect(guardrailUpdate.body.level).toBe('ERROR');
+    expect(guardrailUpdate.body.statusMessage).toBe(
+      'BLOCKED: policy_violation'
+    );
+    expect(
+      (guardrailUpdate.body.metadata as Record<string, unknown>).outcome
+    ).toBe('blocked');
+    expect(
+      (guardrailUpdate.body.metadata as Record<string, unknown>).actionApplied
+    ).toBe(true);
+    expect(
+      (guardrailUpdate.body.metadata as Record<string, unknown>).violations
+    ).toHaveLength(1);
+  });
+
+  it('outcome=blocked with actionApplied=false should note "not enforced"', async () => {
+    const chainRunId = makeUUID();
+    await grHandler.handleChainStart(
+      grSerialized('RunnableSequence'),
+      { input: 'test' },
+      chainRunId
+    );
+
+    grHandler.traceGuardrail({
+      name: 'Input Moderation',
+      outcome: 'blocked',
+      actionApplied: false,
+      reason: 'policy_violation',
+      source: 'INPUT',
+      input: 'Some flagged content',
+    });
+
+    const spanUpdates = grClient.getByType('span-update');
+    const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
+    expect(guardrailUpdate.body.level).toBe('ERROR');
+    expect(guardrailUpdate.body.statusMessage).toContain('not enforced');
+    expect(
+      (guardrailUpdate.body.metadata as Record<string, unknown>).actionApplied
+    ).toBe(false);
+  });
+
+  it('outcome=anonymized should mark with WARNING and include modified content', async () => {
+    const chainRunId = makeUUID();
+    await grHandler.handleChainStart(
+      grSerialized('RunnableSequence'),
+      { input: 'test' },
+      chainRunId
+    );
+
+    grHandler.traceGuardrail({
+      name: 'Output Moderation',
+      outcome: 'anonymized',
+      actionApplied: true,
+      action: 'GUARDRAIL_INTERVENED',
+      reason: 'anonymized',
+      source: 'OUTPUT',
+      input: 'My email is john@example.com and SSN is 123-45-6789',
+      originalContent: 'My email is john@example.com and SSN is 123-45-6789',
+      modifiedContent: 'My email is [EMAIL] and SSN is [SSN]',
+      violations: [
+        { type: 'PII_POLICY', category: 'EMAIL', action: 'ANONYMIZED' },
+        { type: 'PII_POLICY', category: 'SSN', action: 'ANONYMIZED' },
+      ],
+    });
+
+    const spanUpdates = grClient.getByType('span-update');
+    const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
+    expect(guardrailUpdate.body.level).toBe('WARNING');
+    expect(guardrailUpdate.body.statusMessage).toBe(
+      'ANONYMIZED: PII detected and masked'
+    );
+    // output should be the modified content
+    expect(guardrailUpdate.body.output).toBe(
+      'My email is [EMAIL] and SSN is [SSN]'
+    );
+    const meta = guardrailUpdate.body.metadata as Record<string, unknown>;
+    expect(meta.originalContent).toBe(
+      'My email is john@example.com and SSN is 123-45-6789'
+    );
+    expect(meta.modifiedContent).toBe('My email is [EMAIL] and SSN is [SSN]');
+    expect(meta.violations).toHaveLength(2);
+  });
+
+  it('outcome=intervened should mark with WARNING level', async () => {
+    const chainRunId = makeUUID();
+    await grHandler.handleChainStart(
+      grSerialized('RunnableSequence'),
+      { input: 'test' },
+      chainRunId
+    );
+
+    grHandler.traceGuardrail({
+      name: 'Output Moderation',
+      outcome: 'intervened',
+      actionApplied: true,
+      action: 'GUARDRAIL_INTERVENED',
+      reason: 'intervened_passthrough',
+      source: 'OUTPUT',
+      input: 'Original response text',
+      modifiedContent: 'Modified response text',
+    });
+
+    const spanUpdates = grClient.getByType('span-update');
+    const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
+    expect(guardrailUpdate.body.level).toBe('WARNING');
+    expect(guardrailUpdate.body.statusMessage).toBe(
+      'INTERVENED: intervened_passthrough'
+    );
+  });
+
+  it('traceGuardrail() should return null when no trace exists', () => {
+    const freshClient = new GuardrailMockClient();
+    const freshHandler: any = new ObservabilityCallbackHandler({
+      client: freshClient as any,
+    });
+
+    const result = freshHandler.traceGuardrail({
+      name: 'Pre-trace Guardrail',
+      outcome: 'passed',
+      actionApplied: false,
+      source: 'INPUT',
+    });
+
+    expect(result).toBeNull();
+    expect(freshClient.getByType('guardrail-create')).toHaveLength(0);
+  });
+
+  it('traceGuardrail() should use custom startTime and endTime', async () => {
+    const chainRunId = makeUUID();
+    await grHandler.handleChainStart(
+      grSerialized('Chain'),
+      { input: 'test' },
+      chainRunId
+    );
+
+    const startTime = '2024-01-01T00:00:00.000Z';
+    const endTime = '2024-01-01T00:00:01.500Z';
+
+    grHandler.traceGuardrail({
+      name: 'Timed Guardrail',
+      outcome: 'passed',
+      actionApplied: false,
+      source: 'OUTPUT',
+      startTime,
+      endTime,
+    });
+
+    const guardrailEvents = grClient.getByType('guardrail-create');
+    expect(guardrailEvents[0].body.startTime).toBe(startTime);
+
+    const spanUpdates = grClient.getByType('span-update');
+    const lastUpdate = spanUpdates[spanUpdates.length - 1];
+    expect(lastUpdate.body.endTime).toBe(endTime);
+  });
+
+  it('should trace both input and output guardrails on the same trace', async () => {
+    const chainRunId = makeUUID();
+    await grHandler.handleChainStart(
+      grSerialized('RunnableSequence'),
+      { input: 'test' },
+      chainRunId
+    );
+
+    const inputId = grHandler.traceGuardrail({
+      name: 'Input Moderation',
+      outcome: 'passed',
+      actionApplied: false,
+      source: 'INPUT',
+      input: 'What is 2+2?',
+    });
+
+    const outputId = grHandler.traceGuardrail({
+      name: 'Output Moderation',
+      outcome: 'passed',
+      actionApplied: false,
+      source: 'OUTPUT',
+      input: 'The answer is 4.',
+    });
+
+    expect(inputId).not.toBeNull();
+    expect(outputId).not.toBeNull();
+    expect(inputId).not.toBe(outputId);
+
+    const guardrailEvents = grClient.getByType('guardrail-create');
+    expect(guardrailEvents).toHaveLength(2);
+
+    const traceId = guardrailEvents[0].body.traceId;
+    expect(guardrailEvents[1].body.traceId).toBe(traceId);
+
+    expect(
+      (guardrailEvents[0].body.metadata as Record<string, unknown>).source
+    ).toBe('INPUT');
+    expect(
+      (guardrailEvents[1].body.metadata as Record<string, unknown>).source
+    ).toBe('OUTPUT');
+  });
+
+  it('should include full AWS Bedrock assessments and violations in metadata', async () => {
+    const chainRunId = makeUUID();
+    await grHandler.handleChainStart(
+      grSerialized('Chain'),
+      { input: 'test' },
+      chainRunId
+    );
+
+    grHandler.traceGuardrail({
+      name: 'Detailed Guardrail',
+      outcome: 'blocked',
+      actionApplied: true,
+      action: 'GUARDRAIL_INTERVENED',
+      reason: 'policy_violation',
+      source: 'INPUT',
+      violations: [
+        {
+          type: 'CONTENT_POLICY',
+          category: 'VIOLENCE',
+          confidence: 'HIGH',
+          action: 'BLOCKED',
+        },
+        { type: 'PII_POLICY', category: 'SSN', action: 'BLOCKED' },
+      ],
+      assessments: [
+        {
+          contentPolicy: {
+            filters: [
+              { type: 'VIOLENCE', confidence: 'HIGH', action: 'BLOCKED' },
+            ],
+          },
+          sensitiveInformationPolicy: {
+            piiEntities: [{ type: 'SSN', action: 'BLOCKED' }],
+          },
+        },
+      ],
+    });
+
+    const spanUpdates = grClient.getByType('span-update');
+    const lastUpdate = spanUpdates[spanUpdates.length - 1];
+    const meta = lastUpdate.body.metadata as Record<string, unknown>;
+
+    expect(meta.violations).toHaveLength(2);
+    expect(meta.assessments).toHaveLength(1);
+    expect((meta.violations as any[])[0]).toEqual({
+      type: 'CONTENT_POLICY',
+      category: 'VIOLENCE',
+      confidence: 'HIGH',
+      action: 'BLOCKED',
+    });
+  });
+});
+
+// ---------------------------------------------------------------------------
+// 18. Guardrail tracing with debug logging — full pipeline with all 4 outcomes
+// ---------------------------------------------------------------------------
+
+describe('Guardrail tracing with debug output', () => {
+  it('should produce debug-friendly trace output for manual verification', async () => {
+    const debugClient = new GuardrailMockClient();
+    const debugHandler: any = new ObservabilityCallbackHandler({
+      client: debugClient as any,
+      traceName: 'guardrail-debug-test',
+      userId: 'user-123',
+      sessionId: 'session-456',
+      debug: true,
+    });
+
+    // 1. Chain starts (creates trace)
+    const chainRunId = makeUUID();
+    await debugHandler.handleChainStart(
+      grSerialized('RunnableSequence'),
+      { messages: [{ role: 'user', content: 'Explain quantum computing' }] },
+      chainRunId
+    );
+
+    // 2. LLM call
+    const llmRunId = makeUUID();
+    await debugHandler.handleLLMStart(
+      grSerialized('ChatOpenAI'),
+      ['Explain quantum computing'],
+      llmRunId,
+      chainRunId
+    );
+    const llmResult: LLMResult = {
+      generations: [
+        [{ text: 'Quantum computing uses qubits...', generationInfo: {} }],
+      ],
+      llmOutput: {
+        tokenUsage: { promptTokens: 10, completionTokens: 25, totalTokens: 35 },
+        modelName: 'gpt-4o',
+      },
+    };
+    await debugHandler.handleLLMEnd(llmResult, llmRunId);
+    await debugHandler.handleChainEnd(
+      { output: 'Quantum computing uses qubits...' },
+      chainRunId
+    );
+
+    // 3. Input guardrail — passed
+    const inputGuardrailId = debugHandler.traceGuardrail({
+      name: 'Input Moderation (Bedrock)',
+      guardrailId: 'arn:aws:bedrock:us-east-1:123456:guardrail/abc123',
+      guardrailVersion: '3',
+      outcome: 'passed',
+      actionApplied: false,
+      action: 'NONE',
+      reason: 'passed',
+      source: 'INPUT',
+      input: 'Explain quantum computing',
+      startTime: '2024-06-15T10:00:00.000Z',
+      endTime: '2024-06-15T10:00:00.150Z',
+    });
+
+    // 4. Output guardrail — passed
+    const outputGuardrailId = debugHandler.traceGuardrail({
+      name: 'Output Moderation (Bedrock)',
+      guardrailId: 'arn:aws:bedrock:us-east-1:123456:guardrail/abc123',
+      guardrailVersion: '3',
+      outcome: 'passed',
+      actionApplied: false,
+      action: 'NONE',
+      reason: 'passed',
+      source: 'OUTPUT',
+      input: 'Quantum computing uses qubits...',
+      violations: [],
+      assessments: [
+        { topicPolicy: { topics: [] }, contentPolicy: { filters: [] } },
+      ],
+      startTime: '2024-06-15T10:00:01.000Z',
+      endTime: '2024-06-15T10:00:01.200Z',
+    });
+
+    expect(inputGuardrailId).not.toBeNull();
+    expect(outputGuardrailId).not.toBeNull();
+
+    console.log('\n=== SAMPLE GUARDRAIL TRACE OUTPUT ===');
+    debugClient.printEvents();
+    console.log('=== END TRACE ===\n');
+
+    // Verify complete event sequence
+    const types = debugClient.getEventTypes();
+    expect(types).toEqual([
+      'trace-create',
+      'chain-create',
+      'generation-create',
+      'generation-update',
+      'span-update', // chain end
+      'guardrail-create', // input guardrail
+      'span-update', // input guardrail update
+      'guardrail-create', // output guardrail
+      'span-update', // output guardrail update
+    ]);
+
+    // Verify trace ID consistency
+    const traceId = debugClient.getByType('trace-create')[0].body.id;
+    for (const evt of debugClient.events) {
+      const body = evt.body;
+      if (evt.type !== 'trace-create') {
+        expect(body.traceId).toBe(traceId);
+      }
+    }
+  });
+
+  it('should trace all 4 guardrail outcomes with correct levels', async () => {
+    const debugClient = new GuardrailMockClient();
+    const handler: any = new ObservabilityCallbackHandler({
+      client: debugClient as any,
+      debug: true,
+    });
+
+    // Create trace
+    const chainRunId = makeUUID();
+    await handler.handleChainStart(
+      grSerialized('Chain'),
+      { input: 'test' },
+      chainRunId
+    );
+
+    // All 4 outcomes
+    handler.traceGuardrail({
+      name: 'Passed',
+      outcome: 'passed',
+      actionApplied: false,
+      source: 'INPUT',
+    });
+    handler.traceGuardrail({
+      name: 'Blocked',
+      outcome: 'blocked',
+      actionApplied: true,
+      reason: 'policy_violation',
+      source: 'INPUT',
+    });
+    handler.traceGuardrail({
+      name: 'Anonymized',
+      outcome: 'anonymized',
+      actionApplied: true,
+      reason: 'anonymized',
+      source: 'OUTPUT',
+      modifiedContent: '[EMAIL]',
+    });
+    handler.traceGuardrail({
+      name: 'Intervened',
+      outcome: 'intervened',
+      actionApplied: true,
+      reason: 'intervened_passthrough',
+      source: 'OUTPUT',
+    });
+
+    console.log('\n=== ALL 4 GUARDRAIL OUTCOMES ===');
+    debugClient.printEvents();
+    console.log('=== END ===\n');
+
+    const guardrailEvents = debugClient.getByType('guardrail-create');
+    expect(guardrailEvents).toHaveLength(4);
+
+    const spanUpdates = debugClient.getByType('span-update');
+    // 4 guardrail updates (chain not ended in this test)
+    expect(spanUpdates).toHaveLength(4);
+
+    // Check levels: passed=DEFAULT, blocked=ERROR, anonymized=WARNING, intervened=WARNING
+    const guardrailUpdates = spanUpdates;
+    expect(guardrailUpdates[0].body.level).toBe('DEFAULT');
+    expect(guardrailUpdates[1].body.level).toBe('ERROR');
+    expect(guardrailUpdates[2].body.level).toBe('WARNING');
+    expect(guardrailUpdates[3].body.level).toBe('WARNING');
+  });
+});