diff --git a/src/benchmarks/deep-swe/benchmark.ts b/src/benchmarks/deep-swe/benchmark.ts index 03099ab..0a1bd07 100644 --- a/src/benchmarks/deep-swe/benchmark.ts +++ b/src/benchmarks/deep-swe/benchmark.ts @@ -51,6 +51,9 @@ function makeDeepSweLayer( ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), sessionId: input.sessionId, ...(input.modelRetry !== undefined && { retry: input.modelRetry }), + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), }); const sandboxLayer = makeModalSandboxLayer({ appName: "openrouter-deep-swe", diff --git a/src/benchmarks/define-single-turn-benchmark.ts b/src/benchmarks/define-single-turn-benchmark.ts index 5fd5d53..f513c8d 100644 --- a/src/benchmarks/define-single-turn-benchmark.ts +++ b/src/benchmarks/define-single-turn-benchmark.ts @@ -68,6 +68,9 @@ export function defineSingleTurnBenchmark< ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), sessionId: input.sessionId, ...(input.modelRetry !== undefined && { retry: input.modelRetry }), + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), }); const solverLayer = layerEffect(Solver)( gen(function* () { diff --git a/src/benchmarks/swe-atlas/benchmark.ts b/src/benchmarks/swe-atlas/benchmark.ts index 39cdcbe..d9cf9f2 100644 --- a/src/benchmarks/swe-atlas/benchmark.ts +++ b/src/benchmarks/swe-atlas/benchmark.ts @@ -59,6 +59,9 @@ function makeSweAtlasLayer( ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), sessionId: input.sessionId, ...(input.modelRetry !== undefined && { retry: input.modelRetry }), + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), }); const sandboxLayer = makeModalSandboxLayer({ appName: "openrouter-swe-atlas", diff --git a/src/benchmarks/tau-bench-airline/benchmark.ts b/src/benchmarks/tau-bench-airline/benchmark.ts index 3b7646f..d8f560c 100644 --- a/src/benchmarks/tau-bench-airline/benchmark.ts +++ b/src/benchmarks/tau-bench-airline/benchmark.ts @@ -126,12 +126,18 @@ function makeAirlineLayer( ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), sessionId: input.sessionId, ...(input.modelRetry !== undefined && { retry: input.modelRetry }), + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), }); const userModelLayer = makeResponsesModelLayer({ model: benchmarkConfig.userModel, apiKey: input.apiKey, ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), sessionId: input.sessionId, + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), }); const solverLayer = layerEffect(Solver)( gen(function* () { diff --git a/src/benchmarks/tau3-bench-banking/benchmark.ts b/src/benchmarks/tau3-bench-banking/benchmark.ts index b98ec0e..f84432f 100644 --- a/src/benchmarks/tau3-bench-banking/benchmark.ts +++ b/src/benchmarks/tau3-bench-banking/benchmark.ts @@ -81,12 +81,18 @@ function makeBankingLayer( ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), sessionId: input.sessionId, ...(input.modelRetry !== undefined && { retry: input.modelRetry }), + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), }); const userModelLayer = makeResponsesModelLayer({ model: config.userModel, apiKey: input.apiKey, ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), sessionId: input.sessionId, + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), }); const solverLayer = layerEffect(Solver)( gen(function* () { diff --git a/src/benchmarks/types.ts b/src/benchmarks/types.ts index 229d070..a14cd6c 100644 --- a/src/benchmarks/types.ts +++ b/src/benchmarks/types.ts @@ -23,6 +23,7 @@ export interface BenchmarkRunInput< readonly datasetRetry?: RetryConfig; readonly modelRetry?: RetryConfig; readonly maxOutputTokensCeiling?: number; + readonly traceHeaders?: Readonly>; readonly modelLayer?: Layer; readonly responsesModelLayer?: Layer< ResponsesModel, diff --git a/src/benchmarks/wandr/benchmark.ts b/src/benchmarks/wandr/benchmark.ts index 57648e8..b44a5b0 100644 --- a/src/benchmarks/wandr/benchmark.ts +++ b/src/benchmarks/wandr/benchmark.ts @@ -69,6 +69,9 @@ function makeWandrLayer( ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), sessionId: input.sessionId, ...(input.modelRetry !== undefined && { retry: input.modelRetry }), + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), }); const sandboxLayer = makeModalSandboxLayer({ appName: "openrouter-wandr", diff --git a/src/providers/openrouter-model.test.ts b/src/providers/openrouter-model.test.ts index ffb85e5..4238cb2 100644 --- a/src/providers/openrouter-model.test.ts +++ b/src/providers/openrouter-model.test.ts @@ -143,4 +143,55 @@ describe("openrouter-model", () => { }); expect(exit.value.rawResponse).toBeUndefined(); }); + + it("sends trace headers on the OpenRouter request and never overrides auth", async () => { + const stream = await readFile( + new URL( + "../../test/fixtures/advisor-responses-stream.sse", + import.meta.url + ), + "utf8" + ); + const originalFetch = globalThis.fetch; + let request: Request | undefined; + globalThis.fetch = async (input, init) => { + request = input instanceof Request ? input : new Request(input, init); + return new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + }; + restore = () => { + globalThis.fetch = originalFetch; + }; + + const traceparent = + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + const exit = await runPromiseExit( + gen(function* () { + const model = yield* Model; + return yield* model.generate([{ role: "user", content: "question" }], { + temperature: 0, + }); + }).pipe( + provide( + makeOpenRouterModelLayer({ + model: "openai/gpt-5", + apiKey: "sk-test", + baseUrl: "https://example.test", + traceHeaders: { + traceparent, + "x-benchmark-trace": "test-trace-key", + authorization: "Bearer attacker-key", + }, + }).pipe(layerProvide(FetchHttpClient.layer)) + ) + ) + ); + assertSuccess(exit); + expect(request?.url).toBe("https://example.test/api/v1/responses"); + expect(request?.headers.get("traceparent")).toBe(traceparent); + expect(request?.headers.get("x-benchmark-trace")).toBe("test-trace-key"); + expect(request?.headers.get("authorization")).toBe("Bearer sk-test"); + }); }); diff --git a/src/providers/openrouter-model.ts b/src/providers/openrouter-model.ts index 12624ce..bf94339 100644 --- a/src/providers/openrouter-model.ts +++ b/src/providers/openrouter-model.ts @@ -22,6 +22,7 @@ export interface OpenRouterModelConfig { readonly baseUrl?: string; readonly sessionId?: string; readonly retry?: RetryConfig; + readonly traceHeaders?: Readonly>; } export function normalizeBaseUrl(baseUrl: string): string { @@ -40,6 +41,9 @@ export function makeOpenRouterModelLayer( }), ...(config.sessionId !== undefined && { sessionId: config.sessionId }), ...(config.retry !== undefined && { retry: config.retry }), + ...(config.traceHeaders !== undefined && { + traceHeaders: config.traceHeaders, + }), }); return effect(Model)( gen(function* () { diff --git a/src/providers/responses-client.ts b/src/providers/responses-client.ts index 4d71787..4d83eda 100644 --- a/src/providers/responses-client.ts +++ b/src/providers/responses-client.ts @@ -27,6 +27,7 @@ import { ModelError } from "../harness/core"; import { Either } from "../internal/either"; import { isRecord } from "../internal/guards"; import { parseSchema, z } from "../internal/zod"; +import { filterTraceHeaders } from "../runner/trace-headers"; import { recordGenerationId } from "../runtime/generation-ids"; import type { ResponseCacheAttemptState } from "../runtime/response-cache"; import { @@ -103,6 +104,7 @@ export interface ResponsesConfig { readonly apiKey: string; readonly baseUrl?: string; readonly sessionId?: string; + readonly traceHeaders?: Readonly>; } export const VERSION_OVERRIDE_HEADER = @@ -154,6 +156,7 @@ function normalizeBaseUrl(baseUrl: string): string { } export function makeResponsesLayer(config: ResponsesConfig): Layer { + const traceHeaders = filterTraceHeaders(config.traceHeaders); const send = ( body: ResponsesRequest, options: ResponsesSendOptions, @@ -188,6 +191,7 @@ export function makeResponsesLayer(config: ResponsesConfig): Layer { const headers: Record = { "HTTP-Referer": BENCH_HARNESS_APP_REFERRER, "X-OpenRouter-Title": BENCH_HARNESS_APP_TITLE, + ...traceHeaders, ...options.extraHeaders, ...(options.versionOverride ? { [VERSION_OVERRIDE_HEADER]: `api="${options.versionOverride}"` } diff --git a/src/providers/responses-model.ts b/src/providers/responses-model.ts index 9dbb174..815d941 100644 --- a/src/providers/responses-model.ts +++ b/src/providers/responses-model.ts @@ -87,6 +87,7 @@ export interface ResponsesModelConfig { readonly baseUrl?: string; readonly sessionId?: string; readonly retry?: RetryConfig; + readonly traceHeaders?: Readonly>; } export interface ResponsesModelService { @@ -111,6 +112,9 @@ export function makeResponsesModelLayer( apiKey: config.apiKey, baseUrl, ...(config.sessionId !== undefined && { sessionId: config.sessionId }), + ...(config.traceHeaders !== undefined && { + traceHeaders: config.traceHeaders, + }), }); return effect(ResponsesModel)( gen(function* () { diff --git a/src/runner/run-by-id.ts b/src/runner/run-by-id.ts index 59308ac..20ec147 100644 --- a/src/runner/run-by-id.ts +++ b/src/runner/run-by-id.ts @@ -44,6 +44,7 @@ import { } from "../runtime/generation-resolver"; import { withRunAttempt } from "../runtime/response-cache"; import type { RetryConfig } from "../runtime/retry"; +import { filterTraceHeaders } from "./trace-headers"; export interface RunBenchmarkInput { readonly benchmarkId: string; @@ -65,6 +66,7 @@ export interface RunBenchmarkInput { readonly abortSignal?: AbortSignal; readonly resultStore?: ResultStoreService; readonly maxOutputTokensCeiling?: number; + readonly traceHeaders?: Readonly>; } export interface RunBenchmarkOutput { @@ -108,11 +110,13 @@ export function runBenchmarkById( const fullBenchmarkLayer = benchmarkLayer.pipe( layerProvide(FetchHttpClient.layer) ); + const traceHeaders = filterTraceHeaders(input.traceHeaders); const resolverLayer = layerSucceed( GenerationResolver, makeOpenRouterGenerationResolver({ apiKey: input.apiKey, ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), + ...(traceHeaders !== undefined && { traceHeaders }), }) ); const layers = layerMergeAll( @@ -239,10 +243,12 @@ function makeBenchmarkLayer( benchmarkConfig: Config ): ReturnType { const maxRetries = benchmarkConfig.maxRetries; + const traceHeaders = filterTraceHeaders(input.traceHeaders); const benchmarkInput: BenchmarkRunInput = { apiKey: input.apiKey, benchmarkConfig, ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), + ...(traceHeaders !== undefined && { traceHeaders }), sessionId: input.sessionId, ...(input.datasetRetry !== undefined && { datasetRetry: input.datasetRetry, diff --git a/src/runner/trace-headers.test.ts b/src/runner/trace-headers.test.ts new file mode 100644 index 0000000..a5bf03e --- /dev/null +++ b/src/runner/trace-headers.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from "bun:test"; + +import { filterTraceHeaders } from "./trace-headers"; + +const TRACEPARENT = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + +describe("filterTraceHeaders", () => { + it("keeps the allowed trace header names", () => { + expect( + filterTraceHeaders({ + traceparent: TRACEPARENT, + tracestate: "or=bench", + "x-or-traceparent": TRACEPARENT, + "x-benchmark-trace": "test-key", + }) + ).toEqual({ + traceparent: TRACEPARENT, + tracestate: "or=bench", + "x-or-traceparent": TRACEPARENT, + "x-benchmark-trace": "test-key", + }); + }); + + it("normalizes header names to lowercase", () => { + expect(filterTraceHeaders({ Traceparent: TRACEPARENT })).toEqual({ + traceparent: TRACEPARENT, + }); + }); + + it("drops disallowed header names such as authorization", () => { + expect( + filterTraceHeaders({ + traceparent: TRACEPARENT, + authorization: "Bearer attacker-key", + "x-session-id": "spoofed", + }) + ).toEqual({ traceparent: TRACEPARENT }); + }); + + it("returns undefined when nothing survives the filter", () => { + expect(filterTraceHeaders({ authorization: "Bearer x" })).toBeUndefined(); + expect(filterTraceHeaders(undefined)).toBeUndefined(); + }); +}); diff --git a/src/runner/trace-headers.ts b/src/runner/trace-headers.ts new file mode 100644 index 0000000..50946d3 --- /dev/null +++ b/src/runner/trace-headers.ts @@ -0,0 +1,30 @@ +export const ALLOWED_TRACE_HEADER_NAMES = [ + "traceparent", + "tracestate", + "x-or-traceparent", + "x-benchmark-trace", +] as const; + +export type AllowedTraceHeaderName = + (typeof ALLOWED_TRACE_HEADER_NAMES)[number]; + +const ALLOWED_TRACE_HEADER_SET: ReadonlySet = new Set( + ALLOWED_TRACE_HEADER_NAMES +); + +export function filterTraceHeaders( + headers: Readonly> | undefined +): Record | undefined { + if (headers === undefined) { + return undefined; + } + const filtered = Object.fromEntries( + Object.entries(headers) + .map(([name, value]): readonly [string, string] => [ + name.toLowerCase(), + value, + ]) + .filter(([name]) => ALLOWED_TRACE_HEADER_SET.has(name)) + ); + return Object.keys(filtered).length > 0 ? filtered : undefined; +} diff --git a/src/runtime/generation-resolver.test.ts b/src/runtime/generation-resolver.test.ts index 6c2a15c..c6e09ae 100644 --- a/src/runtime/generation-resolver.test.ts +++ b/src/runtime/generation-resolver.test.ts @@ -227,6 +227,34 @@ describe("makeOpenRouterGenerationResolver", () => { "https://example.com/api/v1/generation?id=gen-original", ]); }); + it("sends filtered trace headers on the lookup without touching auth", async () => { + let request: Request | undefined; + globalThis.fetch = ((input: string | URL | Request, init?: RequestInit) => { + request = input instanceof Request ? input : new Request(input, init); + return Promise.resolve( + jsonResponse({ data: { response_cache_source_id: "gen-original" } }) + ); + }) as typeof fetch; + const traceparent = + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + const resolver = makeOpenRouterGenerationResolver({ + apiKey: "test-key", + baseUrl: "https://example.com", + pollIntervalMs: 1, + maxAttempts: 1, + traceHeaders: { + traceparent, + "x-benchmark-trace": "trace-key", + authorization: "Bearer attacker-key", + }, + }); + await runPromise( + resolver.resolveSourceGeneration("gen-dummy", { includeUsage: false }) + ); + expect(request?.headers.get("traceparent")).toBe(traceparent); + expect(request?.headers.get("x-benchmark-trace")).toBe("trace-key"); + expect(request?.headers.get("authorization")).toBe("Bearer test-key"); + }); it("skips the source usage lookup when includeUsage is false", async () => { const getCalls = mockFetch(() => jsonResponse({ data: { response_cache_source_id: "gen-original" } }) diff --git a/src/runtime/generation-resolver.ts b/src/runtime/generation-resolver.ts index aeaff1f..5c2285e 100644 --- a/src/runtime/generation-resolver.ts +++ b/src/runtime/generation-resolver.ts @@ -22,6 +22,7 @@ import { unknownErrorToString } from "../internal/errors"; import { isDefinedAndNotNull } from "../internal/guards"; import { wLog } from "../internal/log"; import { parseSchema, z } from "../internal/zod"; +import { filterTraceHeaders } from "../runner/trace-headers"; import type { GenerationIdEntry } from "./generation-ids"; import { getCollectedGenerationIdEntries } from "./generation-ids"; @@ -77,6 +78,7 @@ export interface GenerationResolverConfig { readonly baseUrl?: string; readonly pollIntervalMs?: number; readonly maxAttempts?: number; + readonly traceHeaders?: Readonly>; } const DEFAULT_POLL_INTERVAL_MS = 5000; @@ -117,6 +119,7 @@ export function makeOpenRouterGenerationResolver( config: GenerationResolverConfig ): GenerationResolverService { const baseUrl = normalizeBaseUrl(config.baseUrl ?? "https://openrouter.ai"); + const traceHeaders = filterTraceHeaders(config.traceHeaders); const pollIntervalMs = config.pollIntervalMs ?? DEFAULT_POLL_INTERVAL_MS; const maxAttempts = config.maxAttempts ?? DEFAULT_MAX_ATTEMPTS; const pollSchedule = (attempts: number) => @@ -134,7 +137,10 @@ export function makeOpenRouterGenerationResolver( const response = await fetch( `${baseUrl}/generation?id=${encodeURIComponent(generationId)}`, { - headers: { Authorization: `Bearer ${config.apiKey}` }, + headers: { + ...traceHeaders, + Authorization: `Bearer ${config.apiKey}`, + }, signal, } );