diff --git a/src/benchmarks/tau-bench-airline/airline.test.ts b/src/benchmarks/tau-bench-airline/airline.test.ts index 7334f39..d4b6097 100644 --- a/src/benchmarks/tau-bench-airline/airline.test.ts +++ b/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()); diff --git a/src/datasets/cached-file.test.ts b/src/datasets/cached-file.test.ts index b8be009..a2ff625 100644 --- a/src/datasets/cached-file.test.ts +++ b/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): { 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,10 +56,12 @@ function run( describe("fetchCachedTextFile", () => { let originalFetch: typeof global.fetch; let requestCount: number; + let sentHeaders: Headers[]; beforeEach(() => { originalFetch = global.fetch; requestCount = 0; + sentHeaders = []; }); afterEach(() => { @@ -59,11 +69,13 @@ describe("fetchCachedTextFile", () => { }); 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(); + }); +}); diff --git a/src/datasets/cached-file.ts b/src/datasets/cached-file.ts index 612c1f1..931a439 100644 --- a/src/datasets/cached-file.ts +++ b/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> | undefined { + if (hfToken === "" || !isHuggingFaceUrl(url)) { + return undefined; + } + return { Authorization: `Bearer ${hfToken}` }; +} + function download( url: string, + hfToken: string, client: HttpClient.HttpClient -): Effect { +): Effect { 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 { 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( + request.retry, + isRetryableCachedFileFailure, + cachedFileRetryAfterMs + ) + ) ); yield* tryPromise(() => store.writeJson(key, { text })).pipe(ignore); return text; diff --git a/src/datasets/huggingface.ts b/src/datasets/huggingface.ts index ed033a7..0c732dc 100644 --- a/src/datasets/huggingface.ts +++ b/src/datasets/huggingface.ts @@ -4,6 +4,7 @@ import { join } from "node:path"; import { FetchHttpClient, HttpClient } from "@effect/platform"; import { fromIterable } from "effect/Chunk"; import { map as configMap, option, string } from "effect/Config"; +import type { DurationInput } from "effect/Duration"; import type { Effect } from "effect/Effect"; import { fail, @@ -24,6 +25,7 @@ import type { Schedule } from "effect/Schedule"; import { exponential, jittered, + modifyDelay, passthrough, whileInput, } from "effect/Schedule"; @@ -92,7 +94,8 @@ function hfPageCacheKey( export function hfFetchRetrySchedule( config: RetryConfig = {}, - isRetryable: (error: E) => boolean = () => true + isRetryable: (error: E) => boolean = () => true, + retryAfterMs: (error: E) => number | undefined = () => undefined ): Schedule< { readonly error: E; @@ -107,7 +110,13 @@ export function hfFetchRetrySchedule( passthrough ); return withRetryAttemptLogging( - whileInput(scheduled, isRetryable).pipe(passthrough), + whileInput(scheduled, isRetryable).pipe( + passthrough, + modifyDelay((error, computed): DurationInput => { + const explicit = retryAfterMs(error); + return explicit !== undefined ? `${explicit} millis` : computed; + }) + ), maxRetries ); }