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

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

const {
  forceFlush,
  shutdown,
  sdkStart,
  sdkShutdown,
  generationUpdate,
  generationEnd,
  childStartObservation,
  startObservation,
  propagateAttributes,
} = 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),
    propagateAttributes: vi.fn().mockImplementation((_attrs: unknown, fn: () => unknown) => fn()),
  };
});

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, propagateAttributes }));

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: { promptTokens: 10, completionTokens: 5, totalTokens: 15 },
    });

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

  it('omits usageDetails entirely when no usage is reported', () => {
    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' });

    expect(generationUpdate).toHaveBeenCalledWith(expect.objectContaining({ usageDetails: undefined }));
  });

  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);
  });

  it('runWithSession propagates sessionId/tags to the callback via propagateAttributes and returns its result', async () => {
    const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
    const fn = vi.fn().mockResolvedValue('done');

    const result = await tracer.runWithSession({ sessionId: 'session-1', tags: ['cron'] }, fn);

    expect(propagateAttributes).toHaveBeenCalledWith({ sessionId: 'session-1', tags: ['cron'] }, fn);
    expect(result).toBe('done');
  });
});