@effect-uai/core 0.5.1 → 0.5.2

package/src/retry/Retry.test.ts

@@ -0,0 +1,163 @@
+import { Cause, Effect, Exit, Option, Ref, Schedule, Stream } from "effect"
+import { describe, expect, it } from "vitest"
+import * as AiError from "../domain/AiError.js"
+import * as Retry from "./Retry.js"
+
+// Each call pulls the next plan entry; the last entry repeats forever
+// (so "succeed on attempt N" plans naturally re-yield success after).
+const planned = <A, E>(
+  attempts: Ref.Ref<number>,
+  plan: ReadonlyArray<Stream.Stream<A, E>>,
+): Stream.Stream<A, E> =>
+  Stream.unwrap(
+    Ref.getAndUpdate(attempts, (n) => n + 1).pipe(
+      Effect.map((n) => plan[Math.min(n, plan.length - 1)]!),
+    ),
+  )
+
+const plannedEffect = <A, E>(
+  attempts: Ref.Ref<number>,
+  plan: ReadonlyArray<Effect.Effect<A, E>>,
+): Effect.Effect<A, E> =>
+  Ref.getAndUpdate(attempts, (n) => n + 1).pipe(
+    Effect.flatMap((n) => plan[Math.min(n, plan.length - 1)]!),
+  )
+
+describe("Retry.stream", () => {
+  it("retries on RateLimited and yields success on the next attempt", async () => {
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const s = planned(attempts, [
+        Stream.fail<AiError.AiError>(new AiError.RateLimited({ provider: "mock", raw: null })),
+        Stream.fromIterable(["a", "b"]),
+      ]).pipe(Retry.stream(Schedule.recurs(3)))
+      const out = yield* Stream.runCollect(s)
+      const count = yield* Ref.get(attempts)
+      return { out: Array.from(out), count }
+    })
+
+    const { out, count } = await Effect.runPromise(program)
+    expect(count).toBe(2)
+    expect(out).toEqual(["a", "b"])
+  })
+
+  it("surfaces the underlying retryable failure when retries are exhausted", async () => {
+    const cause = new AiError.Unavailable({ provider: "mock", raw: null })
+    const s = Stream.fail<AiError.AiError>(cause).pipe(Retry.stream(Schedule.recurs(2)))
+
+    const exit = await Effect.runPromise(Effect.exit(Stream.runCollect(s)))
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(cause))
+    }
+  })
+
+  it("bypasses retry for non-retryable AiError (ContentFiltered)", async () => {
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const cause = new AiError.ContentFiltered({ provider: "mock", raw: null })
+      const s = planned(attempts, [Stream.fail<AiError.AiError>(cause)]).pipe(
+        Retry.stream(Schedule.recurs(5)),
+      )
+      const exit = yield* Effect.exit(Stream.runCollect(s))
+      const count = yield* Ref.get(attempts)
+      return { exit, count, cause }
+    })
+
+    const { exit, count, cause } = await Effect.runPromise(program)
+    expect(count).toBe(1)
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(cause))
+    }
+  })
+
+  it("preserves items emitted before a retryable failure (and replays on retry)", async () => {
+    // Documents the 'replays on retry' caveat: the entire stream re-runs.
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const s = planned(attempts, [
+        Stream.concat(
+          Stream.succeed("partial"),
+          Stream.fail<AiError.AiError>(new AiError.Timeout({ provider: "mock", raw: null })),
+        ),
+        Stream.fromIterable(["partial", "done"]),
+      ]).pipe(Retry.stream(Schedule.recurs(1)))
+      const out = yield* Stream.runCollect(s)
+      return Array.from(out)
+    })
+
+    const out = await Effect.runPromise(program)
+    expect(out).toEqual(["partial", "partial", "done"])
+  })
+})
+
+describe("Retry.effect", () => {
+  it("retries on RateLimited and yields success on the next attempt", async () => {
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const value = yield* plannedEffect(attempts, [
+        Effect.fail<AiError.AiError>(new AiError.RateLimited({ provider: "mock", raw: null })),
+        Effect.succeed("ok"),
+      ]).pipe(Retry.effect(Schedule.recurs(3)))
+      const count = yield* Ref.get(attempts)
+      return { value, count }
+    })
+
+    const { value, count } = await Effect.runPromise(program)
+    expect(count).toBe(2)
+    expect(value).toBe("ok")
+  })
+
+  it("surfaces the underlying retryable failure when retries are exhausted", async () => {
+    const cause = new AiError.Unavailable({ provider: "mock", raw: null })
+    const eff = Effect.fail<AiError.AiError>(cause).pipe(Retry.effect(Schedule.recurs(2)))
+
+    const exit = await Effect.runPromise(Effect.exit(eff))
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(cause))
+    }
+  })
+
+  it("bypasses retry for non-retryable AiError (ContentFiltered)", async () => {
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const cause = new AiError.ContentFiltered({ provider: "mock", raw: null })
+      const exit = yield* Effect.exit(
+        plannedEffect(attempts, [Effect.fail<AiError.AiError>(cause)]).pipe(
+          Retry.effect(Schedule.recurs(5)),
+        ),
+      )
+      const count = yield* Ref.get(attempts)
+      return { exit, count, cause }
+    })
+
+    const { exit, count, cause } = await Effect.runPromise(program)
+    expect(count).toBe(1)
+    expect(Exit.isFailure(exit)).toBe(true)
+    if (Exit.isFailure(exit)) {
+      expect(Cause.findErrorOption(exit.cause)).toEqual(Option.some(cause))
+    }
+  })
+
+  it("does NOT replay successful prior work — at-most-once semantics", async () => {
+    // Contrast with Retry.stream: an Effect.retry only re-invokes the
+    // failing computation. No replay of work done before the failure
+    // (there's nothing to replay — Effects are single-shot).
+    const program = Effect.gen(function* () {
+      const attempts = yield* Ref.make(0)
+      const value = yield* plannedEffect(attempts, [
+        Effect.fail<AiError.AiError>(new AiError.Timeout({ provider: "mock", raw: null })),
+        Effect.fail<AiError.AiError>(new AiError.Timeout({ provider: "mock", raw: null })),
+        Effect.succeed("done"),
+      ]).pipe(Retry.effect(Schedule.recurs(5)))
+      const count = yield* Ref.get(attempts)
+      return { value, count }
+    })
+
+    const { value, count } = await Effect.runPromise(program)
+    expect(count).toBe(3)
+    expect(value).toBe("done")
+  })
+})

