From d2f4fdf8e6fa5fe480c1e7bc2fac01b35efe75b0 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 13 Jul 2026 20:24:48 +0100 Subject: [PATCH] refactor(amazon-bedrock): use exact AWS SDK contracts --- .../amazon-bedrock/aws-credential-refresh.ts | 12 +--- .../amazon-bedrock/embedding-provider.test.ts | 12 ++-- .../amazon-bedrock/embedding-provider.ts | 57 +++++-------------- .../amazon-bedrock/register.sync.runtime.ts | 29 +++------- 4 files changed, 29 insertions(+), 81 deletions(-) diff --git a/extensions/amazon-bedrock/aws-credential-refresh.ts b/extensions/amazon-bedrock/aws-credential-refresh.ts index a0848cf11289..a678c3b2d539 100644 --- a/extensions/amazon-bedrock/aws-credential-refresh.ts +++ b/extensions/amazon-bedrock/aws-credential-refresh.ts @@ -2,10 +2,6 @@ * AWS shared config cache refresh helpers for Bedrock. They nudge the AWS SDK * to re-read profile/SSO config when no static credentials are present. */ -type SharedIniFileLoader = { - loadSharedConfigFiles(init?: { ignoreCache?: boolean }): Promise; -}; - function hasStaticAwsCredentialEnv(env: NodeJS.ProcessEnv): boolean { return Boolean(env.AWS_ACCESS_KEY_ID && env.AWS_SECRET_ACCESS_KEY); } @@ -18,10 +14,6 @@ function shouldRefreshAwsSharedConfigCacheForBedrock(env: NodeJS.ProcessEnv): bo return !hasStaticAwsCredentialEnv(env); } -async function loadSharedIniFileLoader(): Promise { - return (await import("@smithy/shared-ini-file-loader")) as SharedIniFileLoader; -} - /** Refresh Smithy shared config files when Bedrock needs default-chain credentials. */ export async function refreshAwsSharedConfigCacheForBedrock( env: NodeJS.ProcessEnv = process.env, @@ -29,6 +21,6 @@ export async function refreshAwsSharedConfigCacheForBedrock( if (!shouldRefreshAwsSharedConfigCacheForBedrock(env)) { return; } - const loader = await loadSharedIniFileLoader(); - await loader.loadSharedConfigFiles({ ignoreCache: true }); + const { loadSharedConfigFiles } = await import("@smithy/shared-ini-file-loader"); + await loadSharedConfigFiles({ ignoreCache: true }); } diff --git a/extensions/amazon-bedrock/embedding-provider.test.ts b/extensions/amazon-bedrock/embedding-provider.test.ts index 5c1526c97a81..b4447c87b415 100644 --- a/extensions/amazon-bedrock/embedding-provider.test.ts +++ b/extensions/amazon-bedrock/embedding-provider.test.ts @@ -35,9 +35,9 @@ describe("hasAwsCredentials", () => { }); it("requires AWS profile credentials to resolve through the credential chain", async () => { - const loadCredentialProvider = vi.fn().mockResolvedValue({ - defaultProvider: () => async () => ({ accessKeyId: "resolved-access-key" }), - }); + const loadCredentialProvider = vi + .fn() + .mockResolvedValue(() => async () => ({ accessKeyId: "resolved-access-key" })); await expect(hasAwsCredentials({ AWS_PROFILE: "work" }, loadCredentialProvider)).resolves.toBe( true, @@ -47,10 +47,8 @@ describe("hasAwsCredentials", () => { }); it("rejects AWS profile markers when the credential chain cannot resolve", async () => { - const loadCredentialProvider = vi.fn().mockResolvedValue({ - defaultProvider: () => async () => { - throw new Error("Could not load credentials from any providers"); - }, + const loadCredentialProvider = vi.fn().mockResolvedValue(() => async () => { + throw new Error("Could not load credentials from any providers"); }); await expect( diff --git a/extensions/amazon-bedrock/embedding-provider.ts b/extensions/amazon-bedrock/embedding-provider.ts index ab577179ff3c..655b3a2ccebf 100644 --- a/extensions/amazon-bedrock/embedding-provider.ts +++ b/extensions/amazon-bedrock/embedding-provider.ts @@ -118,38 +118,18 @@ function inferFamily(modelId: string): Family { // AWS SDK lazy loader // --------------------------------------------------------------------------- -type SdkClient = import("@aws-sdk/client-bedrock-runtime").BedrockRuntimeClient; -type SdkCommand = import("@aws-sdk/client-bedrock-runtime").InvokeModelCommand; +type AwsSdk = typeof import("@aws-sdk/client-bedrock-runtime"); +type AwsCredentialProvider = typeof import("@aws-sdk/credential-provider-node").defaultProvider; +type AwsCredentialProviderLoader = () => Promise; -interface AwsSdk { - BedrockRuntimeClient: new (config: { region: string }) => SdkClient; - InvokeModelCommand: new (input: { - modelId: string; - body: string; - contentType: string; - accept: string; - }) => SdkCommand; -} - -interface AwsCredentialProviderSdk { - defaultProvider: (init?: { timeout?: number; maxRetries?: number }) => () => Promise<{ - accessKeyId?: string; - }>; -} - -type AwsCredentialProviderLoader = () => Promise; - -let sdkCache: AwsSdk | null = null; -let credentialProviderSdkCache: AwsCredentialProviderSdk | null | undefined; +let sdkPromise: Promise | null = null; +let credentialProviderPromise: Promise | null = null; async function loadSdk(): Promise { - if (sdkCache) { - return sdkCache; - } try { - sdkCache = (await import("@aws-sdk/client-bedrock-runtime")) as unknown as AwsSdk; - return sdkCache; + return await (sdkPromise ??= import("@aws-sdk/client-bedrock-runtime")); } catch { + sdkPromise = null; throw new Error( "No API key found for provider bedrock: @aws-sdk/client-bedrock-runtime is not installed. " + "Install it with: npm install @aws-sdk/client-bedrock-runtime", @@ -157,17 +137,10 @@ async function loadSdk(): Promise { } } -async function loadCredentialProviderSdk(): Promise { - if (credentialProviderSdkCache !== undefined) { - return credentialProviderSdkCache; - } - try { - credentialProviderSdkCache = - (await import("@aws-sdk/credential-provider-node")) as unknown as AwsCredentialProviderSdk; - } catch { - credentialProviderSdkCache = null; - } - return credentialProviderSdkCache; +function loadDefaultCredentialProvider(): Promise { + return (credentialProviderPromise ??= import("@aws-sdk/credential-provider-node") + .then(({ defaultProvider }) => defaultProvider) + .catch(() => null)); } // --------------------------------------------------------------------------- @@ -454,7 +427,7 @@ function resolveBedrockEmbeddingClient( export async function hasAwsCredentials( env: NodeJS.ProcessEnv = process.env, - loadCredentialProvider: AwsCredentialProviderLoader = loadCredentialProviderSdk, + loadCredentialProvider: AwsCredentialProviderLoader = loadDefaultCredentialProvider, ): Promise { if (env.AWS_ACCESS_KEY_ID?.trim() && env.AWS_SECRET_ACCESS_KEY?.trim()) { return true; @@ -462,12 +435,12 @@ export async function hasAwsCredentials( if (env.AWS_BEARER_TOKEN_BEDROCK?.trim()) { return true; } - const credentialProviderSdk = await loadCredentialProvider(); - if (!credentialProviderSdk) { + const defaultProvider = await loadCredentialProvider(); + if (!defaultProvider) { return false; } try { - const credentials = await credentialProviderSdk.defaultProvider({ + const credentials = await defaultProvider({ timeout: 1000, maxRetries: 0, })(); diff --git a/extensions/amazon-bedrock/register.sync.runtime.ts b/extensions/amazon-bedrock/register.sync.runtime.ts index bec0e2ace59a..dde1444f7023 100644 --- a/extensions/amazon-bedrock/register.sync.runtime.ts +++ b/extensions/amazon-bedrock/register.sync.runtime.ts @@ -247,25 +247,6 @@ type BedrockAppProfileTraits = { const appProfileTraitsCache = new Map(); -type BedrockGetInferenceProfileResponse = { - models?: Array<{ modelArn?: string }>; -}; - -type BedrockControlPlane = { - getInferenceProfile: (input: { - inferenceProfileIdentifier: string; - }) => Promise; -}; - -async function createBedrockControlPlane(region: string | undefined): Promise { - await refreshAwsSharedConfigCacheForBedrock(); - const { BedrockClient, GetInferenceProfileCommand } = await import("@aws-sdk/client-bedrock"); - const client = new BedrockClient(region ? { region } : {}); - return { - getInferenceProfile: async (input) => await client.send(new GetInferenceProfileCommand(input)), - }; -} - async function resolveAppProfileTraits( modelId: string, fallbackRegion: string | undefined, @@ -276,10 +257,14 @@ async function resolveAppProfileTraits( } try { const region = extractRegionFromArn(modelId) ?? fallbackRegion; - const controlPlane = await createBedrockControlPlane(region); - const resp = await controlPlane.getInferenceProfile({ inferenceProfileIdentifier: modelId }); + await refreshAwsSharedConfigCacheForBedrock(); + const { BedrockClient, GetInferenceProfileCommand } = await import("@aws-sdk/client-bedrock"); + const client = new BedrockClient(region ? { region } : {}); + const resp = await client.send( + new GetInferenceProfileCommand({ inferenceProfileIdentifier: modelId }), + ); const models = resp.models ?? []; - const modelArns = models.map((m: { modelArn?: string }) => m.modelArn ?? ""); + const modelArns = models.map((model) => model.modelArn ?? ""); const traits = { cacheEligible: models.length > 0 && modelArns.every((modelArn) => resolvedModelSupportsCaching(modelArn)),