sugarr/test/core/llm/LlmPipeline.test.ts

import { describe, expect, it, vi } from 'vitest';
import { LlmPipeline } from '../../../src/core/llm/LlmPipeline';
import type { RunContext } from '../../../src/core/llm/RunContext';
import type { Tracer, TraceHandle } from '../../../src/core/llm/Tracer';
import type { LlmTool } from '../../../src/core/llm/LlmTool';
import type { Logger } from '../../../src/core/logging/logger';

class SumPipeline extends LlmPipeline<number[], number> {
  protected readonly name = 'sum-pipeline';
  protected async runInternal(input: number[], ctx: RunContext): Promise<number> {
    expect(ctx.trace.id).toBe('trace-1');
    return input.reduce((a, b) => a + b, 0);
  }
}

class ThrowingPipeline extends LlmPipeline<void, void> {
  protected readonly name = 'throwing-pipeline';
  protected async runInternal(): Promise<void> {
    throw new Error('pipeline failed');
  }
}

function makeTracer(trace: TraceHandle) {
  const flush = vi.fn().mockResolvedValue(undefined);
  const startTrace = vi.fn().mockReturnValue(trace);
  const runWithSession = vi.fn().mockImplementation((_session: unknown, fn: () => unknown) => fn());
  const tracer = { startTrace, flush, runWithSession } as unknown as Tracer;
  return { tracer, startTrace, flush, runWithSession };
}

describe('LlmPipeline', () => {
  it('runs the whole call inside a session scoped to a fresh session id, opens one trace, runs runInternal, and flushes afterward', async () => {
    const trace = { id: 'trace-1', startGeneration: vi.fn() } as unknown as TraceHandle;
    const { tracer, startTrace, flush, runWithSession } = makeTracer(trace);
    const logger = {} as unknown as Logger;

    const pipeline = new SumPipeline(tracer, logger);
    const result = await pipeline.runAsync([1, 2, 3]);

    expect(result).toBe(6);
    expect(runWithSession).toHaveBeenCalledWith(
      expect.objectContaining({ sessionId: expect.any(String), tags: undefined }),
      expect.any(Function)
    );
    expect(startTrace).toHaveBeenCalledWith('sum-pipeline', { trigger: undefined });
    expect(flush).toHaveBeenCalledTimes(1);
  });

  it('tags the session with the given trigger', async () => {
    const trace = { id: 'trace-1', startGeneration: vi.fn() } as unknown as TraceHandle;
    const { tracer, startTrace, runWithSession } = makeTracer(trace);
    const logger = {} as unknown as Logger;

    const pipeline = new SumPipeline(tracer, logger);
    await pipeline.runAsync([1, 2, 3], { trigger: 'cron' });

    expect(runWithSession).toHaveBeenCalledWith(
      expect.objectContaining({ tags: ['cron'] }),
      expect.any(Function)
    );
    expect(startTrace).toHaveBeenCalledWith('sum-pipeline', { trigger: 'cron' });
  });

  it('still flushes the trace when runInternal throws', async () => {
    const trace = { id: 'trace-2', startGeneration: vi.fn() } as unknown as TraceHandle;
    const { tracer, flush } = makeTracer(trace);
    const logger = {} as unknown as Logger;

    const pipeline = new ThrowingPipeline(tracer, logger);
    await expect(pipeline.runAsync()).rejects.toThrow('pipeline failed');
    expect(flush).toHaveBeenCalledTimes(1);
  });

  it('runTool delegates to the given tool.runAsync with the same context', async () => {
    const trace = { id: 'trace-3', startGeneration: vi.fn() } as unknown as TraceHandle;
    const { tracer } = makeTracer(trace);
    const logger = {} as unknown as Logger;

    class ToolRunnerPipeline extends LlmPipeline<void, string> {
      protected readonly name = 'tool-runner';
      protected async runInternal(_input: void, ctx: RunContext): Promise<string> {
        const tool = { runAsync: vi.fn().mockResolvedValue('tool-output') } as unknown as LlmTool<
          Record<string, never>,
          string
        >;
        return this.runTool(tool, {}, ctx);
      }
    }

    const pipeline = new ToolRunnerPipeline(tracer, logger);
    const result = await pipeline.runAsync();
    expect(result).toBe('tool-output');
  });
});