package/src/retry/Retry.ts

@@ -0,0 +1,84 @@
+import { Data, Effect, type Schedule, Stream } from "effect"
+import type * as AiError from "../domain/AiError.js"
+
+/**
+ * Subset-aware retry helpers for `AiError`.
+ *
+ * Effect's own `Effect.retry` / `Stream.retry` apply the schedule to
+ * every failure. For AI calls that's the wrong policy — replaying a
+ * `ContentFiltered` or `InvalidRequest` won't change the outcome and
+ * just wastes quota. These helpers gate the schedule on the
+ * {@link Retryable} subset (rate limits, transport hiccups, timeouts);
+ * everything else propagates unchanged.
+ *
+ * Two carriers:
+ * - {@link stream} — for `Stream<A, AiError, R>` (e.g. `streamTurn`,
+ *   `streamSynthesis`, `streamTranscriptionFrom`). Caveat: the entire
+ *   stream re-runs on retry, so deltas before the failure replay on
+ *   the next attempt.
+ * - {@link effect} — for `Effect<A, AiError, R>` (e.g. `turn`, `embed`,
+ *   `embedMany`, `synthesize`, `transcribe`). At-most-once.
+ *
+ * Composes with `Stream.retry` / `Effect.retry` directly (no naming
+ * collision since these live under the `Retry` namespace).
+ */
+
+/**
+ * Wrapper around the retryable subset of `AiError`. Internal to the
+ * lift/retry/unlift dance — exported so callers can build their own
+ * variants if needed (e.g. a custom schedule that observes `cause`).
+ */
+export class Retryable extends Data.TaggedError("RetryableAi")<{
+  readonly cause: AiError.RateLimited | AiError.Unavailable | AiError.Timeout
+}> {}
+
+/** Type-narrowing predicate for the retryable subset. */
+export const isRetryable = (
+  e: AiError.AiError,
+): e is AiError.RateLimited | AiError.Unavailable | AiError.Timeout =>
+  e._tag === "RateLimited" || e._tag === "Unavailable" || e._tag === "Timeout"
+
+// Lift events to Items, non-retryable failures to Terminal values
+// (escape retry), retryable failures to wrapped errors (only thing
+// retry sees).
+type Lifted<A> =
+  | { readonly _tag: "Item"; readonly value: A }
+  | { readonly _tag: "Terminal"; readonly cause: AiError.AiError }
+
+/**
+ * Retry a `Stream<A, AiError, R>` on the retryable subset. Other
+ * failures bypass the schedule and propagate unchanged. The whole
+ * stream replays on each attempt — deltas emitted before the failure
+ * are re-emitted on the next try.
+ */
+export const stream =
+  <Out>(schedule: Schedule.Schedule<Out, Retryable>) =>
+  <A, R>(s: Stream.Stream<A, AiError.AiError, R>): Stream.Stream<A, AiError.AiError, R> =>
+    s.pipe(
+      Stream.map((value): Lifted<A> => ({ _tag: "Item", value })),
+      Stream.catchIf(
+        isRetryable,
+        (cause) => Stream.fail(new Retryable({ cause })),
+        (cause) => Stream.succeed<Lifted<A>>({ _tag: "Terminal", cause }),
+      ),
+      Stream.retry(schedule),
+      Stream.catchTag("RetryableAi", (e) => Stream.fail<AiError.AiError>(e.cause)),
+      Stream.flatMap((item) =>
+        item._tag === "Item" ? Stream.succeed(item.value) : Stream.fail(item.cause),
+      ),
+    )
+
+/**
+ * Retry an `Effect<A, AiError, R>` on the retryable subset. Other
+ * failures bypass the schedule and propagate unchanged. At-most-once
+ * — the underlying call only re-runs if the previous attempt failed
+ * with a retryable tag.
+ *
+ * Skips the lift/unlift dance the Stream variant needs: `Effect.retry`
+ * exposes a `while:` predicate so the error channel stays a plain
+ * `AiError` throughout.
+ */
+export const effect =
+  <Out>(schedule: Schedule.Schedule<Out, AiError.AiError>) =>
+  <A, R>(eff: Effect.Effect<A, AiError.AiError, R>): Effect.Effect<A, AiError.AiError, R> =>
+    Effect.retry(eff, { schedule, while: isRetryable })

