fix(ai): use SDK token auth and omit Bedrock display in GovCloud closes #3359

This commit is contained in:
Mario Zechner
2026-04-18 01:24:13 +02:00
parent 2dddc5ba25
commit 454b9619c7
3 changed files with 55 additions and 51 deletions

View File

@@ -4,6 +4,7 @@
### Fixed ### Fixed
- Fixed Bedrock bearer-token authentication to use the SDK's native token auth path and omit Claude `thinking.display` for GovCloud targets, avoiding duplicate `Authorization` headers and GovCloud Converse validation errors ([#3359](https://github.com/badlogic/pi-mono/issues/3359))
- Fixed direct Mistral tool definitions to strip TypeBox symbol metadata before passing schemas to the SDK, restoring tool calls after the SDK's stricter outbound validation ([#3361](https://github.com/badlogic/pi-mono/issues/3361)) - Fixed direct Mistral tool definitions to strip TypeBox symbol metadata before passing schemas to the SDK, restoring tool calls after the SDK's stricter outbound validation ([#3361](https://github.com/badlogic/pi-mono/issues/3361))
## [0.67.67] - 2026-04-17 ## [0.67.67] - 2026-04-17

View File

@@ -15,15 +15,11 @@ import {
type ConverseStreamMetadataEvent, type ConverseStreamMetadataEvent,
ImageFormat, ImageFormat,
type Message, type Message,
type ServiceInputTypes,
type ServiceOutputTypes,
type SystemContentBlock, type SystemContentBlock,
type ToolChoice, type ToolChoice,
type ToolConfiguration, type ToolConfiguration,
ToolResultStatus, ToolResultStatus,
} from "@aws-sdk/client-bedrock-runtime"; } from "@aws-sdk/client-bedrock-runtime";
import type { FinalizeRequestMiddleware } from "@smithy/types";
import { calculateCost } from "../models.js"; import { calculateCost } from "../models.js";
import type { import type {
Api, Api,
@@ -119,8 +115,9 @@ export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOpt
profile: options.profile, profile: options.profile,
}; };
// Resolve bearer token for API key auth (bypasses SigV4) // Resolve bearer token for Bedrock API key auth.
const bearerToken = options.bearerToken || process.env.AWS_BEARER_TOKEN_BEDROCK || undefined; const bearerToken = options.bearerToken || process.env.AWS_BEARER_TOKEN_BEDROCK || undefined;
const useBearerToken = bearerToken !== undefined && process.env.AWS_BEDROCK_SKIP_AUTH !== "1";
// in Node.js/Bun environment only // in Node.js/Bun environment only
if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) { if (typeof process !== "undefined" && (process.versions?.node || process.versions?.bun)) {
@@ -142,15 +139,6 @@ export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOpt
}; };
} }
// Bearer token auth: use API key instead of SigV4 signing.
// Requires bedrock:CallWithBearerToken IAM permission.
if (bearerToken && process.env.AWS_BEDROCK_SKIP_AUTH !== "1") {
config.credentials = {
accessKeyId: "bearer-token-auth",
secretAccessKey: "bearer-token-auth",
};
}
if ( if (
process.env.HTTP_PROXY || process.env.HTTP_PROXY ||
process.env.HTTPS_PROXY || process.env.HTTPS_PROXY ||
@@ -182,37 +170,13 @@ export const streamBedrock: StreamFunction<"bedrock-converse-stream", BedrockOpt
config.region = options.region || "us-east-1"; config.region = options.region || "us-east-1";
} }
if (useBearerToken) {
config.token = { token: bearerToken };
config.authSchemePreference = ["httpBearerAuth"];
}
try { try {
const client = new BedrockRuntimeClient(config); const client = new BedrockRuntimeClient(config);
// Inject bearer token middleware after SigV4 signing
if (bearerToken) {
const bearerTokenAuthMiddleware: FinalizeRequestMiddleware<ServiceInputTypes, ServiceOutputTypes> =
(next) => async (args) => {
const request = args.request;
if (
typeof request === "object" &&
request !== null &&
"headers" in request &&
typeof request.headers === "object" &&
request.headers !== null
) {
const headers = request.headers as Record<string, string>;
headers.authorization = `Bearer ${bearerToken}`;
delete headers["x-amz-date"];
delete headers["x-amz-security-token"];
delete headers["x-amz-content-sha256"];
}
return next(args);
};
client.middlewareStack.addRelativeTo(bearerTokenAuthMiddleware, {
relation: "after",
toMiddleware: "awsAuthMiddleware",
name: "bearerTokenAuth",
});
}
const cacheRetention = resolveCacheRetention(options.cacheRetention); const cacheRetention = resolveCacheRetention(options.cacheRetention);
let commandInput = { let commandInput = {
modelId: model.id, modelId: model.id,
@@ -812,6 +776,16 @@ function mapStopReason(reason: string | undefined): StopReason {
} }
} }
function isGovCloudBedrockTarget(model: Model<"bedrock-converse-stream">, options: BedrockOptions): boolean {
const region = options.region || process.env.AWS_REGION || process.env.AWS_DEFAULT_REGION;
if (region?.toLowerCase().startsWith("us-gov-")) {
return true;
}
const modelId = model.id.toLowerCase();
return modelId.startsWith("us-gov.") || modelId.startsWith("arn:aws-us-gov:");
}
function buildAdditionalModelRequestFields( function buildAdditionalModelRequestFields(
model: Model<"bedrock-converse-stream">, model: Model<"bedrock-converse-stream">,
options: BedrockOptions, options: BedrockOptions,
@@ -821,12 +795,12 @@ function buildAdditionalModelRequestFields(
} }
if (model.id.includes("anthropic.claude") || model.id.includes("anthropic/claude")) { if (model.id.includes("anthropic.claude") || model.id.includes("anthropic/claude")) {
// Default to "summarized" so Opus 4.7 and Mythos Preview behave like // GovCloud Bedrock currently rejects the Claude thinking.display field.
// older Claude 4 models (whose API default is also "summarized"). // Omit it there until the GovCloud Converse schema catches up.
const display: BedrockThinkingDisplay = options.thinkingDisplay ?? "summarized"; const display = isGovCloudBedrockTarget(model, options) ? undefined : (options.thinkingDisplay ?? "summarized");
const result: Record<string, any> = supportsAdaptiveThinking(model.id) const result: Record<string, any> = supportsAdaptiveThinking(model.id)
? { ? {
thinking: { type: "adaptive", display }, thinking: { type: "adaptive", ...(display !== undefined ? { display } : {}) },
output_config: { effort: mapThinkingLevelToEffort(options.reasoning, model.id) }, output_config: { effort: mapThinkingLevelToEffort(options.reasoning, model.id) },
} }
: (() => { : (() => {
@@ -846,7 +820,7 @@ function buildAdditionalModelRequestFields(
thinking: { thinking: {
type: "enabled", type: "enabled",
budget_tokens: budget, budget_tokens: budget,
display, ...(display !== undefined ? { display } : {}),
}, },
}; };
})(); })();

View File

@@ -1,7 +1,7 @@
import { describe, expect, it } from "vitest"; import { describe, expect, it } from "vitest";
import { getModel } from "../src/models.js"; import { getModel } from "../src/models.js";
import { streamBedrock } from "../src/providers/amazon-bedrock.js"; import { type BedrockOptions, streamBedrock } from "../src/providers/amazon-bedrock.js";
import type { Context, Model, SimpleStreamOptions } from "../src/types.js"; import type { Context, Model } from "../src/types.js";
interface BedrockThinkingPayload { interface BedrockThinkingPayload {
additionalModelRequestFields?: { additionalModelRequestFields?: {
@@ -19,7 +19,7 @@ function makeContext(): Context {
async function capturePayload( async function capturePayload(
model: Model<"bedrock-converse-stream">, model: Model<"bedrock-converse-stream">,
options?: SimpleStreamOptions, options?: BedrockOptions,
): Promise<BedrockThinkingPayload> { ): Promise<BedrockThinkingPayload> {
let capturedPayload: BedrockThinkingPayload | undefined; let capturedPayload: BedrockThinkingPayload | undefined;
const s = streamBedrock(model, makeContext(), { const s = streamBedrock(model, makeContext(), {
@@ -75,4 +75,33 @@ describe("Bedrock thinking payload", () => {
expect(payload.additionalModelRequestFields?.output_config).toEqual({ effort: "xhigh" }); expect(payload.additionalModelRequestFields?.output_config).toEqual({ effort: "xhigh" });
expect(payload.additionalModelRequestFields?.anthropic_beta).toBeUndefined(); expect(payload.additionalModelRequestFields?.anthropic_beta).toBeUndefined();
}); });
it("omits display for GovCloud model ids on non-adaptive Claude thinking", async () => {
const baseModel = getModel("amazon-bedrock", "us.anthropic.claude-sonnet-4-5-20250929-v1:0");
const model: Model<"bedrock-converse-stream"> = {
...baseModel,
id: "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0",
name: "Claude Sonnet 4.5 (GovCloud)",
};
const payload = await capturePayload(model);
expect(payload.additionalModelRequestFields?.thinking).toEqual({ type: "enabled", budget_tokens: 16384 });
expect(payload.additionalModelRequestFields?.anthropic_beta).toEqual(["interleaved-thinking-2025-05-14"]);
});
it("omits display for GovCloud regions on adaptive Claude thinking", async () => {
const baseModel = getModel("amazon-bedrock", "global.anthropic.claude-opus-4-6-v1");
const model: Model<"bedrock-converse-stream"> = {
...baseModel,
id: "global.anthropic.claude-opus-4-7-v1",
name: "Claude Opus 4.7 (Global)",
};
const payload = await capturePayload(model, { region: "us-gov-west-1" });
expect(payload.additionalModelRequestFields?.thinking).toEqual({ type: "adaptive" });
expect(payload.additionalModelRequestFields?.output_config).toEqual({ effort: "high" });
expect(payload.additionalModelRequestFields?.anthropic_beta).toBeUndefined();
});
}); });