opencode/packages/ai/test/promise.test.ts

import { describe, expect, test } from "bun:test"
import { mkdtemp, rm } from "node:fs/promises"
import { tmpdir } from "node:os"
import { join } from "node:path"
import { Effect, Layer, Schema } from "effect"
import { HttpClientRequest } from "effect/http"
import { AIError, LLMEvent, Media } from "../src/effect.js"
import { SpeechEvent, TranscriptionEvent } from "../src/experimental.js"
import { RequestExecutor, type StreamOptions } from "../src/route.js"
import { AI } from "../src/index.js"
import { AssemblyAI, OpenAI, Runway } from "../src/providers.js"
import { handlerLayer, json } from "./lib/http.js"
import { finishChunk, toolCallChunk } from "./lib/openai-chunks.js"
import { sseEvents } from "./lib/sse.js"

const openai = OpenAI.configure({ apiKey: "test", baseURL: "https://openai.test/v1" })

const chatBody = sseEvents(
  { choices: [{ delta: { content: "Hello" } }] },
  { choices: [{ delta: { content: " world" } }] },
  { choices: [{ delta: {}, finish_reason: "stop" }] },
)

/**
 * Executor layer that answers chat completions with SSE text, image generations with one base64 PNG, Runway video
 * tasks with a queued submission that succeeds on the second poll, speech with raw audio or SSE audio deltas, OpenAI
 * transcription with JSON or SSE text deltas, AssemblyAI transcripts that complete on the first poll, and `slow.test`
 * chat completions that send one text delta and never finish.
 */
const executor = (seen: Array<string>) =>
  RequestExecutor.layer.pipe(
    Layer.provide(
      handlerLayer((input) =>
        Effect.gen(function* () {
          const web = yield* HttpClientRequest.toWeb(input.request).pipe(Effect.orDie)
          seen.push(web.url)
          if (web.url.endsWith("/images/generations"))
            return JSON.parse(input.text).stream === true
              ? input.respond(
                  sseEvents(
                    {
                      type: "image_generation.partial_image",
                      b64_json: "AQ==",
                      partial_image_index: 0,
                      output_format: "png",
                    },
                    { type: "image_generation.completed", b64_json: "AQID", output_format: "png" },
                  ),
                  { headers: { "content-type": "text/event-stream" } },
                )
              : input.respond(JSON.stringify({ data: [{ b64_json: "AQID" }], output_format: "png" }), {
                  headers: { "content-type": "application/json" },
                })
          if (web.url.startsWith("https://slow.test"))
            return input.respond(
              new ReadableStream({
                start: (controller) =>
                  controller.enqueue(
                    new TextEncoder().encode(
                      `data: ${JSON.stringify({ choices: [{ delta: { content: "Hello" } }] })}\n\n`,
                    ),
                  ),
              }),
              { headers: { "content-type": "text/event-stream" } },
            )
          if (web.url.endsWith("/chat/completions"))
            return input.respond(chatBody, { headers: { "content-type": "text/event-stream" } })
          if (web.url.endsWith("/audio/speech"))
            return JSON.parse(input.text).stream_format === "sse"
              ? input.respond(
                  sseEvents(
                    { type: "speech.audio.delta", audio: "AQI=" },
                    { type: "speech.audio.delta", audio: "Aw==" },
                    { type: "speech.audio.done", usage: { input_tokens: 4, output_tokens: 8, total_tokens: 12 } },
                  ),
                  { headers: { "content-type": "text/event-stream" } },
                )
              : input.respond(Uint8Array.from([1, 2, 3]), { headers: { "content-type": "audio/pcm" } })
          if (web.url.endsWith("/audio/transcriptions"))
            return input.text.includes('name="stream"')
              ? input.respond(
                  sseEvents(
                    { type: "transcript.text.delta", delta: "Hello" },
                    { type: "transcript.text.delta", delta: " there." },
                    { type: "transcript.text.done", text: "Hello there." },
                  ),
                  { headers: { "content-type": "text/event-stream" } },
                )
              : input.respond(JSON.stringify({ text: "Hello there." }), {
                  headers: { "content-type": "application/json" },
                })
          if (web.url.endsWith("/v2/transcript"))
            return input.respond(JSON.stringify({ id: "tr_1", status: "queued" }), {
              headers: { "content-type": "application/json" },
            })
          if (web.url.endsWith("/v2/transcript/tr_1"))
            return input.respond(JSON.stringify({ id: "tr_1", status: "completed", text: "Hello there." }), {
              headers: { "content-type": "application/json" },
            })
          if (web.url.endsWith("/text_to_video"))
            return input.respond(JSON.stringify({ id: "task_1" }), { headers: { "content-type": "application/json" } })
          if (web.url.endsWith("/tasks/task_1")) {
            const polls = seen.filter((url) => url.endsWith("/tasks/task_1")).length
            return input.respond(
              JSON.stringify(
                polls === 1
                  ? { status: "RUNNING", progress: 0.5 }
                  : { status: "SUCCEEDED", output: ["https://runway.test/out.mp4"], cost: { credits: 5 } },
              ),
              { headers: { "content-type": "application/json" } },
            )
          }
          return input.respond(JSON.stringify({ error: { message: "not found" } }), {
            status: 404,
            headers: { "content-type": "application/json" },
          })
        }),
      ),
    ),
  )

