diff --git a/extensions/amazon-bedrock/index.test.ts b/extensions/amazon-bedrock/index.test.ts index 049ebc45810..e4addeb776b 100644 --- a/extensions/amazon-bedrock/index.test.ts +++ b/extensions/amazon-bedrock/index.test.ts @@ -20,6 +20,104 @@ describe("amazon-bedrock provider plugin", () => { ).toBeUndefined(); }); + it("enables prompt caching for Application Inference Profile ARNs with Claude model name", () => { + const provider = registerSingleProviderPlugin(amazonBedrockPlugin); + const baseFn = (_model: never, _context: never, options: Record) => options; + const arn = + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/my-claude-profile"; + const result = provider.wrapStreamFn?.({ + provider: "amazon-bedrock", + modelId: arn, + config: { + models: { + providers: { + "amazon-bedrock": { + models: [{ id: arn, name: "Claude Sonnet 4.6 via Inference Profile" }], + }, + }, + }, + }, + streamFn: baseFn, + } as never); + + // Should return the original streamFn (no no-cache wrapper) + expect(result).toBe(baseFn); + }); + + it("enables prompt caching for inference profile when config uses provider alias", () => { + const provider = registerSingleProviderPlugin(amazonBedrockPlugin); + const baseFn = (_model: never, _context: never, options: Record) => options; + const arn = + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/my-claude-profile"; + const result = provider.wrapStreamFn?.({ + provider: "amazon-bedrock", + modelId: arn, + config: { + models: { + providers: { + // "bedrock" is a known alias that normalizeProviderId maps to "amazon-bedrock" + bedrock: { + models: [{ id: arn, name: "Claude Sonnet 4.6 via Inference Profile" }], + }, + }, + }, + }, + streamFn: baseFn, + } as never); + + // Should return the original streamFn even when config key is "bedrock" alias + expect(result).toBe(baseFn); + }); + + it("disables prompt caching for Application Inference Profile ARNs with non-Claude model name", () => { + const provider = registerSingleProviderPlugin(amazonBedrockPlugin); + const baseFn = (_model: never, _context: never, options: Record) => options; + const arn = + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/llama-profile"; + const wrapped = provider.wrapStreamFn?.({ + provider: "amazon-bedrock", + modelId: arn, + config: { + models: { + providers: { + "amazon-bedrock": { + models: [{ id: arn, name: "Llama 2 via Inference Profile" }], + }, + }, + }, + }, + streamFn: baseFn, + } as never); + + expect( + wrapped?.( + { api: "openai-completions", provider: "amazon-bedrock", id: arn } as never, + { messages: [] } as never, + {}, + ), + ).toMatchObject({ cacheRetention: "none" }); + }); + + it("disables prompt caching for Application Inference Profile ARNs with no config entry", () => { + const provider = registerSingleProviderPlugin(amazonBedrockPlugin); + const baseFn = (_model: never, _context: never, options: Record) => options; + const arn = + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/unknown-profile"; + const wrapped = provider.wrapStreamFn?.({ + provider: "amazon-bedrock", + modelId: arn, + streamFn: baseFn, + } as never); + + expect( + wrapped?.( + { api: "openai-completions", provider: "amazon-bedrock", id: arn } as never, + { messages: [] } as never, + {}, + ), + ).toMatchObject({ cacheRetention: "none" }); + }); + it("disables prompt caching for non-Anthropic Bedrock models", () => { const provider = registerSingleProviderPlugin(amazonBedrockPlugin); const wrapped = provider.wrapStreamFn?.({ @@ -42,4 +140,74 @@ describe("amazon-bedrock provider plugin", () => { cacheRetention: "none", }); }); + + it("injects region from bedrockDiscovery config into stream options", () => { + const provider = registerSingleProviderPlugin(amazonBedrockPlugin); + const baseFn = (_model: never, _context: never, options: Record) => options; + const wrapped = provider.wrapStreamFn?.({ + provider: "amazon-bedrock", + modelId: "eu.anthropic.claude-sonnet-4-6", + config: { + models: { + bedrockDiscovery: { region: "eu-west-1" }, + }, + }, + streamFn: baseFn, + } as never); + + const result = wrapped?.( + { + api: "bedrock-converse-stream", + provider: "amazon-bedrock", + id: "eu.anthropic.claude-sonnet-4-6", + } as never, + { messages: [] } as never, + {}, + ); + expect(result).toMatchObject({ region: "eu-west-1" }); + }); + + it("injects region extracted from provider baseUrl into stream options", () => { + const provider = registerSingleProviderPlugin(amazonBedrockPlugin); + const baseFn = (_model: never, _context: never, options: Record) => options; + const wrapped = provider.wrapStreamFn?.({ + provider: "amazon-bedrock", + modelId: "eu.anthropic.claude-sonnet-4-6", + config: { + models: { + providers: { + "amazon-bedrock": { + baseUrl: "https://bedrock-runtime.eu-central-1.amazonaws.com", + models: [], + }, + }, + }, + }, + streamFn: baseFn, + } as never); + + const result = wrapped?.( + { + api: "bedrock-converse-stream", + provider: "amazon-bedrock", + id: "eu.anthropic.claude-sonnet-4-6", + } as never, + { messages: [] } as never, + {}, + ); + expect(result).toMatchObject({ region: "eu-central-1" }); + }); + + it("does not inject region when neither bedrockDiscovery nor baseUrl is configured", () => { + const provider = registerSingleProviderPlugin(amazonBedrockPlugin); + const baseFn = (_model: never, _context: never, options: Record) => options; + const result = provider.wrapStreamFn?.({ + provider: "amazon-bedrock", + modelId: "anthropic.claude-sonnet-4-6", + streamFn: baseFn, + } as never); + + // Without region config, Claude model returns the base streamFn directly + expect(result).toBe(baseFn); + }); }); diff --git a/extensions/amazon-bedrock/index.ts b/extensions/amazon-bedrock/index.ts index 7c76a5419da..88a7fb622e5 100644 --- a/extensions/amazon-bedrock/index.ts +++ b/extensions/amazon-bedrock/index.ts @@ -1,4 +1,5 @@ import { definePluginEntry } from "openclaw/plugin-sdk/plugin-entry"; +import { normalizeProviderId } from "openclaw/plugin-sdk/provider-models"; import { createBedrockNoCacheWrapper, isAnthropicBedrockModel, @@ -7,6 +8,15 @@ import { const PROVIDER_ID = "amazon-bedrock"; const CLAUDE_46_MODEL_RE = /claude-(?:opus|sonnet)-4(?:\.|-)6(?:$|[-.])/i; +/** Extract the AWS region from a bedrock-runtime baseUrl, e.g. "https://bedrock-runtime.eu-west-1.amazonaws.com". */ +function extractRegionFromBaseUrl(baseUrl: string | undefined): string | undefined { + if (!baseUrl) { + return undefined; + } + const match = /bedrock-runtime\.([a-z0-9-]+)\.amazonaws\.com/.exec(baseUrl); + return match?.[1]; +} + export default definePluginEntry({ id: PROVIDER_ID, name: "Amazon Bedrock Provider", @@ -17,8 +27,60 @@ export default definePluginEntry({ label: "Amazon Bedrock", docsPath: "/providers/models", auth: [], - wrapStreamFn: ({ modelId, streamFn }) => - isAnthropicBedrockModel(modelId) ? streamFn : createBedrockNoCacheWrapper(streamFn), + wrapStreamFn: ({ modelId, config, streamFn }) => { + // Look up model name and region from provider config. + // Use normalized key matching so aliases like "bedrock" / "aws-bedrock" are found. + let modelName: string | undefined; + let providerBaseUrl: string | undefined; + const providers = config?.models?.providers; + if (providers) { + for (const [key, value] of Object.entries(providers)) { + if (normalizeProviderId(key) !== PROVIDER_ID) { + continue; + } + const typedValue = value as { + baseUrl?: string; + models?: Array<{ id?: string; name?: string }>; + }; + if (!providerBaseUrl && typedValue.baseUrl) { + providerBaseUrl = typedValue.baseUrl; + } + const modelDef = typedValue.models?.find((m) => m.id === modelId); + if (modelDef?.name) { + modelName = modelDef.name; + break; + } + } + } + + // Extract region from provider baseUrl or bedrockDiscovery config so the + // pi-ai BedrockRuntimeClient uses the correct endpoint. Without this, the + // gateway process (which may not inherit AWS_REGION) falls back to us-east-1 + // and rejects cross-region inference profile IDs like "eu.anthropic.claude-*". + const region = + config?.models?.bedrockDiscovery?.region ?? extractRegionFromBaseUrl(providerBaseUrl); + + const baseFn = isAnthropicBedrockModel(modelId, modelName) + ? streamFn + : createBedrockNoCacheWrapper(streamFn); + + if (!region) { + return baseFn; + } + + // Wrap to inject the region into every stream call. + const underlying = baseFn ?? streamFn; + if (!underlying) { + return baseFn; + } + return (model, context, options) => { + // pi-ai's bedrock provider reads `options.region` at runtime but the + // StreamFn type does not declare it. Merge via Object.assign to avoid + // an unsafe type assertion. + const merged = Object.assign({}, options, { region }); + return underlying(model, context, merged); + }; + }, resolveDefaultThinkingLevel: ({ modelId }) => CLAUDE_46_MODEL_RE.test(modelId.trim()) ? "adaptive" : undefined, }); diff --git a/src/agents/pi-embedded-runner/anthropic-stream-wrappers.test.ts b/src/agents/pi-embedded-runner/anthropic-stream-wrappers.test.ts new file mode 100644 index 00000000000..2cad3a6440a --- /dev/null +++ b/src/agents/pi-embedded-runner/anthropic-stream-wrappers.test.ts @@ -0,0 +1,156 @@ +import { describe, expect, it } from "vitest"; +import { isAnthropicBedrockModel } from "./anthropic-stream-wrappers.js"; + +describe("isAnthropicBedrockModel", () => { + describe("standard Bedrock model IDs", () => { + it("should return true for standard Anthropic Bedrock model IDs with dot notation", () => { + expect(isAnthropicBedrockModel("anthropic.claude-3-opus-20240229-v1:0")).toBe(true); + expect(isAnthropicBedrockModel("anthropic.claude-3-sonnet-20240229-v1:0")).toBe(true); + expect(isAnthropicBedrockModel("anthropic.claude-3-haiku-20240307-v1:0")).toBe(true); + expect(isAnthropicBedrockModel("anthropic.claude-v2:1")).toBe(true); + expect(isAnthropicBedrockModel("anthropic.claude-instant-v1")).toBe(true); + }); + + it("should return true for standard Anthropic Bedrock model IDs with slash notation", () => { + expect(isAnthropicBedrockModel("anthropic/claude-3-opus-20240229-v1:0")).toBe(true); + expect(isAnthropicBedrockModel("anthropic/claude-3-sonnet-20240229-v1:0")).toBe(true); + }); + + it("should return true for new US Anthropic Bedrock model IDs", () => { + expect(isAnthropicBedrockModel("us.anthropic.claude-opus-4-6-v1:0")).toBe(true); + expect(isAnthropicBedrockModel("us.anthropic.claude-sonnet-4-6-v1:0")).toBe(true); + expect(isAnthropicBedrockModel("us.anthropic.claude-haiku-4-5-v1:0")).toBe(true); + }); + + it("should handle case-insensitive matching", () => { + expect(isAnthropicBedrockModel("ANTHROPIC.CLAUDE-3-OPUS-20240229-V1:0")).toBe(true); + expect(isAnthropicBedrockModel("Anthropic.Claude-3-Opus-20240229-v1:0")).toBe(true); + }); + + it("should return false for non-Anthropic Bedrock model IDs", () => { + expect(isAnthropicBedrockModel("amazon.titan-text-express-v1")).toBe(false); + expect(isAnthropicBedrockModel("amazon.titan-embed-text-v1")).toBe(false); + expect(isAnthropicBedrockModel("meta.llama2-13b-chat-v1")).toBe(false); + expect(isAnthropicBedrockModel("cohere.command-text-v14")).toBe(false); + expect(isAnthropicBedrockModel("ai21.j2-ultra-v1")).toBe(false); + expect(isAnthropicBedrockModel("stability.stable-diffusion-xl-v1")).toBe(false); + }); + }); + + describe("Application Inference Profile ARNs", () => { + it("should return true for Application Inference Profile ARNs when model name contains 'claude'", () => { + const arn = + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/my-claude-profile"; + expect(isAnthropicBedrockModel(arn, "Claude 3 Opus Profile")).toBe(true); + expect(isAnthropicBedrockModel(arn, "My Claude Model")).toBe(true); + expect(isAnthropicBedrockModel(arn, "claude-profile")).toBe(true); + expect(isAnthropicBedrockModel(arn, "CLAUDE-PROFILE")).toBe(true); + }); + + it("should return false for Application Inference Profile ARNs when model name does not contain 'claude'", () => { + const arn = + "arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/llama-profile"; + expect(isAnthropicBedrockModel(arn, "Llama 2 Profile")).toBe(false); + expect(isAnthropicBedrockModel(arn, "Titan Model")).toBe(false); + expect(isAnthropicBedrockModel(arn, "General Model")).toBe(false); + }); + + it("should return false for Application Inference Profile ARNs when neither profile ID nor model name contains 'claude'", () => { + const arn = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/my-profile"; + expect(isAnthropicBedrockModel(arn)).toBe(false); + expect(isAnthropicBedrockModel(arn, undefined)).toBe(false); + expect(isAnthropicBedrockModel(arn, "Titan Model")).toBe(false); + }); + + it("should return true for Application Inference Profile ARNs when profile ID contains 'claude' (no model name needed)", () => { + expect( + isAnthropicBedrockModel( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/my-claude-profile", + ), + ).toBe(true); + expect( + isAnthropicBedrockModel( + "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/claude-sonnet", + ), + ).toBe(true); + }); + + it("should handle Application Inference Profile ARNs with various formats", () => { + expect( + isAnthropicBedrockModel( + "arn:aws:bedrock:eu-west-1:987654321098:application-inference-profile/prod-claude", + "Production Claude Model", + ), + ).toBe(true); + expect( + isAnthropicBedrockModel( + "arn:aws:bedrock:ap-southeast-1:111222333444:application-inference-profile/test", + "Test Claude Instance", + ), + ).toBe(true); + }); + }); + + describe("short Application Inference Profile IDs", () => { + it("should return true when short ID is provided with model name containing 'claude'", () => { + expect(isAnthropicBedrockModel("my-profile-id", "Claude Profile")).toBe(true); + expect(isAnthropicBedrockModel("prod-inference", "Production Claude")).toBe(true); + expect(isAnthropicBedrockModel("test-profile", "claude-test")).toBe(true); + }); + + it("should return false when short ID is provided with model name not containing 'claude'", () => { + expect(isAnthropicBedrockModel("my-profile-id", "Titan Profile")).toBe(false); + expect(isAnthropicBedrockModel("prod-inference", "Production Llama")).toBe(false); + expect(isAnthropicBedrockModel("test-profile", "test-model")).toBe(false); + }); + + it("should return false for short IDs without model name and without 'claude' in ID", () => { + expect(isAnthropicBedrockModel("my-profile-id")).toBe(false); + expect(isAnthropicBedrockModel("prod-inference")).toBe(false); + expect(isAnthropicBedrockModel("test-profile")).toBe(false); + }); + + it("should return true for short IDs containing 'claude' without model name", () => { + expect(isAnthropicBedrockModel("my-claude-profile")).toBe(true); + expect(isAnthropicBedrockModel("claude-sonnet-profile")).toBe(true); + }); + }); + + describe("edge cases", () => { + it("should handle empty strings", () => { + expect(isAnthropicBedrockModel("")).toBe(false); + expect(isAnthropicBedrockModel("", "")).toBe(false); + expect(isAnthropicBedrockModel("", "Claude Model")).toBe(false); + }); + + it("should handle model names with 'claude' in different positions", () => { + expect(isAnthropicBedrockModel("some-id", "claude")).toBe(true); + expect(isAnthropicBedrockModel("some-id", "claude-at-start")).toBe(true); + expect(isAnthropicBedrockModel("some-id", "middle-claude-here")).toBe(true); + expect(isAnthropicBedrockModel("some-id", "ends-with-claude")).toBe(true); + expect(isAnthropicBedrockModel("some-id", "CLAUDE")).toBe(true); + expect(isAnthropicBedrockModel("some-id", "ClAuDe")).toBe(true); + }); + + it("should match claude as a substring in any position within model names", () => { + // These should still match because they contain 'claude' as a substring + expect(isAnthropicBedrockModel("some-id", "unclaude")).toBe(true); + expect(isAnthropicBedrockModel("some-id", "claudette")).toBe(true); + }); + + it("should handle special characters in model IDs and names", () => { + expect(isAnthropicBedrockModel("anthropic.claude-3_opus@v1:0")).toBe(true); + // IDs with special characters don't match short profile ID pattern + expect(isAnthropicBedrockModel("model-with-special-chars!@#", "Claude Model!")).toBe(false); + }); + }); + + describe("backward compatibility", () => { + it("should maintain backward compatibility when modelName parameter is not provided", () => { + // These should work exactly as before when no modelName is provided + expect(isAnthropicBedrockModel("anthropic.claude-3-opus-20240229-v1:0")).toBe(true); + expect(isAnthropicBedrockModel("amazon.titan-text-express-v1")).toBe(false); + expect(isAnthropicBedrockModel("us.anthropic.claude-opus-4-6-v1:0")).toBe(true); + }); + }); +}); diff --git a/src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts b/src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts index 511b70d280d..6800f142e87 100644 --- a/src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts +++ b/src/agents/pi-embedded-runner/anthropic-stream-wrappers.ts @@ -390,7 +390,52 @@ export function createBedrockNoCacheWrapper(baseStreamFn: StreamFn | undefined): }); } -export function isAnthropicBedrockModel(modelId: string): boolean { +export function isAnthropicBedrockModel(modelId: string, modelName?: string): boolean { const normalized = modelId.toLowerCase(); - return normalized.includes("anthropic.claude") || normalized.includes("anthropic/claude"); + + // Direct Anthropic Claude model IDs (e.g., anthropic.claude-sonnet-4-6, global.anthropic.claude-opus-4-6-v1) + if (normalized.includes("anthropic.claude") || normalized.includes("anthropic/claude")) { + return true; + } + + // Application Inference Profile ARN — detect Claude via profile ID segment or model name. + // ARN format: arn::bedrock:::application-inference-profile/ + // Supports all AWS partitions with Bedrock: aws, aws-cn, aws-us-gov. + // Note: model name (`models[].name`) is a user-chosen display label, so this is best-effort. + // A profile ID or name containing "claude" is treated as an Anthropic Claude model; if neither + // contains "claude", the no-cache wrapper is applied (safe default). + if ( + /^arn:aws(-cn|-us-gov)?:bedrock:/.test(normalized) && + normalized.includes(":application-inference-profile/") + ) { + const profileId = normalized.split(":application-inference-profile/")[1] ?? ""; + if (profileId.includes("claude")) { + return true; + } + return modelName ? modelName.toLowerCase().includes("claude") : false; + } + + // Short/opaque inference profile IDs — fall back to ID or model name for IDs + // that look like short profile IDs (alphanumeric, no dots/colons/slashes). + // Excludes standard model IDs, other ARN resource types, and any dotted identifiers. + // Note: the regex is intentionally broad; it is safe because standard Bedrock model IDs + // always contain dots or colons (e.g. "amazon.nova-micro-v1:0") which exclude them here. + if (looksLikeShortProfileId(normalized)) { + if (normalized.includes("claude")) { + return true; + } + return modelName ? modelName.toLowerCase().includes("claude") : false; + } + + return false; +} + +/** + * Returns true when the ID looks like a short Application Inference Profile ID + * (opaque alphanumeric-and-hyphen string, cannot start/end with a hyphen). + * Requires at least 2 characters; single-char IDs are not realistic model IDs. + * Examples: "gdkqufd9flgg", "s3rr0t98ews8", "my-claude-profile" + */ +function looksLikeShortProfileId(normalizedId: string): boolean { + return /^[a-z0-9][a-z0-9-]*[a-z0-9]$/.test(normalizedId); }