@effect-uai/core 0.15.0 → 0.17.0

package/src/streaming/Settle.test.ts

@@ -0,0 +1,344 @@
+import { describe, it } from "@effect/vitest"
+import { Cause, Effect, Fiber, Option, Queue, Ref, Stream } from "effect"
+import { TestClock } from "effect/testing"
+import { expect } from "vitest"
+import { drainBurst, onQuiet, settleBurst } from "./Settle.js"
+
+describe("drainBurst", () => {
+  it.effect("resets the window on every arrival, so a burst outlives one settle", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string>()
+      const burst = yield* Effect.forkChild(drainBurst(queue, "1 second"))
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("900 millis")
+      yield* Queue.offer(queue, "b")
+      yield* TestClock.adjust("900 millis")
+      // 1.8s elapsed, more than one settle, and the burst is still open.
+      expect(burst.pollUnsafe()).toBeUndefined()
+
+      yield* TestClock.adjust("1 second")
+      expect(yield* Fiber.join(burst)).toEqual(["a", "b"])
+    }),
+  )
+
+  it.effect("waits for the first item without a deadline", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string>()
+      const burst = yield* Effect.forkChild(drainBurst(queue, "1 second"))
+
+      yield* TestClock.adjust("1 hour")
+      expect(burst.pollUnsafe()).toBeUndefined()
+
+      yield* Queue.offer(queue, "late")
+      yield* TestClock.adjust("1 second")
+      expect(yield* Fiber.join(burst)).toEqual(["late"])
+    }),
+  )
+
+  it.effect("ends on the first quiet gap and leaves later items for the next drain", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string>()
+      const burst = yield* Effect.forkChild(drainBurst(queue, "1 second"))
+
+      yield* Queue.offer(queue, "first")
+      yield* TestClock.adjust("1 second")
+      expect(yield* Fiber.join(burst)).toEqual(["first"])
+
+      yield* Queue.offer(queue, "second")
+      expect(yield* Queue.poll(queue)).toEqual(Option.some("second"))
+    }),
+  )
+
+  it.effect("yields the batch in hand when the queue ends mid-burst", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const burst = yield* Effect.forkChild(drainBurst(queue, "1 second"))
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("100 millis")
+      yield* Queue.end(queue)
+      yield* TestClock.adjust("10 millis")
+
+      expect(yield* Fiber.join(burst)).toEqual(["a"])
+    }),
+  )
+
+  it.effect("propagates a queue failure rather than treating it as the end", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, string | Cause.Done>()
+      const burst = yield* Effect.forkChild(drainBurst(queue, "1 second"))
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("100 millis")
+      yield* Queue.fail(queue, "boom")
+      yield* TestClock.adjust("10 millis")
+
+      const exit = yield* Effect.exit(Fiber.join(burst))
+      expect(exit._tag).toBe("Failure")
+      expect(JSON.stringify(exit)).toContain("boom")
+    }),
+  )
+
+  it.effect("does not mistake an interrupt mid-burst for the end of the queue", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const burst = yield* Effect.forkChild(drainBurst(queue, "1 second"))
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("100 millis")
+      yield* Fiber.interrupt(burst)
+
+      // Returning ["a"] here would mean interruption silently looked like a
+      // settled burst.
+      const exit = yield* Effect.exit(Fiber.join(burst))
+      expect(exit._tag).toBe("Failure")
+    }),
+  )
+})
+
+describe("settleBurst", () => {
+  const collecting = <A>(stream: Stream.Stream<A, never>) =>
+    Effect.gen(function* () {
+      const seen = yield* Ref.make<ReadonlyArray<A>>([])
+      const fiber = yield* Effect.forkChild(
+        Stream.runForEach(stream, (a) => Ref.update(seen, (xs) => [...xs, a])),
+      )
+      return { seen, fiber }
+    })
+
+  it.effect("holds a growing burst back, then emits it as one batch", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const { seen } = yield* collecting(Stream.fromQueue(queue).pipe(settleBurst("1 second")))
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("900 millis")
+      yield* Queue.offer(queue, "b")
+      yield* TestClock.adjust("900 millis")
+      // Each arrival reset the window, so nothing has been emitted yet.
+      expect(yield* Ref.get(seen)).toEqual([])
+
+      yield* TestClock.adjust("1 second")
+      expect(yield* Ref.get(seen)).toEqual([["a", "b"]])
+    }),
+  )
+
+  it.effect("emits the in-flight batch when the source ends mid-burst", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const { seen } = yield* collecting(Stream.fromQueue(queue).pipe(settleBurst("1 second")))
+
+      yield* Queue.offer(queue, "only")
+      yield* TestClock.adjust("100 millis")
+      yield* Queue.end(queue)
+      yield* TestClock.adjust("10 millis")
+
+      expect(yield* Ref.get(seen)).toEqual([["only"]])
+    }),
+  )
+
+  it.effect("splits bursts on every silence, keeping order across a conversation", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const { seen } = yield* collecting(Stream.fromQueue(queue).pipe(settleBurst("1 second")))
+
+      // Three quick lines, a pause, two more, a pause, then one alone.
+      yield* Queue.offer(queue, "hey")
+      yield* TestClock.adjust("100 millis")
+      yield* Queue.offer(queue, "are you there")
+      yield* TestClock.adjust("100 millis")
+      yield* Queue.offer(queue, "?")
+      yield* TestClock.adjust("2 seconds")
+
+      yield* Queue.offer(queue, "actually")
+      yield* TestClock.adjust("100 millis")
+      yield* Queue.offer(queue, "never mind")
+      yield* TestClock.adjust("2 seconds")
+
+      yield* Queue.offer(queue, "bye")
+      yield* TestClock.adjust("2 seconds")
+
+      expect(yield* Ref.get(seen)).toEqual([
+        ["hey", "are you there", "?"],
+        ["actually", "never mind"],
+        ["bye"],
+      ])
+    }),
+  )
+
+  it.effect("fails the stream when the source fails instead of ending quietly", () =>
+    Effect.gen(function* () {
+      const source: Stream.Stream<string, string> = Stream.make("a").pipe(
+        Stream.concat(Stream.fail("boom")),
+      )
+      const exit = yield* Effect.exit(Stream.runCollect(source.pipe(settleBurst("1 second"))))
+
+      expect(exit._tag).toBe("Failure")
+      expect(JSON.stringify(exit)).toContain("boom")
+    }),
+  )
+
+  it.effect("stops pulling the source while the consumer is blocked", () =>
+    Effect.gen(function* () {
+      const pulled = yield* Ref.make(0)
+      const source = Stream.fromIterable(Array.from({ length: 50 }, (_, i) => i)).pipe(
+        Stream.rechunk(1),
+        Stream.tap(() => Ref.update(pulled, (n) => n + 1)),
+      )
+      yield* Effect.forkChild(
+        source.pipe(
+          settleBurst("1 second", { capacity: 2 }),
+          Stream.runForEach(() => Effect.never),
+        ),
+      )
+      yield* TestClock.adjust("10 seconds")
+
+      // An unbounded buffer would have drained all 50 into memory by now.
+      expect(yield* Ref.get(pulled)).toBeLessThan(50)
+    }),
+  )
+})
+
+describe("onQuiet", () => {
+  const mark = Effect.succeed(Option.some("<quiet>"))
+
+  const collecting = (stream: Stream.Stream<string, never>) =>
+    Effect.gen(function* () {
+      const seen = yield* Ref.make<ReadonlyArray<string>>([])
+      yield* Effect.forkChild(
+        Stream.runForEach(stream, (a) => Ref.update(seen, (xs) => [...xs, a])),
+      )
+      return seen
+    })
+
+  it.effect("passes elements downstream without waiting for the quiet period", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const seen = yield* collecting(Stream.fromQueue(queue).pipe(onQuiet("1 second", mark)))
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("10 millis")
+      yield* Queue.offer(queue, "b")
+      yield* TestClock.adjust("10 millis")
+
+      // Both are downstream well inside the settle window, and no mark yet.
+      expect(yield* Ref.get(seen)).toEqual(["a", "b"])
+    }),
+  )
+
+  it.effect("emits once the source goes quiet, and not again until it speaks", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const seen = yield* collecting(Stream.fromQueue(queue).pipe(onQuiet("1 second", mark)))
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("2 seconds")
+      expect(yield* Ref.get(seen)).toEqual(["a", "<quiet>"])
+
+      // Still silent: the mark does not repeat while nothing arrives.
+      yield* TestClock.adjust("10 seconds")
+      expect(yield* Ref.get(seen)).toEqual(["a", "<quiet>"])
+
+      // A new element re-arms it.
+      yield* Queue.offer(queue, "b")
+      yield* TestClock.adjust("2 seconds")
+      expect(yield* Ref.get(seen)).toEqual(["a", "<quiet>", "b", "<quiet>"])
+    }),
+  )
+
+  it.effect("marks an element that arrives in the same instant as the mark", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const seen = yield* collecting(Stream.fromQueue(queue).pipe(onQuiet("1 second", mark)))
+
+      yield* Queue.offer(queue, "a")
+      // The window expires exactly here, so "b" races the re-arm.
+      yield* TestClock.adjust("1 second")
+      yield* Queue.offer(queue, "b")
+      yield* TestClock.adjust("2 seconds")
+
+      // Losing that race would leave "b" unmarked until something else spoke.
+      expect(yield* Ref.get(seen)).toEqual(["a", "<quiet>", "b", "<quiet>"])
+    }),
+  )
+
+  it.effect("keeps the window open while elements keep arriving", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const seen = yield* collecting(Stream.fromQueue(queue).pipe(onQuiet("1 second", mark)))
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("900 millis")
+      yield* Queue.offer(queue, "b")
+      yield* TestClock.adjust("900 millis")
+      // 1.8s elapsed, but never 1s of silence.
+      expect(yield* Ref.get(seen)).toEqual(["a", "b"])
+
+      yield* TestClock.adjust("2 seconds")
+      expect(yield* Ref.get(seen)).toEqual(["a", "b", "<quiet>"])
+    }),
+  )
+
+  it.effect("marks each utterance in a run of speech, pauses and more speech", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const seen = yield* collecting(Stream.fromQueue(queue).pipe(onQuiet("1 second", mark)))
+
+      // Deltas of one sentence, arriving faster than the window.
+      for (const token of ["what", " about", " Paris"]) {
+        yield* Queue.offer(queue, token)
+        yield* TestClock.adjust("200 millis")
+      }
+      expect(yield* Ref.get(seen)).toEqual(["what", " about", " Paris"])
+
+      // The speaker stops: one mark, and it does not repeat.
+      yield* TestClock.adjust("3 seconds")
+      // A second sentence re-arms and marks again.
+      for (const token of ["in", " winter"]) {
+        yield* Queue.offer(queue, token)
+        yield* TestClock.adjust("200 millis")
+      }
+      yield* TestClock.adjust("3 seconds")
+
+      expect(yield* Ref.get(seen)).toEqual([
+        "what",
+        " about",
+        " Paris",
+        "<quiet>",
+        "in",
+        " winter",
+        "<quiet>",
+      ])
+    }),
+  )
+
+  it.effect("emits nothing when there is nothing to emit, and waits for the next element", () =>
+    Effect.gen(function* () {
+      const queue = yield* Queue.unbounded<string, Cause.Done>()
+      const emitted = yield* Ref.make(0)
+      // Stands in for an accumulator that is already empty.
+      const emitIfSecond = Effect.map(
+        Ref.updateAndGet(emitted, (n) => n + 1),
+        (n) => (n === 1 ? Option.none<string>() : Option.some("<quiet>")),
+      )
+      const seen = yield* collecting(
+        Stream.fromQueue(queue).pipe(onQuiet("1 second", emitIfSecond)),
+      )
+
+      yield* Queue.offer(queue, "a")
+      yield* TestClock.adjust("3 seconds")
+      // First quiet period asked, got None, and stayed silent.
+      expect(yield* Ref.get(seen)).toEqual(["a"])
+      expect(yield* Ref.get(emitted)).toBe(1)
+
+      // It is disarmed, so silence alone does not ask again.
+      yield* TestClock.adjust("10 seconds")
+      expect(yield* Ref.get(emitted)).toBe(1)
+
+      yield* Queue.offer(queue, "b")
+      yield* TestClock.adjust("3 seconds")
+      expect(yield* Ref.get(seen)).toEqual(["a", "b", "<quiet>"])
+    }),
+  )
+})

