packages/ai/test/testing.test.ts

Rama dev · réplica de lectura del repositorio original.

import { describe, expect } from "bun:test"
import {
  ai,
  AIError,
  CompactionCheckpointResponse,
  CompactionPart,
  CompactionResponse,
  LanguageModel,
  LLMClient,
  LLMEvent,
  type LLMRequest,
  Message,
  ProviderID,
  RateLimitError,
} from "../src/effect.js"
import { OpenAI } from "../src/providers.js"
import { OpenAIChat } from "../src/protocols/openai-chat.js"
import { TestLLM } from "../src/testing.js"
import { Effect, Fiber, Latch, Stream } from "effect"
import { testEffect } from "./lib/effect.js"

const request = ai.request({
  model: LanguageModel.make({ id: "fictional-model", provider: "fixture", route: OpenAIChat.route }),
  prompt: "Say hello",
})
const it = testEffect(TestLLM.testLayer())

describe("TestLLM first-class client", () => {
  it.effect("returns scripted checkpoints only for trigger compaction", () =>
    Effect.gen(function* () {
      const client = yield* TestLLM.Test
      const request = ai.request({ model: OpenAI.configure().responses("fixture"), prompt: "hello" })
      const checkpoint = new CompactionCheckpointResponse({
        checkpoint: { type: "compaction", provider: ProviderID.make("openai"), encrypted: "opaque" },
        responseID: "resp_fixture",
      })
      const endpoint = new CompactionResponse({ replacement: [] })
      yield* client.push(checkpoint, endpoint, checkpoint, checkpoint)
      const operation = ai.compact(request, { mechanism: "trigger" })
      expect(yield* client.requests()).toEqual([])
      expect(yield* operation).toBe(checkpoint)
      expect(yield* client.compact(request, { mechanism: "trigger" }).pipe(Effect.catchDefect(Effect.succeed))).toBe(
        "TestLLM trigger compaction requires a CompactionCheckpointResponse",
      )
      expect(yield* client.compact(request).pipe(Effect.catchDefect(Effect.succeed))).toBe(
        "TestLLM compaction requires a CompactionResponse",
      )
      expect(yield* client.generate(request).pipe(Effect.catchDefect(Effect.succeed))).toBe(
        "TestLLM generation requires an event response",
      )
    }),
  )

  it.effect("rejects response fixtures for the wrong operation", () =>
    Effect.gen(function* () {
      const client = yield* TestLLM.Test
      const request = ai.request({
        model: OpenAI.configure({ apiKey: "test" }).responses("fixture"),
        prompt: "hello",
      })
      yield* client.push(TestLLM.stop(), new CompactionResponse({ replacement: [] }))
      expect(yield* client.compact(request).pipe(Effect.catchDefect(Effect.succeed))).toBe(
        "TestLLM compaction requires a CompactionResponse",
      )
      expect(yield* client.generate(request).pipe(Effect.catchDefect(Effect.succeed))).toBe(
        "TestLLM generation requires an event response",
      )
    }),
  )

  it.effect("scripts replacement windows with the same lazy recording, gates, and fallback controls", () =>
    Effect.gen(function* () {
      const client = yield* TestLLM.Test
      const request = ai.request({
        model: OpenAI.configure({ apiKey: "test" }).responses("fixture"),
        prompt: "hello",
      })
      const compacted = new CompactionResponse({
        replacement: [
          Message.user("retained input"),
          Message.assistant(CompactionPart.make({ provider: ProviderID.make("openai"), encrypted: "checkpoint" })),
          Message.user("retained tail"),
        ],
      })
      yield* client.push(compacted, TestLLM.text("continued", "answer"))
      const operation = ai.compact(request)
      expect(yield* client.requests()).toEqual([])
      const gate = yield* client.gate()
      const fiber = yield* operation.pipe(Effect.forkChild({ startImmediately: true }))
      yield* gate.started
      yield* client.wait(1)
      expect(fiber.pollUnsafe()).toBeUndefined()
      yield* gate.release
      expect(yield* Fiber.join(fiber)).toBe(compacted)
      const next = { ...request, messages: compacted.replacement }
      expect((yield* LLMClient.generate(next)).text).toBe("continued")
      yield* client.serve((observed) => {
        expect(observed).toEqual(next)
        return compacted
      })
      expect(yield* ai.compact(next)).toBe(compacted)
      yield* client.always(compacted)
      expect(yield* ai.compact(next)).toBe(compacted)
      expect(yield* client.requests()).toEqual([request, next, next, next])
    }),
  )

  it.effect("provides the same object under normal and test tags with snapshot observations", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      const client = yield* LLMClient.Service
      expect(client).toBe(llm)
      const before = yield* llm.requests()
      yield* llm.push(TestLLM.text("hello", "answer"))
      const generate = client.generate(request)
      client.stream(request)
      expect(yield* llm.requests()).toEqual([])

      expect((yield* generate).text).toBe("hello")
      expect(before).toEqual([])
      expect(yield* llm.requests()).toEqual([request])
      expect(yield* llm.requests()).not.toBe(yield* llm.requests())
    }),
  )

  it.effect("prioritizes queued replies over request-dependent and constant fallbacks", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      const served: LLMRequest[] = []
      yield* llm.always(TestLLM.text("old fallback", "old"))
      yield* llm.push(TestLLM.text("first", "first"), TestLLM.text("second", "second"))
      yield* llm.serve((request) => {
        served.push(request)
        return TestLLM.text(request.promptCacheKey ?? "default", "served")
      })

      expect((yield* LLMClient.generate(request)).text).toBe("first")
      expect((yield* LLMClient.generate(request)).text).toBe("second")
      expect(served).toEqual([])
      const selected = { ...request, promptCacheKey: "selected" }
      expect((yield* LLMClient.generate(selected)).text).toBe("selected")
      expect((yield* LLMClient.generate(request)).text).toBe("default")
      expect(served).toEqual([selected, request])

      yield* llm.push(TestLLM.text("queued again", "queued"))
      yield* llm.always(TestLLM.text("constant", "constant"))
      expect((yield* LLMClient.generate(request)).text).toBe("queued again")
      expect((yield* LLMClient.generate(request)).text).toBe("constant")
      expect((yield* LLMClient.generate(request)).text).toBe("constant")
      expect(served).toEqual([selected, request])
    }),
  )

  testEffect(
    TestLLM.testLayer({
      transformRequest: (request) => ({ ...request, promptCacheKey: "observation" }),
    }),
  ).effect("transforms observations without changing the request passed to the responder", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      yield* llm.serve((input) => {
        expect(input).toBe(request)
        return TestLLM.text("original", "answer")
      })
      const generate = llm.generate(request)
      expect(yield* llm.requests()).toEqual([])
      expect((yield* generate).text).toBe("original")
      expect(yield* llm.requests()).toEqual([{ ...request, promptCacheKey: "observation" }])
    }),
  )

  it.effect("broadcasts request-arrival waits and satisfies waits registered afterward", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      yield* llm.always(TestLLM.stop())
      const first = yield* llm.wait(2).pipe(Effect.forkChild({ startImmediately: true }))
      const second = yield* llm.wait(2).pipe(Effect.forkChild({ startImmediately: true }))
      yield* llm.generate(request)
      expect(first.pollUnsafe()).toBeUndefined()
      expect(second.pollUnsafe()).toBeUndefined()
      yield* llm.generate(request)
      yield* Fiber.join(first)
      yield* Fiber.join(second)
      yield* llm.wait(2)
      expect(yield* llm.requests()).toHaveLength(2)
    }),
  )
  ;(["queued", "served"] as const).forEach((mode) => {
    it.effect(`assigns ${mode} replies before resuming request-arrival continuations`, () =>
      Effect.gen(function* () {
        const llm = yield* TestLLM.Test
        const responses = [TestLLM.text("first", "first"), TestLLM.text("second", "second")]
        yield* mode === "queued" ? llm.push(...responses) : llm.serve(() => responses.shift() ?? [])
        const later = yield* llm
          .wait(1)
          .pipe(Effect.andThen(llm.generate(request)), Effect.forkChild({ startImmediately: true }))

        expect((yield* llm.generate(request)).text).toBe("first")
        expect((yield* Fiber.join(later)).text).toBe("second")
        expect(yield* llm.requests()).toEqual([request, request])
      }),
    )
  })

  it.effect("notifies arrival waiters even when the responder defects", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      const defect = new Error("Broken fixture responder")
      yield* llm.serve(() => {
        throw defect
      })
      const waiter = yield* llm.wait(1).pipe(Effect.forkChild({ startImmediately: true }))
      expect(yield* llm.generate(request).pipe(Effect.catchDefect(Effect.succeed))).toBe(defect)
      yield* Fiber.join(waiter)
    }),
  )

  it.effect("builds independent state even when the same layer is provided concurrently", () => {
    const layer = TestLLM.testLayer()
    const run = Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      expect(yield* llm.requests()).toEqual([])
      yield* llm.push(TestLLM.text("one", "answer"))
      expect((yield* LLMClient.generate(request)).text).toBe("one")
      return yield* llm.requests()
    }).pipe(Effect.provide(layer))
    return Effect.gen(function* () {
      expect(yield* Effect.all([run, run], { concurrency: "unbounded" })).toEqual([[request], [request]])
    })
  })

  it.effect("counts concurrent starts on one gate without serializing their response assignment", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      yield* llm.push(TestLLM.text("first", "first"), TestLLM.text("second", "second"))
      const generate = llm.generate(request)
      const gate = yield* llm.gate()
      const first = yield* generate.pipe(Effect.forkChild({ startImmediately: true }))
      yield* gate.started
      const second = yield* generate.pipe(Effect.forkChild({ startImmediately: true }))
      yield* gate.started
      yield* llm.wait(2)
      expect(first.pollUnsafe()).toBeUndefined()
      expect(second.pollUnsafe()).toBeUndefined()
      yield* gate.release
      expect((yield* Fiber.join(first)).text).toBe("first")
      expect((yield* Fiber.join(second)).text).toBe("second")
    }),
  )

  it.effect("does not clear a replacement gate when the previous gate is released", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      yield* llm.always(TestLLM.stop())
      const previous = yield* llm.gate()
      const first = yield* llm.generate(request).pipe(Effect.forkChild({ startImmediately: true }))
      yield* previous.started
      const next = yield* llm.gate()
      yield* previous.release
      yield* Fiber.join(first)

      const second = yield* llm.generate(request).pipe(Effect.forkChild({ startImmediately: true }))
      yield* next.started
      expect(second.pollUnsafe()).toBeUndefined()
      yield* next.release
      yield* Fiber.join(second)
    }),
  )

  it.effect("releases a gate when its deliberately narrower scope closes", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      yield* llm.always(TestLLM.stop())
      // Only the gate is scoped here; its release must happen before the test ends.
      const run = yield* Effect.scoped(
        Effect.gen(function* () {
          const gate = yield* llm.gate()
          const run = yield* llm.generate(request).pipe(Effect.forkChild({ startImmediately: true }))
          yield* gate.started
          return run
        }),
      )
      yield* Fiber.join(run)
      yield* llm.generate(request)
      expect(yield* llm.requests()).toHaveLength(2)
    }),
  )

  it.effect("keeps an executed response consumed after interruption and permits later requests", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      yield* llm.push(TestLLM.text("interrupted", "first"), TestLLM.text("next", "second"))
      const gate = yield* llm.gate()
      const run = yield* llm.generate(request).pipe(Effect.forkChild({ startImmediately: true }))
      yield* gate.started
      yield* Fiber.interrupt(run)
      yield* gate.release
      expect((yield* llm.generate(request)).text).toBe("next")
      expect(yield* llm.requests()).toHaveLength(2)
    }),
  )

  it.effect("consumes a supplied stream's post-finish tail and runs its finalizer", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      const tail = yield* Latch.make()
      const release = yield* Latch.make()
      const finalized = yield* Latch.make()
      yield* llm.push(
        Stream.unwrap(
          Effect.gen(function* () {
            yield* Effect.addFinalizer(() => finalized.open)
            return Stream.fromIterable(TestLLM.text("complete", "answer")).pipe(
              Stream.concat(Stream.fromEffect(tail.open.pipe(Effect.andThen(release.await))).pipe(Stream.drain)),
            )
          }),
        ),
      )
      const run = yield* llm.generate(request).pipe(Effect.forkChild({ startImmediately: true }))
      yield* tail.await
      expect(run.pollUnsafe()).toBeUndefined()
      yield* release.open
      expect((yield* Fiber.join(run)).text).toBe("complete")
      yield* finalized.await
    }),
  )

  it.effect("preserves irregular events, ordinary EOF, typed failures, and responder defects", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      const events = [LLMEvent.textDelta({ id: "without-start", text: "partial" })]
      yield* llm.push(events, [])
      expect(yield* Stream.runCollect(llm.stream(request))).toEqual(events)
      expect(yield* Stream.runCollect(llm.stream(request))).toEqual([])

      const failure = new AIError({ reason: new RateLimitError({ message: "Try later" }) })
      const observed: LLMEvent[] = []
      yield* llm.serve(() => TestLLM.failAfter(failure, ...events))
      expect(
        yield* llm.stream(request).pipe(
          Stream.runForEach((event) => Effect.sync(() => observed.push(event))),
          Effect.flip,
        ),
      ).toBe(failure)
      expect(observed).toEqual(events)
      const generated = yield* llm.generate(request).pipe(Effect.flip)
      expect(generated.reason).toBe(failure.reason)
      expect(generated.partial?.text).toBe("partial")

      const defect = new Error("Broken fixture responder")
      yield* llm.serve(() => {
        throw defect
      })
      expect(yield* llm.generate(request).pipe(Effect.catchDefect(Effect.succeed))).toBe(defect)
      yield* llm.push(TestLLM.text("recovered", "answer"))
      expect((yield* llm.generate(request)).text).toBe("recovered")
    }),
  )

  it.effect("defects on unexpected requests instead of waiting for a late script", () =>
    Effect.gen(function* () {
      const llm = yield* TestLLM.Test
      const defect = yield* llm.generate(request).pipe(Effect.catchDefect(Effect.succeed))
      expect(defect).toBeInstanceOf(Error)
      if (!(defect instanceof Error)) return
      expect(defect.message).toBe("TestLLM has no response for request 1")
      expect(yield* llm.requests()).toEqual([request])
      yield* llm.push(TestLLM.stop())
      yield* llm.generate(request)
    }),
  )
})