fix(amazon-bedrock): skip auto memory embeddings without credentials (#71245)

Co-authored-by: bitloi <raphaelaloi.eth@gmail.com>

extensions/amazon-bedrock/embedding-provider.test.ts

@@ -0,0 +1,65 @@
+import { describe, expect, it, vi } from "vitest";
+import { hasAwsCredentials } from "./embedding-provider.js";
+
+describe("hasAwsCredentials", () => {
+  it("accepts static AWS key credentials without loading the credential chain", async () => {
+    const loadCredentialProvider = vi.fn();
+
+    await expect(
+      hasAwsCredentials(
+        {
+          AWS_ACCESS_KEY_ID: "access-key",
+          AWS_SECRET_ACCESS_KEY: "secret-key",
+        },
+        loadCredentialProvider,
+      ),
+    ).resolves.toBe(true);
+
+    expect(loadCredentialProvider).not.toHaveBeenCalled();
+  });
+
+  it("accepts the Bedrock bearer token without loading the credential chain", async () => {
+    const loadCredentialProvider = vi.fn();
+
+    await expect(
+      hasAwsCredentials(
+        {
+          AWS_BEARER_TOKEN_BEDROCK: "bearer-token",
+        },
+        loadCredentialProvider,
+      ),
+    ).resolves.toBe(true);
+
+    expect(loadCredentialProvider).not.toHaveBeenCalled();
+  });
+
+  it("requires AWS profile credentials to resolve through the credential chain", async () => {
+    const loadCredentialProvider = vi.fn().mockResolvedValue({
+      defaultProvider: () => async () => ({ accessKeyId: "resolved-access-key" }),
+    });
+
+    await expect(hasAwsCredentials({ AWS_PROFILE: "work" }, loadCredentialProvider)).resolves.toBe(
+      true,
+    );
+
+    expect(loadCredentialProvider).toHaveBeenCalledOnce();
+  });
+
+  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");
+      },
+    });
+
+    await expect(
+      hasAwsCredentials({ AWS_PROFILE: "missing" }, loadCredentialProvider),
+    ).resolves.toBe(false);
+  });
+
+  it("returns false when the AWS credential provider package is unavailable", async () => {
+    const loadCredentialProvider = vi.fn().mockResolvedValue(null);
+
+    await expect(hasAwsCredentials({}, loadCredentialProvider)).resolves.toBe(false);
+  });
+});

extensions/amazon-bedrock/embedding-provider.ts

@@ -122,6 +122,8 @@ interface AwsCredentialProviderSdk {
   }>;
 }
 
+type AwsCredentialProviderLoader = () => Promise<AwsCredentialProviderSdk | null>;
+
 let sdkCache: AwsSdk | null = null;
 let credentialProviderSdkCache: AwsCredentialProviderSdk | null | undefined;
 