package/src/streaming/Settle.ts

@@ -0,0 +1,196 @@
+/**
+ * Acting on a pause in a stream of arrivals.
+ *
+ * A chat user types three lines then stops; a transcriber emits token deltas
+ * then goes quiet. Both want a resetting window: every arrival pushes the
+ * deadline out. Effect's `groupedWithin` measures a fixed window from the first
+ * element of a batch and `debounce` drops all but the last element, so neither
+ * expresses this.
+ *
+ * `settleBurst` withholds elements until a burst ends; `onQuiet` holds nothing
+ * back and only adds an event once things go still.
+ */
+import {
+  Array as Arr,
+  Cause,
+  Clock,
+  Deferred,
+  Duration,
+  Effect,
+  Function,
+  Option,
+  Pull,
+  Queue,
+  Ref,
+  Result,
+  Stream,
+} from "effect"
+
+/** Buffer between a source stream and the batcher. 64 matches the realtime adapters. */
+const DEFAULT_CAPACITY = 64
+
+/**
+ * Block for the first item, then keep taking while the next one lands within
+ * `settle` of the previous. The window resets on every arrival, so a burst
+ * becomes one batch and a lone message followed by silence returns at once.
+ *
+ * A take that loses to the window is interrupted before it removes anything,
+ * so late arrivals stay queued for the next drain. The queue is the caller's:
+ * it owns the backpressure strategy.
+ */
+export const drainBurst = <A, E>(
+  queue: Queue.Dequeue<A, E>,
+  settle: Duration.Input,
+): Effect.Effect<Arr.NonEmptyReadonlyArray<A>, E> => {
+  const settled = (
+    batch: Arr.NonEmptyReadonlyArray<A>,
+  ): Effect.Effect<Arr.NonEmptyReadonlyArray<A>, E> =>
+    Queue.take(queue).pipe(
+      Effect.timeoutOption(settle),
+      Effect.flatMap(
+        Option.match({
+          onNone: () => Effect.succeed(batch),
+          onSome: (item) => settled(Arr.append(batch, item)),
+        }),
+      ),
+      // A queue that ends mid-burst yields the batch in hand. Interrupts and
+      // real failures still propagate.
+      Effect.catchCause((cause) =>
+        Pull.isDoneCause(cause) ? Effect.succeed(batch) : Effect.failCause(cause),
+      ),
+    )
+  return Effect.flatMap(Queue.take(queue), (first) => settled([first]))
+}
+
+export type SettleBurstOptions = {
+  /** Elements buffered ahead of the batcher. A full buffer suspends the source. */
+  readonly capacity?: number
+}
+
+/**
+ * `drainBurst` over a stream: one array per burst, nothing while a burst is
+ * still growing.
+ *
+ * Detecting a pause means reading ahead of the consumer, so this buffers by
+ * construction. The buffer is bounded, so a slow consumer suspends the source
+ * instead of growing memory.
+ */
+export const settleBurst: {
+  (
+    settle: Duration.Input,
+    options?: SettleBurstOptions,
+  ): <A, E, R>(self: Stream.Stream<A, E, R>) => Stream.Stream<ReadonlyArray<A>, E, R>
+  <A, E, R>(
+    self: Stream.Stream<A, E, R>,
+    settle: Duration.Input,
+    options?: SettleBurstOptions,
+  ): Stream.Stream<ReadonlyArray<A>, E, R>
+} = Function.dual(
+  (args) => Stream.isStream(args[0]),
+  <A, E, R>(
+    self: Stream.Stream<A, E, R>,
+    settle: Duration.Input,
+    options?: SettleBurstOptions,
+  ): Stream.Stream<ReadonlyArray<A>, E, R> =>
+    Stream.unwrap(
+      Effect.gen(function* () {
+        // `Cause.Done` as the failure type is what makes `Queue.end` available;
+        // `Queue.shutdown` would clear items still waiting to be batched.
+        const queue = yield* Queue.bounded<A, E | Cause.Done>(options?.capacity ?? DEFAULT_CAPACITY)
+        yield* Stream.runForEach(self, (a) => Queue.offer(queue, a)).pipe(
+          // A failing source must reach the consumer. Ending the queue instead
+          // would present the failure as a clean end of stream.
+          Effect.matchCauseEffect({
+            onFailure: (cause) =>
+              Cause.hasFails(cause) ? Queue.failCause(queue, cause) : Queue.end(queue),
+            onSuccess: () => Queue.end(queue),
+          }),
+          Effect.forkScoped,
+        )
+        return Stream.fromPull(
+          Effect.succeed(Effect.map(drainBurst(queue, settle), (batch) => [batch] as const)),
+        )
+      }),
+    ),
+)
+
+/**
+ * Pass every element through the moment it arrives, and emit whatever `emit`
+ * yields once nothing has arrived for `settle`. At most one emission per quiet
+ * period: the next element re-arms it.
+ *
+ * Nothing is held back, so this suits live output that wants a boundary marker
+ * rather than batching. Ends with the source.
+ */
+export const onQuiet: {
+  <B>(
+    settle: Duration.Input,
+    emit: Effect.Effect<Option.Option<B>>,
+  ): <A, E, R>(self: Stream.Stream<A, E, R>) => Stream.Stream<A | B, E, R>
+  <A, E, R, B>(
+    self: Stream.Stream<A, E, R>,
+    settle: Duration.Input,
+    emit: Effect.Effect<Option.Option<B>>,
+  ): Stream.Stream<A | B, E, R>
+} = Function.dual(
+  3,
+  <A, E, R, B>(
+    self: Stream.Stream<A, E, R>,
+    settle: Duration.Input,
+    emit: Effect.Effect<Option.Option<B>>,
+  ): Stream.Stream<A | B, E, R> =>
+    Stream.unwrap(
+      Effect.gen(function* () {
+        const settleMs = Duration.toMillis(settle)
+        const lastAt = yield* Ref.make(0)
+        // Opened by the first arrival after each emission. An idle stream never
+        // wakes the metronome at all.
+        const gate = yield* Ref.make(yield* Deferred.make<void>())
+
+        const source = self.pipe(
+          Stream.tap(() =>
+            Effect.gen(function* () {
+              // Stamp before opening: the metronome re-reads the stamp after it
+              // swaps the gate, so an arrival that races the swap still counts.
+              const now = yield* Clock.currentTimeMillis
+              yield* Ref.set(lastAt, now)
+              const open = yield* Ref.get(gate)
+              yield* Deferred.succeed(open, undefined)
+            }),
+          ),
+        )
+
+        const quietFor = Effect.map(
+          Effect.all([Clock.currentTimeMillis, Ref.get(lastAt)]),
+          ([now, last]) => now - last,
+        )
+
+        // Waits for the gate, then sleeps out whatever is left of the window.
+        // Never polls, whether the source is busy or idle.
+        const tick = Effect.gen(function* () {
+          yield* Effect.flatMap(Ref.get(gate), Deferred.await)
+          const elapsed = yield* quietFor
+          if (elapsed < settleMs) {
+            yield* Effect.sleep(Duration.millis(settleMs - elapsed))
+            return Option.none<B>()
+          }
+          const next = yield* Deferred.make<void>()
+          yield* Ref.set(gate, next)
+          // An arrival between the check and the swap opened the old gate, so
+          // carry it over rather than sleeping through the burst it started.
+          if ((yield* quietFor) < settleMs) {
+            yield* Deferred.succeed(next, undefined)
+            return Option.none<B>()
+          }
+          return yield* emit
+        })
+
+        const metronome = Stream.fromEffect(tick).pipe(
+          Stream.forever,
+          Stream.filterMap((o) => (Option.isSome(o) ? Result.succeed(o.value) : Result.failVoid)),
+        )
+
+        return source.pipe(Stream.merge(metronome, { haltStrategy: "left" }))
+      }),
+    ),
+)

