diff --git a/packages/ai/src/providers/amazon-bedrock.ts b/packages/ai/src/providers/amazon-bedrock.ts index c0da66e2..b5c88bc5 100644 --- a/packages/ai/src/providers/amazon-bedrock.ts +++ b/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): void { + const middleware: BuildMiddleware = (next) => async (args) => { + const request = args.request; + if (request && typeof request === "object" && "headers" in request) { + const requestHeaders = (request as { headers: Record }).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, diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index bf997317..a99d1416 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -114,8 +114,10 @@ export interface StreamOptions { onResponse?: (response: ProviderResponse, model: Model) => void | Promise; /** * 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; /** diff --git a/packages/ai/test/bedrock-custom-headers.test.ts b/packages/ai/test/bedrock-custom-headers.test.ts new file mode 100644 index 00000000..1017d089 --- /dev/null +++ b/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) => (args: unknown) => Promise; + +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 { + 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 { + 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 } }; + 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, + }, + }; + 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 } }; + await reg.handler(nextSpy)(fakeArgs); + + expect(fakeArgs.request.headers["x-custom"]).toBe("v"); + }); +});