upgrade langfuse

test/core/langfuse/LangfuseTracer.test.ts

@@ -1,42 +1,77 @@
 import { describe, expect, it, vi } from 'vitest';
 
-const generationEnd = vi.fn();
-const generation = vi.fn().mockReturnValue({ end: generationEnd });
-const trace = vi.fn().mockReturnValue({ id: 'lf-trace-1', generation });
-const flushAsync = vi.fn().mockResolvedValue(undefined);
+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', () => ({
-  Langfuse: vi.fn().mockImplementation(() => ({ trace, flushAsync })),
+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 a Langfuse trace and exposes its id', () => {
+  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(trace).toHaveBeenCalledWith({ name: 'my-pipeline', metadata: { foo: 'bar' } });
+    expect(startObservation).toHaveBeenCalledWith('my-pipeline', { metadata: { foo: 'bar' } });
     expect(handle.id).toBe('lf-trace-1');
   });
 
-  it('startGeneration creates a Langfuse generation and end() forwards output/usage', () => {
+  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(generation).toHaveBeenCalledWith(
-      expect.objectContaining({ name: 'step-1', model: 'llama3', input: ['hi'] })
+    expect(childStartObservation).toHaveBeenCalledWith(
+      'step-1',
+      expect.objectContaining({ model: 'llama3', input: ['hi'] }),
+      { asType: 'generation' }
     );
-    expect(generationEnd).toHaveBeenCalledWith(
-      expect.objectContaining({ output: 'done', usage: { totalTokens: 5 } })
+    expect(generationUpdate).toHaveBeenCalledWith(
+      expect.objectContaining({ output: 'done', usageDetails: { totalTokens: 5 } })
     );
+    expect(generationEnd).toHaveBeenCalledTimes(1);
   });
 
-  it('flush() calls flushAsync on the underlying client', async () => {
+  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(flushAsync).toHaveBeenCalledTimes(1);
+    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);
   });
 });