packages/ai/test/tool-stream.test.ts

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

import { describe, expect, test } from "bun:test"
import { Effect } from "effect"
import { AIError } from "../src/schema/index.js"
import { ToolStream } from "../src/protocols/utils/tool-stream.js"
import { it } from "./lib/effect.js"

const ADAPTER = "test-route"

describe("ToolStream", () => {
  it.effect("starts from OpenAI-style deltas and finalizes parsed input", () =>
    Effect.gen(function* () {
      const first = ToolStream.appendOrStart(
        ADAPTER,
        ToolStream.empty<number>(),
        0,
        { id: "call_1", name: "lookup", text: '{"query"' },
        "missing tool",
      )
      if (ToolStream.isError(first)) return yield* first
      const second = ToolStream.appendOrStart(ADAPTER, first.tools, 0, { text: ':"weather"}' }, "missing tool")
      if (ToolStream.isError(second)) return yield* second
      const finished = yield* ToolStream.finish(ADAPTER, second.tools, 0)

      expect(first.events).toEqual([
        { type: "tool-input-start", id: "call_1", name: "lookup" },
        { type: "tool-input-delta", id: "call_1", name: "lookup", text: '{"query"', input: {} },
      ])
      expect(second.events).toEqual([
        { type: "tool-input-delta", id: "call_1", name: "lookup", text: ':"weather"}', input: { query: "weather" } },
      ])
      expect(finished).toEqual({
        tools: {},
        events: [
          { type: "tool-input-end", id: "call_1", name: "lookup" },
          { type: "tool-call", id: "call_1", name: "lookup", input: { query: "weather" } },
        ],
      })
    }),
  )

  test("exposes cumulative partial string values", () => {
    const result = ToolStream.appendOrStart(
      ADAPTER,
      ToolStream.empty<number>(),
      0,
      { id: "call_1", name: "lookup", text: '{"query":"wea' },
      "missing tool",
    )
    if (ToolStream.isError(result)) throw result

    expect(result.events.at(-1)).toEqual({
      type: "tool-input-delta",
      id: "call_1",
      name: "lookup",
      text: '{"query":"wea',
      input: { query: "wea" },
    })
  })

  test("defaults partial input to an empty object when the accumulated value cannot be parsed", () => {
    const result = ToolStream.appendOrStart(
      ADAPTER,
      ToolStream.empty<number>(),
      0,
      { id: "call_1", name: "lookup", text: "x" },
      "missing tool",
    )
    if (ToolStream.isError(result)) throw result

    expect(result.events).toEqual([
      { type: "tool-input-start", id: "call_1", name: "lookup" },
      { type: "tool-input-delta", id: "call_1", name: "lookup", text: "x", input: {} },
    ])
  })

  it.effect("keeps accumulated identity when later deltas contain empty strings", () =>
    Effect.gen(function* () {
      const first = ToolStream.appendOrStart(
        ADAPTER,
        ToolStream.empty<number>(),
        0,
        { id: "call_1", name: "lookup", text: '{"query"' },
        "missing tool",
      )
      if (ToolStream.isError(first)) return yield* first
      const second = ToolStream.appendOrStart(
        ADAPTER,
        first.tools,
        0,
        { id: "", name: "", text: ':"weather"}' },
        "missing tool",
      )
      if (ToolStream.isError(second)) return yield* second
      const finished = yield* ToolStream.finish(ADAPTER, second.tools, 0)

      expect(finished.events).toEqual([
        { type: "tool-input-end", id: "call_1", name: "lookup" },
        { type: "tool-call", id: "call_1", name: "lookup", input: { query: "weather" } },
      ])
    }),
  )

  test("fails appendExisting when the provider skipped the tool start", () => {
    const error = ToolStream.appendExisting(ADAPTER, ToolStream.empty<number>(), 0, "{}", "missing tool")

    expect(error).toBeInstanceOf(AIError)
    if (ToolStream.isError(error)) expect(error.message).toBe("missing tool")
  })

  it.effect("uses final input override without losing accumulated deltas", () =>
    Effect.gen(function* () {
      const tools = ToolStream.start(ToolStream.empty<string>(), "item_1", {
        id: "call_1",
        name: "lookup",
        input: '{"query":"partial"}',
      })
      const finished = yield* ToolStream.finish(ADAPTER, tools, "item_1", '{"query":"final"}')

      expect(finished).toEqual({
        tools: {},
        events: [
          { type: "tool-input-end", id: "call_1", name: "lookup" },
          { type: "tool-call", id: "call_1", name: "lookup", input: { query: "final" } },
        ],
      })
    }),
  )

  it.effect("rejects incomplete final local input", () =>
    Effect.gen(function* () {
      const tools = ToolStream.start(ToolStream.empty<string>(), "item_1", {
        id: "call_1",
        name: "lookup",
        input: '{"query":"partial',
      })
      const finished = yield* ToolStream.finish(ADAPTER, tools, "item_1")

      expect(finished).toEqual({
        tools: {},
        events: [
          { type: "tool-input-end", id: "call_1", name: "lookup" },
          { type: "tool-input-error", id: "call_1", name: "lookup", raw: '{"query":"partial' },
        ],
      })
    }),
  )

  it.effect("rejects malformed string escapes in final local input", () =>
    Effect.gen(function* () {
      const tools = ToolStream.start(ToolStream.empty<string>(), "item_1", {
        id: "call_1",
        name: "lookup",
        input: '{"path":"A\\H","text":"first\tsecond"}',
      })
      const finished = yield* ToolStream.finish(ADAPTER, tools, "item_1")

      expect(finished.events).toEqual([
        { type: "tool-input-end", id: "call_1", name: "lookup" },
        { type: "tool-input-error", id: "call_1", name: "lookup", raw: '{"path":"A\\H","text":"first\tsecond"}' },
      ])
    }),
  )

  it.effect("rejects invalid final local input instead of defaulting to an empty object", () =>
    Effect.gen(function* () {
      const tools = ToolStream.start(ToolStream.empty<string>(), "item_1", {
        id: "call_1",
        name: "lookup",
        input: "invalid",
      })
      const finished = yield* ToolStream.finish(ADAPTER, tools, "item_1")

      expect(finished.events).toEqual([
        { type: "tool-input-end", id: "call_1", name: "lookup" },
        { type: "tool-input-error", id: "call_1", name: "lookup", raw: "invalid" },
      ])
    }),
  )

  it.effect("rejects incomplete input without dropping valid parallel tool calls", () =>
    Effect.gen(function* () {
      const valid = ToolStream.start(ToolStream.empty<number>(), 0, {
        id: "call_valid",
        name: "lookup",
        input: '{"query":"weather"}',
      })
      const tools = ToolStream.start(valid, 1, {
        id: "call_invalid",
        name: "lookup",
        input: '{"query":"partial',
      })
      const finished = yield* ToolStream.finishAll(ADAPTER, tools)

      expect(finished).toEqual({
        tools: {},
        events: [
          { type: "tool-input-end", id: "call_valid", name: "lookup" },
          { type: "tool-call", id: "call_valid", name: "lookup", input: { query: "weather" } },
          { type: "tool-input-end", id: "call_invalid", name: "lookup" },
          { type: "tool-input-error", id: "call_invalid", name: "lookup", raw: '{"query":"partial' },
        ],
      })
    }),
  )

  it.effect("keeps malformed provider-executed input terminal", () =>
    Effect.gen(function* () {
      const tools = ToolStream.start(ToolStream.empty<string>(), "item_1", {
        id: "call_1",
        name: "web_search",
        input: '{"query":"partial',
        providerExecuted: true,
      })
      const result = yield* Effect.exit(ToolStream.finish(ADAPTER, tools, "item_1"))

      expect(result._tag).toBe("Failure")
    }),
  )

  it.effect("preserves providerExecuted and clears all tools", () =>
    Effect.gen(function* () {
      const first: ToolStream.State<number> = ToolStream.start(ToolStream.empty<number>(), 0, {
        id: "call_1",
        name: "lookup",
        input: "{}",
      })
      const tools = ToolStream.start(first, 1, {
        id: "call_2",
        name: "web_search",
        input: '{"query":"docs"}',
        providerExecuted: true,
      })
      const finished = yield* ToolStream.finishAll(ADAPTER, tools)

      expect(finished).toEqual({
        tools: {},
        events: [
          { type: "tool-input-end", id: "call_1", name: "lookup" },
          { type: "tool-call", id: "call_1", name: "lookup", input: {} },
          { type: "tool-input-end", id: "call_2", name: "web_search" },
          {
            type: "tool-call",
            id: "call_2",
            name: "web_search",
            input: { query: "docs" },
            providerExecuted: true,
          },
        ],
      })
    }),
  )
})