From 7e82aee466746d7aa576f4b01b7831d7b5f62040 Mon Sep 17 00:00:00 2001 From: James Sterling Date: Thu, 20 Aug 2026 07:09:44 +0000 Subject: [PATCH 1/4] feat(runner): accept caller trace headers for OpenRouter API requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/runner/run-by-id.ts | 13 ++++++---- src/runner/trace-headers.test.ts | 30 +++++++++++++++++++++++ src/runner/trace-headers.ts | 41 ++++++++++++++++++++++++++++++++ 3 files changed, 80 insertions(+), 4 deletions(-) create mode 100644 src/runner/trace-headers.test.ts create mode 100644 src/runner/trace-headers.ts diff --git a/src/runner/run-by-id.ts b/src/runner/run-by-id.ts index a035212..377ca61 100644 --- a/src/runner/run-by-id.ts +++ b/src/runner/run-by-id.ts @@ -1,4 +1,3 @@ -import { FetchHttpClient } from "@effect/platform"; import { flatMap, provide } from "effect/Effect"; import { mergeAll as layerMergeAll, @@ -33,6 +32,7 @@ import { } from "../runtime/generation-resolver"; import { withRunAttempt } from "../runtime/response-cache"; import type { RetryConfig } from "../runtime/retry"; +import { makeHttpClientLayer } from "./trace-headers"; export interface RunBenchmarkInput { readonly benchmarkId: string; @@ -53,6 +53,7 @@ export interface RunBenchmarkInput { readonly abortSignal?: AbortSignal; readonly resultStore?: ResultStoreService; readonly maxOutputTokensCeiling?: number; + readonly traceHeaders?: Readonly>; } export interface RunBenchmarkOutput { @@ -108,9 +109,13 @@ export function runBenchmarkById( }), }, }; - const fullBenchmarkLayer = benchmarkLayer.pipe( - layerProvide(FetchHttpClient.layer) - ); + const httpClientLayer = makeHttpClientLayer({ + ...(input.traceHeaders !== undefined && { + traceHeaders: input.traceHeaders, + }), + ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), + }); + const fullBenchmarkLayer = benchmarkLayer.pipe(layerProvide(httpClientLayer)); const resolverLayer = layerSucceed( GenerationResolver, makeOpenRouterGenerationResolver({ diff --git a/src/runner/trace-headers.test.ts b/src/runner/trace-headers.test.ts new file mode 100644 index 0000000..b7c49bf --- /dev/null +++ b/src/runner/trace-headers.test.ts @@ -0,0 +1,30 @@ +import { describe, expect, it } from "bun:test"; + +import { HttpClientRequest } from "@effect/platform"; + +import { applyTraceHeaders } from "./trace-headers"; + +const API_PREFIX = "https://openrouter.ai/api/v1"; + +const TRACE_HEADERS = { + traceparent: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + "x-benchmark-trace": "test-key", +} as const; + +describe("applyTraceHeaders", () => { + it("adds the headers to a request for the OpenRouter API", () => { + const request = HttpClientRequest.post(`${API_PREFIX}/chat/completions`); + const result = applyTraceHeaders(request, API_PREFIX, TRACE_HEADERS); + expect(result.headers.traceparent).toBe(TRACE_HEADERS.traceparent); + expect(result.headers["x-benchmark-trace"]).toBe("test-key"); + }); + + it("leaves requests to other hosts untouched", () => { + const request = HttpClientRequest.get( + "https://datasets-server.huggingface.co/rows" + ); + const result = applyTraceHeaders(request, API_PREFIX, TRACE_HEADERS); + expect(result.headers.traceparent).toBeUndefined(); + expect(result.headers["x-benchmark-trace"]).toBeUndefined(); + }); +}); diff --git a/src/runner/trace-headers.ts b/src/runner/trace-headers.ts new file mode 100644 index 0000000..240cb2f --- /dev/null +++ b/src/runner/trace-headers.ts @@ -0,0 +1,41 @@ +import { + FetchHttpClient, + HttpClient, + HttpClientRequest, +} from "@effect/platform"; +import { map } from "effect/Effect"; +import type { Layer } from "effect/Layer"; +import { effect as layerEffect, provide as layerProvide } from "effect/Layer"; + +import { normalizeBaseUrl } from "../providers/openrouter-model"; + +const DEFAULT_BASE_URL = "https://openrouter.ai/api/v1"; + +export function applyTraceHeaders( + request: HttpClientRequest.HttpClientRequest, + apiPrefix: string, + traceHeaders: Readonly> +): HttpClientRequest.HttpClientRequest { + return request.url.startsWith(apiPrefix) + ? HttpClientRequest.setHeaders(request, traceHeaders) + : request; +} + +export function makeHttpClientLayer(opts: { + readonly traceHeaders?: Readonly>; + readonly baseUrl?: string; +}): Layer { + const traceHeaders = opts.traceHeaders; + if (traceHeaders === undefined) { + return FetchHttpClient.layer; + } + const apiPrefix = normalizeBaseUrl(opts.baseUrl ?? DEFAULT_BASE_URL); + return layerEffect( + HttpClient.HttpClient, + map(HttpClient.HttpClient, (client) => + HttpClient.mapRequest(client, (request) => + applyTraceHeaders(request, apiPrefix, traceHeaders) + ) + ) + ).pipe(layerProvide(FetchHttpClient.layer)); +} From 595718292c300dccc2356a8c86c6f3b911ba2e0b Mon Sep 17 00:00:00 2001 From: James Sterling Date: Wed, 26 Aug 2026 18:58:22 +0000 Subject: [PATCH 2/4] fix(runner): match trace-header gate on exact URL or path-segment boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/runner/trace-headers.test.ts | 9 +++++++++ src/runner/trace-headers.ts | 2 +- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/runner/trace-headers.test.ts b/src/runner/trace-headers.test.ts index b7c49bf..7606b4d 100644 --- a/src/runner/trace-headers.test.ts +++ b/src/runner/trace-headers.test.ts @@ -19,6 +19,15 @@ describe("applyTraceHeaders", () => { expect(result.headers["x-benchmark-trace"]).toBe("test-key"); }); + it("leaves requests to a lookalike prefix untouched", () => { + const request = HttpClientRequest.post( + "https://openrouter.ai/api/v10/chat/completions" + ); + const result = applyTraceHeaders(request, API_PREFIX, TRACE_HEADERS); + expect(result.headers.traceparent).toBeUndefined(); + expect(result.headers["x-benchmark-trace"]).toBeUndefined(); + }); + it("leaves requests to other hosts untouched", () => { const request = HttpClientRequest.get( "https://datasets-server.huggingface.co/rows" diff --git a/src/runner/trace-headers.ts b/src/runner/trace-headers.ts index 240cb2f..e462d1d 100644 --- a/src/runner/trace-headers.ts +++ b/src/runner/trace-headers.ts @@ -16,7 +16,7 @@ export function applyTraceHeaders( apiPrefix: string, traceHeaders: Readonly> ): HttpClientRequest.HttpClientRequest { - return request.url.startsWith(apiPrefix) + return request.url === apiPrefix || request.url.startsWith(`${apiPrefix}/`) ? HttpClientRequest.setHeaders(request, traceHeaders) : request; } From 4079d1471ba32dd9a1ecc282a65cde53652a1bd1 Mon Sep 17 00:00:00 2001 From: James Sterling Date: Wed, 26 Aug 2026 19:06:06 +0000 Subject: [PATCH 3/4] fix(runner): thread trace headers through the SDK request path and restrict allowed header names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/benchmarks/deep-swe/benchmark.ts | 3 + .../define-single-turn-benchmark.ts | 3 + src/benchmarks/swe-atlas/benchmark.ts | 3 + src/benchmarks/tau-bench-airline/benchmark.ts | 6 ++ .../tau3-bench-banking/benchmark.ts | 6 ++ src/benchmarks/types.ts | 1 + src/benchmarks/wandr/benchmark.ts | 3 + src/providers/openrouter-model.test.ts | 51 ++++++++++++++++ src/providers/openrouter-model.ts | 4 ++ src/providers/responses-client.ts | 4 ++ src/providers/responses-model.ts | 4 ++ src/runner/run-by-id.ts | 17 +++--- src/runner/trace-headers.test.ts | 61 ++++++++++--------- src/runner/trace-headers.ts | 61 ++++++++----------- src/runtime/generation-resolver.ts | 8 ++- 15 files changed, 162 insertions(+), 73 deletions(-) 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 09c07ea..20ec147 100644 --- a/src/runner/run-by-id.ts +++ b/src/runner/run-by-id.ts @@ -1,3 +1,4 @@ +import { FetchHttpClient } from "@effect/platform"; import { flatMap, provide } from "effect/Effect"; import { mergeAll as layerMergeAll, @@ -43,7 +44,7 @@ import { } from "../runtime/generation-resolver"; import { withRunAttempt } from "../runtime/response-cache"; import type { RetryConfig } from "../runtime/retry"; -import { makeHttpClientLayer } from "./trace-headers"; +import { filterTraceHeaders } from "./trace-headers"; export interface RunBenchmarkInput { readonly benchmarkId: string; @@ -106,18 +107,16 @@ export function runBenchmarkById( }), }, }; - const httpClientLayer = makeHttpClientLayer({ - ...(input.traceHeaders !== undefined && { - traceHeaders: input.traceHeaders, - }), - ...(input.baseUrl !== undefined && { baseUrl: input.baseUrl }), - }); - const fullBenchmarkLayer = benchmarkLayer.pipe(layerProvide(httpClientLayer)); + 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( @@ -244,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 index 7606b4d..a5bf03e 100644 --- a/src/runner/trace-headers.test.ts +++ b/src/runner/trace-headers.test.ts @@ -1,39 +1,44 @@ import { describe, expect, it } from "bun:test"; -import { HttpClientRequest } from "@effect/platform"; +import { filterTraceHeaders } from "./trace-headers"; -import { applyTraceHeaders } from "./trace-headers"; +const TRACEPARENT = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; -const API_PREFIX = "https://openrouter.ai/api/v1"; - -const TRACE_HEADERS = { - traceparent: "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", - "x-benchmark-trace": "test-key", -} as const; +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", + }); + }); -describe("applyTraceHeaders", () => { - it("adds the headers to a request for the OpenRouter API", () => { - const request = HttpClientRequest.post(`${API_PREFIX}/chat/completions`); - const result = applyTraceHeaders(request, API_PREFIX, TRACE_HEADERS); - expect(result.headers.traceparent).toBe(TRACE_HEADERS.traceparent); - expect(result.headers["x-benchmark-trace"]).toBe("test-key"); + it("normalizes header names to lowercase", () => { + expect(filterTraceHeaders({ Traceparent: TRACEPARENT })).toEqual({ + traceparent: TRACEPARENT, + }); }); - it("leaves requests to a lookalike prefix untouched", () => { - const request = HttpClientRequest.post( - "https://openrouter.ai/api/v10/chat/completions" - ); - const result = applyTraceHeaders(request, API_PREFIX, TRACE_HEADERS); - expect(result.headers.traceparent).toBeUndefined(); - expect(result.headers["x-benchmark-trace"]).toBeUndefined(); + it("drops disallowed header names such as authorization", () => { + expect( + filterTraceHeaders({ + traceparent: TRACEPARENT, + authorization: "Bearer attacker-key", + "x-session-id": "spoofed", + }) + ).toEqual({ traceparent: TRACEPARENT }); }); - it("leaves requests to other hosts untouched", () => { - const request = HttpClientRequest.get( - "https://datasets-server.huggingface.co/rows" - ); - const result = applyTraceHeaders(request, API_PREFIX, TRACE_HEADERS); - expect(result.headers.traceparent).toBeUndefined(); - expect(result.headers["x-benchmark-trace"]).toBeUndefined(); + 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 index e462d1d..50946d3 100644 --- a/src/runner/trace-headers.ts +++ b/src/runner/trace-headers.ts @@ -1,41 +1,30 @@ -import { - FetchHttpClient, - HttpClient, - HttpClientRequest, -} from "@effect/platform"; -import { map } from "effect/Effect"; -import type { Layer } from "effect/Layer"; -import { effect as layerEffect, provide as layerProvide } from "effect/Layer"; +export const ALLOWED_TRACE_HEADER_NAMES = [ + "traceparent", + "tracestate", + "x-or-traceparent", + "x-benchmark-trace", +] as const; -import { normalizeBaseUrl } from "../providers/openrouter-model"; +export type AllowedTraceHeaderName = + (typeof ALLOWED_TRACE_HEADER_NAMES)[number]; -const DEFAULT_BASE_URL = "https://openrouter.ai/api/v1"; +const ALLOWED_TRACE_HEADER_SET: ReadonlySet = new Set( + ALLOWED_TRACE_HEADER_NAMES +); -export function applyTraceHeaders( - request: HttpClientRequest.HttpClientRequest, - apiPrefix: string, - traceHeaders: Readonly> -): HttpClientRequest.HttpClientRequest { - return request.url === apiPrefix || request.url.startsWith(`${apiPrefix}/`) - ? HttpClientRequest.setHeaders(request, traceHeaders) - : request; -} - -export function makeHttpClientLayer(opts: { - readonly traceHeaders?: Readonly>; - readonly baseUrl?: string; -}): Layer { - const traceHeaders = opts.traceHeaders; - if (traceHeaders === undefined) { - return FetchHttpClient.layer; +export function filterTraceHeaders( + headers: Readonly> | undefined +): Record | undefined { + if (headers === undefined) { + return undefined; } - const apiPrefix = normalizeBaseUrl(opts.baseUrl ?? DEFAULT_BASE_URL); - return layerEffect( - HttpClient.HttpClient, - map(HttpClient.HttpClient, (client) => - HttpClient.mapRequest(client, (request) => - applyTraceHeaders(request, apiPrefix, traceHeaders) - ) - ) - ).pipe(layerProvide(FetchHttpClient.layer)); + 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.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, } ); From ddf045bc9d676012efde32646757fbb8a75824b8 Mon Sep 17 00:00:00 2001 From: James Sterling Date: Wed, 26 Aug 2026 19:24:22 +0000 Subject: [PATCH 4/4] test(runtime): assert generation lookup sends filtered trace headers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- src/runtime/generation-resolver.test.ts | 28 +++++++++++++++++++++++++ 1 file changed, 28 insertions(+) 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" } })