describe("AI promise client", () => {
  test("generates text, images, and streams over one managed runtime, and queues images", async () => {
    const seen: Array<string> = []
    const ai = AI.make({ layer: executor(seen) })

    const request = ai.request({ model: openai.chat("gpt-4o-mini"), prompt: "Say hello." })
    const text = await ai.generate(request)
    expect(text.text).toBe("Hello world")
    expect((await ai.generate({ model: openai.chat("gpt-4o-mini"), prompt: "Say hello." })).text).toBe("Hello world")

    const image = await ai.image.generate({ model: openai.image("gpt-image-2"), prompt: "A lighthouse" })
    expect(image.image).toBeInstanceOf(Media.Asset)
    expect(image.image.mediaType).toBe("image/png")
    expect(await ai.run(image.image.bytes())).toEqual(Uint8Array.from([1, 2, 3]))
    const requested = await ai.image.generate(
      ai.image.request({ model: openai.image("gpt-image-2"), prompt: "A lighthouse" }),
    )
    expect(requested.image.mediaType).toBe("image/png")

    const deltas: Array<string> = []
    for await (const event of ai.stream(request)) {
      if (LLMEvent.is.textDelta(event)) deltas.push(event.text)
    }
    expect(deltas).toEqual(["Hello", " world"])
    const directDeltas: Array<string> = []
    for await (const event of ai.stream({ model: openai.chat("gpt-4o-mini"), prompt: "Say hello." })) {
      if (LLMEvent.is.textDelta(event)) directDeltas.push(event.text)
    }
    expect(directDeltas).toEqual(deltas)

    const imageEvents: Array<string> = []
    for await (const event of ai.image.stream({ model: openai.image("gpt-image-2"), prompt: "A lighthouse" })) {
      imageEvents.push(event.type)
    }
    expect(imageEvents).toEqual(["image-partial", "image", "finish"])

    expect(seen).toEqual([
      "https://openai.test/v1/chat/completions",
      "https://openai.test/v1/chat/completions",
      "https://openai.test/v1/images/generations",
      "https://openai.test/v1/images/generations",
      "https://openai.test/v1/chat/completions",
      "https://openai.test/v1/chat/completions",
      "https://openai.test/v1/images/generations",
    ])

    await ai.dispose()
  })

  test("starts, awaits, and resumes video generations over the same runtime", async () => {
    const seen: Array<string> = []
    const ai = AI.make({ layer: executor(seen) })
    const model = Runway.configure({ apiKey: "test", baseURL: "https://runway.test/v1" }).video("gen4.5")

    const generation = await ai.experimental.video.start({ model, prompt: "A kite" })
    expect(generation.id).toBe("task_1")
    expect(generation.status).toBe("queued")
    expect(generation.token).toEqual({ taskID: "task_1" })

    const refreshed = await generation.refresh()
    expect(refreshed.status).toBe("running")
    expect(refreshed.progress).toBe(0.5)

    const response = await refreshed.await({ poll: { interval: 10 } })
    expect(response.video).toBeInstanceOf(Media.Asset)
    expect(response.video.source).toMatchObject({ type: "url", url: "https://runway.test/out.mp4" })
    expect(response.usage).toEqual({ type: "credits", credits: 5 })

    const resumed = await ai.experimental.video.resume(model, JSON.parse(JSON.stringify(generation.token)))
    expect(resumed.status).toBe("completed")
    expect((await resumed.await()).videos).toHaveLength(1)

    const events: Array<string> = []
    for await (const event of ai.experimental.video.stream({ model, prompt: "A kite" }, { poll: { interval: 10 } })) {
      events.push(event.type)
    }
    expect(events).toEqual(["video", "finish"])

    expect(seen[0]).toBe("https://runway.test/v1/text_to_video")
    expect(seen.filter((url) => url.endsWith("/tasks/task_1")).length).toBeGreaterThanOrEqual(5)
    await ai.dispose()
  })

  test("forwards per-call options to generate, stream, and loop, and checks compaction support", async () => {
    const seen: Array<string> = []
    const ai = AI.make({ layer: executor([]) })
    const options: StreamOptions = {
      http: (request, handler) => {
        seen.push(request.url)
        return handler(request)
      },
    }
    const request = { model: openai.chat("gpt-4o-mini"), prompt: "Say hello." }

    expect((await ai.generate(request, options)).text).toBe("Hello world")
    await Array.fromAsync(ai.stream(request, options))
    await Array.fromAsync(ai.loop({ ...request, tools: [] }, options))
    expect(seen).toEqual(Array(3).fill("https://openai.test/v1/chat/completions"))

    const responses = ai.request({ model: openai.responses("gpt-5"), prompt: "Hello" })
    expect(ai.canCompact(responses)).toBe(true)
    expect(ai.canCompact(responses, { mechanism: "trigger" })).toBe(true)
    expect(ai.canCompact(ai.request(request))).toBe(false)
    await ai.dispose()
  })

  test("exposes a generation's snapshot and terminal state", async () => {
    const ai = AI.make({ layer: executor([]) })
    const model = Runway.configure({ apiKey: "test", baseURL: "https://runway.test/v1" }).video("gen4.5")
    const generation = await ai.experimental.video.start({ model, prompt: "A kite" })
    expect(generation.snapshot).toEqual({ id: "task_1", status: "queued", progress: undefined, position: undefined })
    expect(generation.terminal).toBe(false)
    const refreshed = await generation.refresh()
    expect(refreshed.terminal).toBe(false)
    expect((await refreshed.refresh()).terminal).toBe(true)
    await ai.dispose()
  })

  test("observes a started generation's events and fetches its result", async () => {
    const ai = AI.make({ layer: executor([]) })
    const model = Runway.configure({ apiKey: "test", baseURL: "https://runway.test/v1" }).video("gen4.5")
    const generation = await ai.experimental.video.start({ model, prompt: "A kite" })

    const events: Array<string> = []
    for await (const event of generation.events({ poll: { interval: 10 } })) events.push(event.type)
    expect(events).toEqual(["generation-progress", "generation-finished"])
    expect(generation.status).toBe("queued")
    expect((await generation.result()).video.source).toMatchObject({ url: "https://runway.test/out.mp4" })
    await ai.dispose()
  })

  test("reads, writes, and decodes assets without leaving promises", async () => {
    const ai = AI.make({ layer: executor([]) })
    const dir = await mkdtemp(join(tmpdir(), "ai-promise-"))
    try {
      const png = Uint8Array.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 1, 2, 3])
      await Bun.write(join(dir, "source.bin"), png)

      const asset = await ai.file(join(dir, "source.bin"))
      expect(asset.mediaType).toBe("image/png")
      await ai.write(asset, join(dir, "copy.png"))
      expect(await ai.bytes(await ai.file(join(dir, "copy.png")))).toEqual(png)
    } finally {
      await rm(dir, { recursive: true })
      await ai.dispose()
    }
  })

  test("generates, streams, and starts transcriptions over the same runtime", async () => {
    const ai = AI.make({ layer: executor([]) })
    const audio = Media.url("https://audio.test/hello.mp3")
    const bytes = Media.bytes(Uint8Array.from([0x49, 0x44, 0x33]), "audio/mpeg")

    expect(
      (
        await ai.experimental.transcription.generate({
          model: openai.transcription("gpt-4o-mini-transcribe"),
          audio: bytes,
        })
      ).text,
    ).toBe("Hello there.")

    const deltas: Array<string> = []
    for await (const event of ai.experimental.transcription.stream({
      model: openai.transcription("gpt-4o-mini-transcribe"),
      audio: bytes,
    }))
      if (TranscriptionEvent.is.textDelta(event)) deltas.push(event.delta)
    expect(deltas).toEqual(["Hello", " there."])

    const model = AssemblyAI.configure({ apiKey: "test", baseURL: "https://assemblyai.test" }).transcription(
      "universal-2",
    )
    const generation = await ai.experimental.transcription.start({ model, audio })
    expect(generation.token).toEqual({ transcriptID: "tr_1" })
    expect((await generation.await({ poll: { interval: 10 } })).text).toBe("Hello there.")
    await ai.dispose()
  })

  test("generates and streams speech over the same runtime", async () => {
    const ai = AI.make({ layer: executor([]) })
    const model = openai.speech("gpt-4o-mini-tts")

    const response = await ai.experimental.speech.generate({ model, text: "Hello", voice: "coral", format: "pcm" })
    expect(await ai.run(response.audio.bytes())).toEqual(Uint8Array.from([1, 2, 3]))

    const events: Array<SpeechEvent> = []
    for await (const event of ai.experimental.speech.stream({ model, text: "Hello", voice: "coral" }))
      events.push(event)
    expect(events.map((event) => event.type)).toEqual(["audio-delta", "audio-delta", "finish"])
    const finish = events.find(SpeechEvent.is.finish)
    expect(await ai.run(finish!.audio.bytes())).toEqual(Uint8Array.from([1, 2, 3]))
    await ai.dispose()
  })

  test("rethrows AIError unchanged and honors abort signals", async () => {
    const ai = AI.make({ layer: executor([]) })

    const failure = await ai
      .generate(ai.request({ model: openai.responses("gpt-5"), prompt: "Hello" }))
      .then(() => undefined)
      .catch((error: unknown) => error)
    expect(failure).toBeInstanceOf(AIError)
    expect(failure instanceof AIError && failure.reason.http?.status).toBe(404)

    const invalidLLM = await ai
      // @ts-expect-error Invalid input must reject with AIError instead of throwing synchronously.
      .generate({ model: openai.responses("gpt-5"), messages: [{ role: "bogus" }] })
      .catch((error: unknown) => error)
    expect(invalidLLM instanceof AIError && invalidLLM.reason._tag).toBe("InvalidRequest")

    const invalid = await ai.image
      .generate({ model: openai.image("gpt-image-2"), prompt: "A lighthouse", n: 1.5 })
      .catch((error: unknown) => error)
    expect(invalid instanceof AIError && invalid.reason._tag).toBe("InvalidRequest")

    const controller = new AbortController()
    controller.abort()
    const aborted = await ai
      .generate(ai.request({ model: openai.chat("gpt-4o-mini"), prompt: "Hello" }), { signal: controller.signal })
      .then(() => "completed")
      .catch(() => "aborted")
    expect(aborted).toBe("aborted")

    await ai.dispose()
  })

  test("aborted calls reject and aborted streams throw with the signal's reason", async () => {
    const ai = AI.make({ layer: executor([]) })
    const slow = OpenAI.configure({ apiKey: "test", baseURL: "https://slow.test/v1" }).chat("gpt-4o-mini")
    const aborted = new AbortController()
    aborted.abort()
    const reason = new Error("mine")

    const rejected = await ai.run(Effect.never, { signal: aborted.signal }).catch((error: unknown) => error)
    expect(rejected).toBe(aborted.signal.reason)
    expect(rejected).toMatchObject({ name: "AbortError" })

    const inFlight = new AbortController()
    setTimeout(() => inFlight.abort(reason), 10)
    expect(
      await ai.generate({ model: slow, prompt: "Hello" }, { signal: inFlight.signal }).catch((error: unknown) => error),
    ).toBe(reason)

    const preAborted = await Array.fromAsync(
      ai.experimental.speech.stream(
        { model: openai.speech("gpt-4o-mini-tts"), text: "Hello" },
        { signal: aborted.signal },
      ),
    ).catch((error: unknown) => error)
    expect(preAborted).toBe(aborted.signal.reason)
    expect(preAborted).toMatchObject({ name: "AbortError" })

    const midStream = new AbortController()
    const deltas: Array<string> = []
    const midStreamFailure = await Array.fromAsync(
      ai.stream({ model: slow, prompt: "Hello" }, { signal: midStream.signal }),
      (event) => {
        if (!LLMEvent.is.textDelta(event)) return
        deltas.push(event.text)
        midStream.abort()
      },
    ).catch((error: unknown) => error)
    expect(deltas).toEqual(["Hello"])
    expect(midStreamFailure).toBe(midStream.signal.reason)
    expect(midStreamFailure).toMatchObject({ name: "AbortError" })

    const model = Runway.configure({ apiKey: "test", baseURL: "https://runway.test/v1" }).video("gen4.5")
    const generation = await ai.experimental.video.start({ model, prompt: "A kite" })
    const polling = new AbortController()
    const events: Array<string> = []
    const eventsFailure = await Array.fromAsync(
      generation.events({ poll: { interval: 60_000 }, signal: polling.signal }),
      (event) => {
        events.push(event.type)
        polling.abort(reason)
      },
    ).catch((error: unknown) => error)
    expect(events).toEqual(["generation-progress"])
    expect(eventsFailure).toBe(reason)

    await ai.dispose()
  })

  test("breaking out of an abortable stream cleans up without throwing", async () => {
    const ai = AI.make({ layer: executor([]) })
    const slow = OpenAI.configure({ apiKey: "test", baseURL: "https://slow.test/v1" }).chat("gpt-4o-mini")
    const controller = new AbortController()
    const deltas: Array<string> = []
    for await (const event of ai.stream({ model: slow, prompt: "Hello" }, { signal: controller.signal })) {
      if (!LLMEvent.is.textDelta(event)) continue
      deltas.push(event.text)
      break
    }
    controller.abort()
    expect(deltas).toEqual(["Hello"])
    await ai.dispose()
  })

  test("loops with async tools and sends rejections back to the model as error results", async () => {
    const bodies: Array<string> = []
    const ai = AI.make({
      layer: RequestExecutor.layer.pipe(
        Layer.provide(
          handlerLayer((input) => {
            bodies.push(input.text)
            return Effect.succeed(
              input.respond(
                bodies.length === 1
                  ? sseEvents(
                      toolCallChunk("call-1", "read_file", JSON.stringify({ path: "a.txt" })),
                      toolCallChunk("call-2", "read_file", JSON.stringify({ path: "missing.txt" }), 1),
                      toolCallChunk("call-3", "read_file", JSON.stringify({ path: "empty.txt" }), 2),
                      finishChunk("tool_calls"),
                    )
                  : chatBody,
                { headers: { "content-type": "text/event-stream" } },
              ),
            )
          }),
        ),
      ),
    })
    const read_file = AI.tool({
      name: "read_file",
      description: "Read a file.",
      parameters: Schema.Struct({ path: Schema.String }),
      execute: async ({ path }) => {
        if (path === "missing.txt") throw new Error(`No such file: ${path}`)
        if (path === "empty.txt") return
        return "hello"
      },
    })

    const events = await Array.fromAsync(
      ai.loop({ model: openai.chat("gpt-4o-mini"), prompt: "Read the files.", tools: [read_file] }),
    )

    expect(events.filter(LLMEvent.is.toolResult).map((event) => event.result)).toEqual([
      { type: "text", value: "hello" },
      { type: "error", value: "No such file: missing.txt" },
      { type: "json", value: null },
    ])
    expect(JSON.parse(bodies[1]).messages.filter((message: { role: string }) => message.role === "tool")).toEqual([
      expect.objectContaining({ tool_call_id: "call-1", content: "hello" }),
      expect.objectContaining({ tool_call_id: "call-2", content: "No such file: missing.txt" }),
      expect.objectContaining({ tool_call_id: "call-3", content: "null" }),
    ])
    expect(events.filter(LLMEvent.is.textDelta).map((event) => event.text)).toEqual(["Hello", " world"])
    await ai.dispose()
  })

  test("generates a decoded object", async () => {
    const ai = AI.make({
      layer: RequestExecutor.layer.pipe(
        Layer.provide(
          handlerLayer((input) =>
            Effect.succeed(
              input.respond(
                sseEvents(
                  toolCallChunk("call-1", "generate_object", JSON.stringify({ city: "Paris" })),
                  finishChunk("tool_calls"),
                ),
                { headers: { "content-type": "text/event-stream" } },
              ),
            ),
          ),
        ),
      ),
    })

    const response = await ai.generate({
      model: openai.chat("gpt-4o-mini"),
      prompt: "Pick a city.",
      schema: Schema.Struct({ city: Schema.String }),
    })

    expect(response.object.city).toBe("Paris")
    await ai.dispose()
  })

  test("compacts through the endpoint and the trigger, and rejects unsupported routes with AIError", async () => {
    const checkpoint = { type: "compaction", id: "cmp_1", encrypted_content: "opaque" }
    const paths: Array<string> = []
    const ai = AI.make({
      layer: RequestExecutor.layer.pipe(
        Layer.provide(
          handlerLayer((input) =>
            Effect.gen(function* () {
              const path = new URL((yield* HttpClientRequest.toWeb(input.request).pipe(Effect.orDie)).url).pathname
              paths.push(path)
              if (path === "/v1/responses/compact")
                return json(input, {
                  object: "response.compaction",
                  output: [checkpoint],
                  usage: { input_tokens: 100, output_tokens: 10, total_tokens: 110 },
                })
              expect(JSON.parse(input.text).input.at(-1)).toEqual({ type: "compaction_trigger" })
              return input.respond(
                sseEvents({ type: "response.completed", response: { id: "resp_1", output: [checkpoint] } }),
                { headers: { "content-type": "text/event-stream" } },
              )
            }),
          ),
        ),
      ),
    })
    const request = ai.request({ model: openai.responses("gpt-5"), prompt: "Hello" })

    const compacted = await ai.compact(request)
    expect(compacted.replacement.flatMap((message) => message.content.map((part) => part.type))).toEqual(["compaction"])
    expect(compacted.usage?.totalTokens).toBe(110)

    const triggered = await ai.compact(request, { mechanism: "trigger" })
    expect(triggered.checkpoint).toMatchObject({ id: "cmp_1", encrypted: "opaque" })
    expect(triggered.responseID).toBe("resp_1")
    expect(paths).toEqual(["/v1/responses/compact", "/v1/responses"])

    const unsupported = await ai
      // @ts-expect-error Untyped callers on routes without a compact endpoint must reject with AIError.
      .compact(ai.request({ model: openai.chat("gpt-4o-mini"), prompt: "Hello" }))
      .catch((error: unknown) => error)
    expect(unsupported instanceof AIError && unsupported.reason._tag).toBe("UnsupportedOperation")
    expect(paths).toHaveLength(2)
    await ai.dispose()
  })

  test("the default client is created lazily", () => {
    expect(typeof AI.ai.generate).toBe("function")
    expect(typeof AI.ai.image.generate).toBe("function")
    expect(typeof AI.ai.experimental.video.start).toBe("function")
    expect(typeof AI.ai.experimental.speech.generate).toBe("function")
  })
})