insights-js/packages/ai/tests/processor.test.ts

import type { SpanProcessor, ReadableSpan, Span } from '@opentelemetry/sdk-trace-base'
import type { Context } from '@opentelemetry/api'
import { InsightsSpanProcessor } from '../src/otel'
import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-http'
import { BatchSpanProcessor } from '@opentelemetry/sdk-trace-base'

jest.mock('@opentelemetry/exporter-trace-otlp-http', () => ({
  OTLPTraceExporter: jest.fn(),
}))

jest.mock('@opentelemetry/sdk-trace-base', () => ({
  BatchSpanProcessor: jest.fn().mockImplementation(() => ({
    onStart: jest.fn(),
    onEnd: jest.fn(),
    shutdown: jest.fn().mockResolvedValue(undefined),
    forceFlush: jest.fn().mockResolvedValue(undefined),
  })),
}))

function mockProcessor(): SpanProcessor & { onStart: jest.Mock; onEnd: jest.Mock } {
  return {
    onStart: jest.fn(),
    onEnd: jest.fn(),
    shutdown: jest.fn().mockResolvedValue(undefined),
    forceFlush: jest.fn().mockResolvedValue(undefined),
  }
}

function makeSpan(name: string, attributes: Record<string, unknown> = {}): ReadableSpan {
  return { name, attributes } as unknown as ReadableSpan
}

describe('InsightsSpanProcessor', () => {
  beforeEach(() => {
    jest.clearAllMocks()
  })

  it.each([
    {
      name: 'default host',
      apiKey: 'phc_test123',
      host: undefined,
      expectedUrl: 'https://us.i.insights.hanzo.ai/i/v0/ai/otel',
      expectedApiKey: 'phc_test123',
    },
    {
      name: 'custom host',
      apiKey: 'phc_test456',
      host: 'https://eu.i.insights.hanzo.ai',
      expectedUrl: 'https://eu.i.insights.hanzo.ai/i/v0/ai/otel',
      expectedApiKey: 'phc_test456',
    },
    {
      name: 'trailing slash',
      apiKey: 'phc_test789',
      host: 'https://custom.insights.hanzo.ai/',
      expectedUrl: 'https://custom.insights.hanzo.ai/i/v0/ai/otel',
      expectedApiKey: 'phc_test789',
    },
    {
      name: 'trimmed whitespace-sensitive values',
      apiKey: '  phc_test999\t ',
      host: '  https://custom.insights.hanzo.ai/\n',
      expectedUrl: 'https://custom.insights.hanzo.ai/i/v0/ai/otel',
      expectedApiKey: 'phc_test999',
    },
  ])('configures the OTLP exporter correctly with $name', ({ apiKey, host, expectedUrl, expectedApiKey }) => {
    new InsightsSpanProcessor({ apiKey, host })

    expect(OTLPTraceExporter).toHaveBeenCalledWith({
      url: expectedUrl,
      headers: { Authorization: `Bearer ${expectedApiKey}` },
    })
    expect(BatchSpanProcessor).toHaveBeenCalledWith(expect.any(Object))
  })

  it.each([
    ['missing', {}],
    ['empty', { apiKey: '' }],
    ['blank', { apiKey: '  \n\t ' }],
  ])('disables and no-ops when apiKey is %s', async (_case, options) => {
    const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => {})
    const processor = new InsightsSpanProcessor(options as any)

    processor.onStart({} as Span, {} as Context)
    processor.onEnd(makeSpan('gen_ai.chat'))
    await expect(processor.shutdown()).resolves.toBeUndefined()
    await expect(processor.forceFlush()).resolves.toBeUndefined()

    expect(OTLPTraceExporter).not.toHaveBeenCalled()
    expect(BatchSpanProcessor).not.toHaveBeenCalled()
    expect(warnSpy).toHaveBeenCalledWith(
      '[InsightsSpanProcessor] apiKey is missing or blank; the processor will be disabled.'
    )
    warnSpy.mockRestore()
  })

  it('delegates onStart to the inner processor', () => {
    const inner = mockProcessor()
    const processor = new InsightsSpanProcessor({ apiKey: 'phc_test', _spanProcessor: inner })

    const span = {} as Span
    const ctx = {} as Context
    processor.onStart(span, ctx)

    expect(inner.onStart).toHaveBeenCalledWith(span, ctx)
  })

  it('delegates shutdown', async () => {
    const inner = mockProcessor()
    const processor = new InsightsSpanProcessor({ apiKey: 'phc_test', _spanProcessor: inner })

    await processor.shutdown()
    expect(inner.shutdown).toHaveBeenCalled()
  })

  it('delegates forceFlush', async () => {
    const inner = mockProcessor()
    const processor = new InsightsSpanProcessor({ apiKey: 'phc_test', _spanProcessor: inner })

    await processor.forceFlush()
    expect(inner.forceFlush).toHaveBeenCalled()
  })
})

describe('InsightsSpanProcessor AI span filtering', () => {
  it('forwards spans with AI name prefixes', () => {
    const inner = mockProcessor()
    const processor = new InsightsSpanProcessor({ apiKey: 'phc_test', _spanProcessor: inner })

    processor.onEnd(makeSpan('gen_ai.chat'))
    processor.onEnd(makeSpan('llm.completion'))
    processor.onEnd(makeSpan('ai.invoke'))
    processor.onEnd(makeSpan('traceloop.workflow'))

    expect(inner.onEnd).toHaveBeenCalledTimes(4)
  })

  it('drops non-AI spans', () => {
    const inner = mockProcessor()
    const processor = new InsightsSpanProcessor({ apiKey: 'phc_test', _spanProcessor: inner })

    processor.onEnd(makeSpan('http.request'))
    processor.onEnd(makeSpan('db.query'))
    processor.onEnd(makeSpan('custom.operation'))

    expect(inner.onEnd).not.toHaveBeenCalled()
  })

  it('detects AI spans by attribute keys', () => {
    const inner = mockProcessor()
    const processor = new InsightsSpanProcessor({ apiKey: 'phc_test', _spanProcessor: inner })

    processor.onEnd(makeSpan('some.operation', { 'gen_ai.model': 'gpt-4' }))
    processor.onEnd(makeSpan('other.operation', { 'http.method': 'GET' }))

    expect(inner.onEnd).toHaveBeenCalledTimes(1)
  })
})