package/src/testing/MockProvider.ts

@@ -1,7 +1,11 @@
 import { Array as Arr, Duration, Effect, Layer, Match, Option, Ref, Schedule, Stream } from "effect"
 import * as AiError from "../domain/AiError.js"
 import { type Item, isOutputText } from "../domain/Items.js"
-import { LanguageModel, type LanguageModelService } from "../language-model/LanguageModel.js"
+import {
+  LanguageModel,
+  type LanguageModelService,
+  turnFromStream,
+} from "../language-model/LanguageModel.js"
 import { type Turn, TurnEvent } from "../domain/Turn.js"
 
 export type MockOptions = {
@@ -81,8 +85,8 @@ const buildService = (
   options: MockOptions | undefined,
   cursor: Ref.Ref<number>,
   record: (call: Call) => Effect.Effect<void>,
-): LanguageModelService => ({
-  streamTurn: (request) =>
+): LanguageModelService => {
+  const streamTurn: LanguageModelService["streamTurn"] = (request) =>
     Stream.unwrap(
       Ref.getAndUpdate(cursor, (n) => n + 1).pipe(
         Effect.flatMap(
@@ -96,8 +100,9 @@ const buildService = (
             }),
         ),
       ),
-    ),
-})
+    )
+  return { streamTurn, turn: turnFromStream(streamTurn) }
+}
 
 // ---------------------------------------------------------------------------
 // Recorder handle. Unsafe Ref is local: it backs both the `record` write