package/src/streaming/WebSocketSession.test.ts

@@ -0,0 +1,52 @@
+import { describe, it } from "@effect/vitest"
+import * as Socket from "effect/unstable/socket/Socket"
+import { expect } from "vitest"
+import * as AiError from "../domain/AiError.js"
+import { toAiError } from "./WebSocketSession.js"
+
+const KEY = "AIzaSyEXAMPLEKEY"
+const URL = `wss://generativelanguage.googleapis.com/ws/BidiGenerateContent?key=${KEY}`
+
+/**
+ * What Bun hands the `error` listener when the upgrade is rejected: the whole
+ * socket URL, in prose, on a prototype getter. Node and Deno say less, but the
+ * adapter cannot know which runtime it is on.
+ */
+class BunErrorEvent {
+  get message(): string {
+    return `WebSocket connection to '${URL}' failed: Connection ended`
+  }
+}
+
+/** Deno names the status but not the URL, also on a prototype getter. */
+class DenoErrorEvent {
+  constructor(private readonly detail: string) {}
+  get message(): string {
+    return `NetworkError: failed to connect to WebSocket: ${this.detail}`
+  }
+}
+
+const openFailure = (cause: unknown) =>
+  new Socket.SocketError({ reason: new Socket.SocketOpenError({ kind: "Unknown", cause }) })
+
+describe("WebSocketSession errors", () => {
+  it("keeps the key out of the error a rejected upgrade produces", () => {
+    const error = toAiError("google")(openFailure(new BunErrorEvent()))
+
+    expect(JSON.stringify(error)).not.toContain(KEY)
+    expect(AiError.describe(error)).not.toContain(KEY)
+    // The host is still there: an error that says nothing is its own problem.
+    expect(JSON.stringify(error)).toContain("generativelanguage.googleapis.com")
+  })
+
+  it("reads the rejected status out of the prose, wherever the runtime put it", () => {
+    // Deno's wording. The status is the only hint a rejected upgrade gives.
+    const forbidden = toAiError("google")(
+      openFailure(new DenoErrorEvent("Invalid status code: 403")),
+    )
+    const ended = toAiError("google")(openFailure(new BunErrorEvent()))
+
+    expect(forbidden._tag).toBe("AuthFailed")
+    expect(ended._tag).toBe("Unavailable")
+  })
+})