sugarr/test/core/langfuse/LangfuseTracer.test.ts

import { describe, expect, it, vi } from 'vitest';

const { forceFlush, shutdown, sdkStart, sdkShutdown, generationUpdate, generationEnd, childStartObservation, startObservation } =
  vi.hoisted(() => {
    const generationUpdate = vi.fn();
    const generationEnd = vi.fn();
    const generation = { update: generationUpdate, end: generationEnd };
    const childStartObservation = vi.fn().mockReturnValue(generation);
    const rootSpan = { traceId: 'lf-trace-1', startObservation: childStartObservation };
    return {
      forceFlush: vi.fn().mockResolvedValue(undefined),
      shutdown: vi.fn().mockResolvedValue(undefined),
      sdkStart: vi.fn(),
      sdkShutdown: vi.fn().mockResolvedValue(undefined),
      generationUpdate,
      generationEnd,
      childStartObservation,
      startObservation: vi.fn().mockReturnValue(rootSpan),
    };
  });

vi.mock('@langfuse/otel', () => ({
  LangfuseSpanProcessor: vi.fn().mockImplementation(() => ({ forceFlush, shutdown })),
}));

vi.mock('@opentelemetry/sdk-node', () => ({
  NodeSDK: vi.fn().mockImplementation(() => ({ start: sdkStart, shutdown: sdkShutdown })),
}));

vi.mock('@langfuse/tracing', () => ({ startObservation }));

import { LangfuseTracer } from '../../../src/core/langfuse/LangfuseTracer';

describe('LangfuseTracer', () => {
  it('starts the OTEL SDK with a LangfuseSpanProcessor on construction', () => {
    new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    expect(sdkStart).toHaveBeenCalled();
  });

  it('starts a root observation and exposes its trace id', () => {
    const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    const handle = tracer.startTrace('my-pipeline', { foo: 'bar' });

    expect(startObservation).toHaveBeenCalledWith('my-pipeline', { metadata: { foo: 'bar' } });
    expect(handle.id).toBe('lf-trace-1');
  });

  it('startGeneration creates a child generation observation and end() forwards output/usage', () => {
    const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    const handle = tracer.startTrace('my-pipeline');
    const gen = handle.startGeneration({ name: 'step-1', model: 'llama3', input: ['hi'] });
    gen.end({ output: 'done', usage: { totalTokens: 5 } });

    expect(childStartObservation).toHaveBeenCalledWith(
      'step-1',
      expect.objectContaining({ model: 'llama3', input: ['hi'] }),
      { asType: 'generation' }
    );
    expect(generationUpdate).toHaveBeenCalledWith(
      expect.objectContaining({ output: 'done', usageDetails: { totalTokens: 5 } })
    );
    expect(generationEnd).toHaveBeenCalledTimes(1);
  });

  it('flush() force-flushes the span processor without shutting down the SDK', async () => {
    const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    await tracer.flush();
    expect(forceFlush).toHaveBeenCalledTimes(1);
    expect(sdkShutdown).not.toHaveBeenCalled();
  });

  it('shutdown() shuts down the underlying OTEL SDK', async () => {
    const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    await tracer.shutdown();
    expect(sdkShutdown).toHaveBeenCalledTimes(1);
  });
});