@@ -368,24 +370,17 @@ export function resolveBedrockEmbeddingClient(
 // Credential detection
 // ---------------------------------------------------------------------------
 
-const CREDENTIAL_ENV_VARS = [
-  "AWS_PROFILE",
-  "AWS_BEARER_TOKEN_BEDROCK",
-  "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI",
-  "AWS_CONTAINER_CREDENTIALS_FULL_URI",
-  "AWS_EC2_METADATA_SERVICE_ENDPOINT",
-  "AWS_WEB_IDENTITY_TOKEN_FILE",
-  "AWS_ROLE_ARN",
-] as const;
-
-export async function hasAwsCredentials(env: NodeJS.ProcessEnv = process.env): Promise<boolean> {
+export async function hasAwsCredentials(
+  env: NodeJS.ProcessEnv = process.env,
+  loadCredentialProvider: AwsCredentialProviderLoader = loadCredentialProviderSdk,
+): Promise<boolean> {
   if (env.AWS_ACCESS_KEY_ID?.trim() && env.AWS_SECRET_ACCESS_KEY?.trim()) {
     return true;
   }
-  if (CREDENTIAL_ENV_VARS.some((k) => env[k]?.trim())) {
+  if (env.AWS_BEARER_TOKEN_BEDROCK?.trim()) {
     return true;
   }
-  const credentialProviderSdk = await loadCredentialProviderSdk();
+  const credentialProviderSdk = await loadCredentialProvider();
   if (!credentialProviderSdk) {
     return false;
   }

extensions/amazon-bedrock/memory-embedding-adapter.test.ts

@@ -0,0 +1,100 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const hasAwsCredentialsMock = vi.hoisted(() => vi.fn());
+const createBedrockEmbeddingProviderMock = vi.hoisted(() => vi.fn());
+
+vi.mock("./embedding-provider.js", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("./embedding-provider.js")>();
+  return {
+    ...actual,
+    hasAwsCredentials: hasAwsCredentialsMock,
+    createBedrockEmbeddingProvider: createBedrockEmbeddingProviderMock,
+  };
+});
+
+import { bedrockMemoryEmbeddingProviderAdapter } from "./memory-embedding-adapter.js";
+
+function defaultCreateOptions() {
+  return {
+    config: {} as Record<string, unknown>,
+    agentDir: "/tmp/test-agent",
+    model: "",
+  };
+}
+
+function stubCreate(client: { region: string; model: string; dimensions?: number }) {
+  createBedrockEmbeddingProviderMock.mockResolvedValue({
+    provider: {
+      id: "bedrock",
+      model: client.model,
+      embedQuery: async () => [],
+      embedBatch: async () => [],
+    },
+    client,
+  });
+}
+
+describe("bedrockMemoryEmbeddingProviderAdapter", () => {
+  beforeEach(() => {
+    hasAwsCredentialsMock.mockReset();
+    createBedrockEmbeddingProviderMock.mockReset();
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("registers the expected adapter metadata", () => {
+    expect(bedrockMemoryEmbeddingProviderAdapter.id).toBe("bedrock");
+    expect(bedrockMemoryEmbeddingProviderAdapter.transport).toBe("remote");
+    expect(bedrockMemoryEmbeddingProviderAdapter.authProviderId).toBe("amazon-bedrock");
+    expect(bedrockMemoryEmbeddingProviderAdapter.autoSelectPriority).toBe(60);
+    expect(bedrockMemoryEmbeddingProviderAdapter.allowExplicitWhenConfiguredAuto).toBe(true);
+  });
+
+  it("throws a missing-api-key sentinel error when AWS credentials are unavailable", async () => {
+    hasAwsCredentialsMock.mockResolvedValue(false);
+
+    await expect(
+      bedrockMemoryEmbeddingProviderAdapter.create(defaultCreateOptions()),
+    ).rejects.toThrow(/No API key found for provider "bedrock"/);
+    await expect(
+      bedrockMemoryEmbeddingProviderAdapter.create(defaultCreateOptions()),
+    ).rejects.toThrow(/AWS credentials are not available/);
+
+    expect(createBedrockEmbeddingProviderMock).not.toHaveBeenCalled();
+  });
+
+  it("creates the provider when AWS credentials are available", async () => {
+    hasAwsCredentialsMock.mockResolvedValue(true);
+    stubCreate({ region: "us-east-1", model: "amazon.titan-embed-text-v2:0", dimensions: 1024 });
+
+    const result = await bedrockMemoryEmbeddingProviderAdapter.create(defaultCreateOptions());
+
+    expect(result.provider?.id).toBe("bedrock");
+    expect(result.runtime).toEqual({
+      id: "bedrock",
+      cacheKeyData: {
+        provider: "bedrock",
+        region: "us-east-1",
+        model: "amazon.titan-embed-text-v2:0",
+        dimensions: 1024,
+      },
+    });
+    expect(createBedrockEmbeddingProviderMock).toHaveBeenCalledOnce();
+  });
+
+  it("lets the auto-select loop skip bedrock when credentials are unavailable", async () => {
+    hasAwsCredentialsMock.mockResolvedValue(false);
+
+    let thrown: unknown;
+    try {
+      await bedrockMemoryEmbeddingProviderAdapter.create(defaultCreateOptions());
+    } catch (err) {
+      thrown = err;
+    }
+
+    expect(thrown).toBeInstanceOf(Error);
+    expect(bedrockMemoryEmbeddingProviderAdapter.shouldContinueAutoSelection?.(thrown)).toBe(true);
+  });
+});

extensions/amazon-bedrock/memory-embedding-adapter.ts

@@ -5,6 +5,7 @@ import {
 import {
   createBedrockEmbeddingProvider,
   DEFAULT_BEDROCK_EMBEDDING_MODEL,
+  hasAwsCredentials,
 } from "./embedding-provider.js";
 
 export const bedrockMemoryEmbeddingProviderAdapter: MemoryEmbeddingProviderAdapter = {
@@ -16,6 +17,15 @@ export const bedrockMemoryEmbeddingProviderAdapter: MemoryEmbeddingProviderAdapt
   allowExplicitWhenConfiguredAuto: true,
   shouldContinueAutoSelection: isMissingEmbeddingApiKeyError,
   create: async (options) => {
+    if (!(await hasAwsCredentials())) {
+      throw new Error(
+        'No API key found for provider "bedrock". ' +
+          "AWS credentials are not available. " +
+          "Set AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY, AWS_PROFILE, or AWS_BEARER_TOKEN_BEDROCK, " +
+          "configure an EC2/ECS/EKS role, " +
+          "or set agents.defaults.memorySearch.provider to another provider.",
+      );
+    }
     const { provider, client } = await createBedrockEmbeddingProvider({
       ...options,
       provider: "bedrock",