fix(datasets): send HF token, retry only 429/5xx, honor Retry-After for cached file fetches

src/benchmarks/tau-bench-airline/airline.test.ts

@@ -1,6 +1,7 @@
-import { beforeAll, describe, expect, it } from "bun:test";
+import { afterEach, beforeAll, describe, expect, it } from "bun:test";
 
-import { runSync } from "effect/Effect";
+import { FetchHttpClient } from "@effect/platform";
+import { makeSemaphore, provide, runPromise, runSync } from "effect/Effect";
 
 import { MessageRole, ScoreValue } from "../../harness/core";
 import { isRecord } from "../../internal/guards";
@@ -10,7 +11,11 @@ import { TauBenchAirlineConfigSchema } from "../benchmark-config";
 import { benchmarkIds, getBenchmark } from "../registry";
 import { compareActionWithToolCall } from "./action-match";
 import { airlineRecordToSample, TAU_BENCH_AIRLINE_ID } from "./benchmark";
-import { seedAirlineDataCache } from "./environment";
+import {
+  ensureAirlineData,
+  loadAirlineData,
+  seedAirlineDataCache,
+} from "./environment";
 import { evaluateSimulation } from "./evaluator";
 import { airlineScorer } from "./scorer";
 import { AIRLINE_TOOL_DEFINITIONS } from "./tools/definitions";
@@ -447,6 +452,32 @@ describe("evaluateSimulation", () => {
     ).toBe(0);
   });
 });
