Merge pull request #5178 from stephanmck/bedrock-custom-headers-pr

ai: add custom-header support to Bedrock provider

packages/ai/src/providers/amazon-bedrock.ts

@@ -21,7 +21,7 @@ import {
 	ToolResultStatus,
 } from "@aws-sdk/client-bedrock-runtime";
 import { NodeHttpHandler } from "@smithy/node-http-handler";
-import type { DocumentType } from "@smithy/types";
+import type { BuildMiddleware, DocumentType, MetadataBearer } from "@smithy/types";
 import { calculateCost } from "../models.ts";
 import type {
 	Api,
@@ -182,6 +182,9 @@ export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOpt
 
 		try {
 			const client = new BedrockRuntimeClient(config);
+			if (options.headers && Object.keys(options.headers).length > 0) {
+				addCustomHeadersMiddleware(client, options.headers);
+			}
 			const cacheRetention = resolveCacheRetention(options.cacheRetention);
 			const inferenceMaxTokens = options.maxTokens ?? (isAnthropicClaudeModel(model) ? model.maxTokens : undefined);
 			let commandInput = {
@@ -296,6 +299,42 @@ function formatBedrockError(error: unknown): string {
 	return message;
 }
 
+/**
+ * Header keys that must never be overwritten by caller-supplied headers.
+ * `host` and `x-amz-*` participate in the SigV4 canonical request; `authorization`
+ * is owned by SigV4 or the bearer-token path (config.token + authSchemePreference).
+ * Compared case-insensitively (caller key is lower-cased before lookup).
+ */
+const RESERVED_HEADER_EXACT = new Set(["authorization", "host"]);
+
+function isReservedHeader(key: string): boolean {
+	const lower = key.toLowerCase();
+	return lower.startsWith("x-amz-") || RESERVED_HEADER_EXACT.has(lower);
+}
+
+/**
+ * Attach caller-supplied headers to the outgoing Bedrock request via a Smithy
+ * `build`-step middleware. The `build` step runs after request serialisation but
+ * before SigV4 signing, so injected headers are covered by the signature. Reserved
+ * SigV4 / auth headers (`x-amz-*`, `authorization`, `host`) are silently skipped;
+ * all other caller headers override any existing same-named header on the request.
+ */
+function addCustomHeadersMiddleware(client: BedrockRuntimeClient, headers: Record<string, string>): void {
+	const middleware: BuildMiddleware<object, MetadataBearer> = (next) => async (args) => {
+		const request = args.request;
+		if (request && typeof request === "object" && "headers" in request) {
+			const requestHeaders = (request as { headers: Record<string, string> }).headers;
+			for (const [key, value] of Object.entries(headers)) {
+				if (!isReservedHeader(key)) {
+					requestHeaders[key] = value;
+				}
+			}
+		}
+		return next(args);
+	};
+	client.middlewareStack.add(middleware, { step: "build", name: "pi-ai-custom-headers", priority: "low" });
+}
+
 export const streamSimpleBedrock: StreamFunction<"bedrock-converse-stream", SimpleStreamOptions> = (
 	model: Model<"bedrock-converse-stream">,
 	context: Context,

packages/ai/src/types.ts

@@ -114,8 +114,10 @@ export interface StreamOptions {
 	onResponse?: (response: ProviderResponse, model: Model<Api>) => void | Promise<void>;
 	/**
 	 * Optional custom HTTP headers to include in API requests.
-	 * Merged with provider defaults; can override default headers.
-	 * Not supported by all providers (e.g., AWS Bedrock uses SDK auth).
+	 * Merged with provider defaults; caller values override default headers.
+	 * On AWS Bedrock these are injected via a Smithy `build`-step middleware so
+	 * they are covered by SigV4 signing; reserved headers (`x-amz-*`,
+	 * `authorization`, `host`) are silently ignored to preserve SigV4 / bearer auth.
 	 */
 	headers?: Record<string, string>;
 	/**

packages/ai/test/bedrock-custom-headers.test.ts

@@ -0,0 +1,202 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+type MiddlewareHandler = (next: (args: unknown) => Promise<unknown>) => (args: unknown) => Promise<unknown>;
+
+const bedrockMock = vi.hoisted(() => ({
+	middlewareRegistrations: [] as Array<{
+		handler: MiddlewareHandler;
+		opts: { step?: string; name?: string; priority?: string };
+	}>,
+}));
+
+vi.mock("@aws-sdk/client-bedrock-runtime", () => {
+	class BedrockRuntimeServiceException extends Error {}
+
+	class BedrockRuntimeClient {
+		middlewareStack = {
+			add: (handler: MiddlewareHandler, opts: { step?: string; name?: string; priority?: string }) => {
+				bedrockMock.middlewareRegistrations.push({ handler, opts });
+			},
+		};
+
+		send(): Promise<never> {
+			return Promise.reject(new Error("mock send"));
+		}
+	}
+
+	class ConverseStreamCommand {
+		readonly input: unknown;
+
+		constructor(input: unknown) {
+			this.input = input;
+		}
+	}
+
+	return {
+		BedrockRuntimeClient,
+		BedrockRuntimeServiceException,
+		ConverseStreamCommand,
+		StopReason: {
+			END_TURN: "end_turn",
+			STOP_SEQUENCE: "stop_sequence",
+			MAX_TOKENS: "max_tokens",
+			MODEL_CONTEXT_WINDOW_EXCEEDED: "model_context_window_exceeded",
+			TOOL_USE: "tool_use",
+		},
+		CachePointType: { DEFAULT: "default" },
+		CacheTTL: { ONE_HOUR: "ONE_HOUR" },
+		ConversationRole: { ASSISTANT: "assistant", USER: "user" },
+		ImageFormat: { JPEG: "jpeg", PNG: "png", GIF: "gif", WEBP: "webp" },
+		ToolResultStatus: { ERROR: "error", SUCCESS: "success" },
+	};
+});
+
+import { getModel } from "../src/models.ts";
+import type { BedrockOptions } from "../src/providers/amazon-bedrock.ts";
+import { streamBedrock, streamSimpleBedrock } from "../src/providers/amazon-bedrock.ts";
+import type { Context, Model } from "../src/types.ts";
+
+const context: Context = {
+	messages: [{ role: "user", content: "hello", timestamp: Date.now() }],
+};
+
+const MIDDLEWARE_NAME = "pi-ai-custom-headers";
+
+function getModelFixture(): Model<"bedrock-converse-stream"> {
+	return getModel("amazon-bedrock", "us.anthropic.claude-opus-4-8");
+}
+
+/**
+ * Drive a stream to completion so the middleware (registered before `client.send`)
+ * is captured even though the mocked `send()` rejects. Errors are swallowed because
+ * the rejecting mock is expected — we only care about the recorded registrations.
+ */
+async function driveBedrock(options: BedrockOptions): Promise<void> {
+	await streamBedrock(getModelFixture(), context, options)
+		.result()
+		.catch(() => undefined);
+}
+
+function findCustomHeadersRegistration() {
+	const matches = bedrockMock.middlewareRegistrations.filter((r) => r.opts.name === MIDDLEWARE_NAME);
+	return matches;
+}
+
+beforeEach(() => {
+	bedrockMock.middlewareRegistrations.length = 0;
+});
+
+describe("bedrock custom headers middleware", () => {
+	it("VC1: registers a build-step middleware that injects the caller header (happy path)", async () => {
+		await driveBedrock({ cacheRetention: "none", headers: { "x-custom": "v" } });
+
+		const registrations = findCustomHeadersRegistration();
+		expect(registrations).toHaveLength(1);
+
+		const [reg] = registrations;
+		expect(reg.opts.step).toBe("build");
+		expect(reg.opts.priority).toBe("low");
+		expect(reg.opts.name).toBe(MIDDLEWARE_NAME);
+
+		const nextSpy = vi.fn(async (a: unknown) => a);
+		const fakeArgs = { request: { headers: {} as Record<string, string> } };
+		await reg.handler(nextSpy)(fakeArgs);
+
+		expect(fakeArgs.request.headers["x-custom"]).toBe("v");
+		expect(nextSpy).toHaveBeenCalledTimes(1);
+		expect(nextSpy).toHaveBeenCalledWith(fakeArgs);
+	});
+
+	it("VC2: skips reserved headers case-insensitively while applying allowed ones", async () => {
+		await driveBedrock({
+			cacheRetention: "none",
+			headers: {
+				authorization: "evil",
+				"x-amz-date": "evil",
+				"x-allowed": "ok",
+				Authorization: "evil2",
+				"X-Amz-Date": "evil2",
+				HOST: "evil3",
+			},
+		});
+
+		const [reg] = findCustomHeadersRegistration();
+		expect(reg).toBeDefined();
+
+		const nextSpy = vi.fn(async (a: unknown) => a);
+		const fakeArgs = {
+			request: {
+				headers: {
+					authorization: "real-auth",
+					"x-amz-date": "real-date",
+					host: "real-host",
+				} as Record<string, string>,
+			},
+		};
+		await reg.handler(nextSpy)(fakeArgs);
+
+		expect(fakeArgs.request.headers.authorization).toBe("real-auth");
+		expect(fakeArgs.request.headers["x-amz-date"]).toBe("real-date");
+		expect(fakeArgs.request.headers.host).toBe("real-host");
+		expect(fakeArgs.request.headers["x-allowed"]).toBe("ok");
+		// Mixed-case reserved keys must be skipped too: a case-sensitive guard would
+		// add them back as distinct capitalised keys. Assert no such leak occurred and
+		// that the only new key beyond the three pre-existing ones is `x-allowed`.
+		expect(fakeArgs.request.headers.Authorization).toBeUndefined();
+		expect(fakeArgs.request.headers["X-Amz-Date"]).toBeUndefined();
+		expect(fakeArgs.request.headers.HOST).toBeUndefined();
+		expect(Object.keys(fakeArgs.request.headers).sort()).toEqual(
+			["authorization", "host", "x-allowed", "x-amz-date"].sort(),
+		);
+		expect(nextSpy).toHaveBeenCalledTimes(1);
+	});
+
+	it("VC3: registers no middleware when headers is undefined", async () => {
+		await driveBedrock({ cacheRetention: "none" });
+
+		expect(findCustomHeadersRegistration()).toHaveLength(0);
+	});
+
+	it("VC3: registers no middleware when headers is empty", async () => {
+		await driveBedrock({ cacheRetention: "none", headers: {} });
+
+		expect(findCustomHeadersRegistration()).toHaveLength(0);
+	});
+
+	it("VC3 (structural guard): passes through unchanged when the request has no headers", async () => {
+		await driveBedrock({ cacheRetention: "none", headers: { "x-custom": "v" } });
+
+		const [reg] = findCustomHeadersRegistration();
+		expect(reg).toBeDefined();
+
+		const nextSpy = vi.fn(async (a: unknown) => a);
+
+		const argsNoHeaders = { request: {} };
+		await expect(reg.handler(nextSpy)(argsNoHeaders)).resolves.toBeDefined();
+		expect(nextSpy).toHaveBeenCalledWith(argsNoHeaders);
+
+		const argsUndefinedRequest = { request: undefined };
+		await expect(reg.handler(nextSpy)(argsUndefinedRequest)).resolves.toBeDefined();
+		expect(nextSpy).toHaveBeenCalledWith(argsUndefinedRequest);
+
+		expect(nextSpy).toHaveBeenCalledTimes(2);
+	});
+
+	it("VC4: streamSimpleBedrock forwards headers end-to-end (regression guard)", async () => {
+		await streamSimpleBedrock(getModelFixture(), context, { headers: { "x-custom": "v" } })
+			.result()
+			.catch(() => undefined);
+
+		const registrations = findCustomHeadersRegistration();
+		expect(registrations).toHaveLength(1);
+
+		const [reg] = registrations;
+		expect(reg.opts.step).toBe("build");
+
+		const nextSpy = vi.fn(async (a: unknown) => a);
+		const fakeArgs = { request: { headers: {} as Record<string, string> } };
+		await reg.handler(nextSpy)(fakeArgs);
+
+		expect(fakeArgs.request.headers["x-custom"]).toBe("v");
+	});
+});