refactor(amazon-bedrock): use exact AWS SDK contracts

This commit is contained in:
Peter Steinberger
2026-07-13 20:24:48 +01:00
parent a15c6af48e
commit d2f4fdf8e6
4 changed files with 29 additions and 81 deletions

View File

@@ -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<unknown>;
};
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<SharedIniFileLoader> {
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 });
}

View File

@@ -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(

View File

@@ -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<AwsCredentialProvider | null>;
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<AwsCredentialProviderSdk | null>;
let sdkCache: AwsSdk | null = null;
let credentialProviderSdkCache: AwsCredentialProviderSdk | null | undefined;
let sdkPromise: Promise<AwsSdk> | null = null;
let credentialProviderPromise: Promise<AwsCredentialProvider | null> | null = null;
async function loadSdk(): Promise<AwsSdk> {
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<AwsSdk> {
}
}
async function loadCredentialProviderSdk(): Promise<AwsCredentialProviderSdk | null> {
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<AwsCredentialProvider | null> {
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<boolean> {
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,
})();

View File

@@ -247,25 +247,6 @@ type BedrockAppProfileTraits = {
const appProfileTraitsCache = new Map<string, BedrockAppProfileTraits>();
type BedrockGetInferenceProfileResponse = {
models?: Array<{ modelArn?: string }>;
};
type BedrockControlPlane = {
getInferenceProfile: (input: {
inferenceProfileIdentifier: string;
}) => Promise<BedrockGetInferenceProfileResponse>;
};
async function createBedrockControlPlane(region: string | undefined): Promise<BedrockControlPlane> {
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)),