+describe("ensureAirlineData", () => {
+  const originalFetch = global.fetch;
+
+  afterEach(() => {
+    global.fetch = originalFetch;
+  });
+
+  it("serves the in-process copy without a network request", async () => {
+    let requestCount = 0;
+    const stub: typeof global.fetch = () => {
+      requestCount += 1;
+      return Promise.resolve(new Response("unexpected", { status: 500 }));
+    };
+    global.fetch = stub;
+    seedAirlineDataCache(makeTestData());
+
+    await runPromise(
+      ensureAirlineData(runSync(makeSemaphore(1))).pipe(
+        provide(FetchHttpClient.layer)
+      )
+    );
+
+    expect(requestCount).toBe(0);
+    expect(loadAirlineData()).toEqual(makeTestData());
+  });
+});
 describe("airlineScorer", () => {
   beforeAll(() => {
     seedAirlineDataCache(makeTestData());

src/datasets/cached-file.test.ts

@@ -6,7 +6,11 @@ import { either, provide, runPromise } from "effect/Effect";
 
 import { Either } from "../internal/either";
 import type { CacheStore } from "./cache-store";
-import { fetchCachedTextFile } from "./cached-file";
+import {
+  fetchCachedTextFile,
+  isHuggingFaceUrl,
+  parseRetryAfterMs,
+} from "./cached-file";
 
 function makeMemoryStore(overrides?: Partial<CacheStore>): {
   readonly store: CacheStore;
@@ -35,6 +39,10 @@ const REQUEST = {
   url: "https://example.test/datasets/owner/name/resolve/abc123/db.json",
 } as const;
 
+const HF_REQUEST = {
+  url: "https://huggingface.co/datasets/owner/name/resolve/abc123/db.json",
+} as const;
+
 const CACHE_KEY = `files/${encodeURIComponent(REQUEST.url)}.json`;
 
 function run(
@@ -48,22 +56,26 @@ function run(
 describe("fetchCachedTextFile", () => {
   let originalFetch: typeof global.fetch;
   let requestCount: number;
+  let sentHeaders: Headers[];
 
   beforeEach(() => {
     originalFetch = global.fetch;
     requestCount = 0;
+    sentHeaders = [];
   });
 
   afterEach(() => {
     global.fetch = originalFetch;
   });
 
   function stubFetch(responses: readonly Response[]): void {
-    global.fetch = (() => {
+    const stub: typeof global.fetch = (_input, init) => {
       const response = responses[Math.min(requestCount, responses.length - 1)];
       requestCount += 1;
+      sentHeaders.push(new Headers(init?.headers));
       return Promise.resolve(response.clone());
-    }) as typeof global.fetch;
+    };
+    global.fetch = stub;
   }
 
   it("stores the downloaded body under the url key", async () => {
@@ -132,6 +144,89 @@ describe("fetchCachedTextFile", () => {
     expect(result.left.status).toBe(429);
   });
 
+  it("does not retry a non-retryable status", async () => {
+    stubFetch([
+      new Response("missing", { status: 404 }),
+      new Response("never", { status: 200 }),
+    ]);
+    const { store, entries } = makeMemoryStore();
+
+    const result = await runPromise(
+      fetchCachedTextFile({
+        ...REQUEST,
+        cacheStore: store,
+        retry: { maxRetries: 3, baseDelayMs: 1 },
+      }).pipe(either, provide(FetchHttpClient.layer))
+    );
+
+    assert(Either.isLeft(result));
+    assert(result.left._tag === "CachedFileError");
+    expect(result.left.status).toBe(404);
+    expect(requestCount).toBe(1);
+    expect(entries.size).toBe(0);
+  });
+
+  it("retries a 5xx response", async () => {
+    stubFetch([
+      new Response("upstream", { status: 503 }),
+      new Response("recovered", { status: 200 }),
+    ]);
+    const { store } = makeMemoryStore();
+
+    await expect(
+      run({
+        ...REQUEST,
+        cacheStore: store,
+        retry: { maxRetries: 1, baseDelayMs: 1 },
+      })
+    ).resolves.toBe("recovered");
+    expect(requestCount).toBe(2);
+  });
+
+  it("waits for retry-after before retrying a 429", async () => {
+    stubFetch([
+      new Response("slow down", {
+        status: 429,
+        headers: { "retry-after": "1" },
+      }),
+      new Response("recovered", { status: 200 }),
+    ]);
+    const { store } = makeMemoryStore();
+    const startedAt = performance.now();
+
+    await expect(
+      run({
+        ...REQUEST,
+        cacheStore: store,
+        retry: { maxRetries: 1, baseDelayMs: 1 },
+      })
+    ).resolves.toBe("recovered");
+    expect(performance.now() - startedAt).toBeGreaterThanOrEqual(900);
+    expect(requestCount).toBe(2);
+  });
+
+  it("sends the HF token only to huggingface.co", async () => {
+    stubFetch([new Response("body", { status: 200 })]);
+    const { store } = makeMemoryStore();
+
+    await run({ ...HF_REQUEST, cacheStore: store, hfToken: "hf_test" });
+    await run({ ...REQUEST, cacheStore: store, hfToken: "hf_test" });
+
+    expect(sentHeaders.map((h) => h.get("authorization"))).toEqual([
+      "Bearer hf_test",
+      null,
+    ]);
+  });
+
+  it("sends no authorization header when the token is empty", async () => {
+    stubFetch([new Response("body", { status: 200 })]);
+    const { store } = makeMemoryStore();
+
+    await run({ ...HF_REQUEST, cacheStore: store, hfToken: "" });
+
+    expect(sentHeaders[0]?.get("authorization")).toBeNull();
+  });
+
   it("treats a 300 response as a failure rather than a body", async () => {
     stubFetch([new Response("moved", { status: 300 })]);
     const { store, entries } = makeMemoryStore();
@@ -161,3 +256,27 @@ describe("fetchCachedTextFile", () => {
     await expect(run({ ...REQUEST, cacheStore: store })).resolves.toBe("body");
   });
 });
+
+describe("isHuggingFaceUrl", () => {
+  it("matches huggingface.co and its subdomains only", () => {
+    expect(isHuggingFaceUrl(HF_REQUEST.url)).toBe(true);
+    expect(isHuggingFaceUrl("https://cdn-lfs.huggingface.co/x")).toBe(true);
+    expect(isHuggingFaceUrl("https://nothuggingface.co/x")).toBe(false);
+    expect(isHuggingFaceUrl(REQUEST.url)).toBe(false);
+    expect(isHuggingFaceUrl("not a url")).toBe(false);
+  });
+});
+
+describe("parseRetryAfterMs", () => {
+  it("parses delay seconds and http dates", () => {
+    const now = Date.UTC(2026, 0, 1, 0, 0, 0);
+    expect(parseRetryAfterMs(undefined, now)).toBeUndefined();
+    expect(parseRetryAfterMs("", now)).toBeUndefined();
+    expect(parseRetryAfterMs("  ", now)).toBeUndefined();
+    expect(parseRetryAfterMs("2", now)).toBe(2000);
+    expect(parseRetryAfterMs("-1", now)).toBeUndefined();
+    expect(parseRetryAfterMs("Thu, 01 Jan 2026 00:00:05 GMT", now)).toBe(5000);
+    expect(parseRetryAfterMs("Wed, 31 Dec 2025 23:59:00 GMT", now)).toBe(0);
+    expect(parseRetryAfterMs("soon", now)).toBeUndefined();
+  });
+});

src/datasets/cached-file.ts

@@ -5,38 +5,101 @@ import type { Effect } from "effect/Effect";
 import { fail, gen, ignore, promise, retry, tryPromise } from "effect/Effect";
 
 import { Either } from "../internal/either";
+import { definedValues } from "../internal/guards";
 import { parseSchema, z } from "../internal/zod";
 import type { RetryConfig } from "../runtime/retry";
 import type { CacheStore } from "./cache-store";
 import { resolveCacheStore } from "./cache-store";
-import { hfFetchRetrySchedule } from "./huggingface";
+import { hfFetchRetrySchedule, resolveHfToken } from "./huggingface";
 import { encodeCacheKeySegment } from "./local-cache";
 
 export class CachedFileError extends TaggedError("CachedFileError")<{
   readonly message: string;
   readonly status?: number;
+  readonly retryAfterMs?: number;
 }> {}
 
 export interface CachedTextFileRequest {
   readonly url: string;
   readonly retry?: RetryConfig;
   readonly cacheStore?: CacheStore;
+  readonly hfToken?: string;
 }
 
+type CachedFileFailure = CachedFileError | HttpClientError.HttpClientError;
+
 const CachedTextSchema = z.object({ text: z.string() });
 
+const HF_HOST = "huggingface.co";
+
+export function isHuggingFaceUrl(url: string): boolean {
+  const parsed = Either.try(() => new URL(url));
+  if (Either.isLeft(parsed)) {
+    return false;
+  }
+  const { hostname } = parsed.right;
+  return hostname === HF_HOST || hostname.endsWith(`.${HF_HOST}`);
+}
+
+export function parseRetryAfterMs(
+  value: string | undefined,
+  now: number = Date.now()
+): number | undefined {
+  const normalized = value?.trim();
+  if (normalized === undefined || normalized === "") {
+    return undefined;
+  }
+  const seconds = Number(normalized);
+  if (Number.isFinite(seconds)) {
+    return seconds >= 0 ? seconds * 1e3 : undefined;
+  }
+  const at = Date.parse(normalized);
+  return Number.isFinite(at) ? Math.max(0, at - now) : undefined;
+}
+
+export function isRetryableCachedFileFailure(
+  error: CachedFileFailure
+): boolean {
+  if (error._tag !== "CachedFileError") {
+    return true;
+  }
+  return error.status === 429 || (error.status ?? 0) >= 500;
+}
+
+function cachedFileRetryAfterMs(error: CachedFileFailure): number | undefined {
+  return error._tag === "CachedFileError" ? error.retryAfterMs : undefined;
+}
+
+function authorizationHeaders(
+  url: string,
+  hfToken: string
+): Readonly<Record<string, string>> | undefined {
+  if (hfToken === "" || !isHuggingFaceUrl(url)) {
+    return undefined;
+  }
+  return { Authorization: `Bearer ${hfToken}` };
+}
+
 function download(
   url: string,
+  hfToken: string,
   client: HttpClient.HttpClient
-): Effect<string, CachedFileError | HttpClientError.HttpClientError> {
+): Effect<string, CachedFileFailure> {
   return gen(function* () {
-    const response = yield* client.get(url);
+    const headers = authorizationHeaders(url, hfToken);
+    const response = yield* client.get(
+      url,
+      headers !== undefined ? { headers } : undefined
+    );
     if (response.status < 200 || response.status >= 300) {
       return yield* fail(
-        new CachedFileError({
-          message: `HTTP ${response.status} for ${url}`,
-          status: response.status,
-        })
+        new CachedFileError(
+          definedValues({
+            message: `HTTP ${response.status} for ${url}`,
+            status: response.status,
+            retryAfterMs: parseRetryAfterMs(response.headers["retry-after"]),
+          })
+        )
       );
     }
     return yield* response.text;
@@ -45,11 +108,7 @@ function download(
 
 export function fetchCachedTextFile(
   request: CachedTextFileRequest
-): Effect<
-  string,
-  CachedFileError | HttpClientError.HttpClientError,
-  HttpClient.HttpClient
-> {
+): Effect<string, CachedFileFailure, HttpClient.HttpClient> {
   return gen(function* () {
     const client = yield* HttpClient.HttpClient;
     const store = request.cacheStore ?? resolveCacheStore();
@@ -61,8 +120,15 @@ export function fetchCachedTextFile(
     if (Either.isRight(cached)) {
       return cached.right.text;
     }
-    const text = yield* download(request.url, client).pipe(
-      retry(hfFetchRetrySchedule(request.retry))
+    const hfToken = request.hfToken ?? (yield* resolveHfToken());
+    const text = yield* download(request.url, hfToken, client).pipe(
+      retry(
+        hfFetchRetrySchedule<CachedFileFailure>(
+          request.retry,
+          isRetryableCachedFileFailure,
+          cachedFileRetryAfterMs
+        )
+      )
     );
     yield* tryPromise(() => store.writeJson(key, { text })).pipe(ignore);
     return text;