From be520b858c40735baaf79b4d2312e97f92c1015a Mon Sep 17 00:00:00 2001 From: Vidit Shah Date: Sun, 25 Oct 2026 23:43:52 +0530 Subject: [PATCH] feat(cli): support global AWS profile selection --- README.md | 17 +- bun.lock | 1 + command.md | 1 + package.json | 1 + src/awsProfile.test.ts | 192 ++++++++++++++++++ src/awsProfile.ts | 38 ++++ src/core/dev/container.test.ts | 64 ++++++ src/core/dev/container.ts | 34 +++- src/core/index.tsx | 3 + src/core/project/backends/cdk/environment.ts | 8 +- src/core/project/backends/cdk/toolkit.ts | 3 + src/core/project/bedrockAgentImport/loader.ts | 3 +- src/core/region.test.ts | 13 +- src/core/region.ts | 27 ++- src/handlers/index.tsx | 8 +- src/handlers/keys.tsx | 9 + src/io/exec.ts | 6 +- src/middleware/index.tsx | 1 + src/middleware/withProfile.test.ts | 127 ++++++++++++ src/middleware/withProfile.ts | 28 +++ src/middleware/withRegion.tsx | 4 +- src/telemetry/client.test.tsx | 48 ++++- src/telemetry/client.tsx | 7 +- 23 files changed, 617 insertions(+), 26 deletions(-) create mode 100644 src/awsProfile.test.ts create mode 100644 src/awsProfile.ts create mode 100644 src/middleware/withProfile.test.ts create mode 100644 src/middleware/withProfile.ts diff --git a/README.md b/README.md index 07988524a2..945cb0bebb 100644 --- a/README.md +++ b/README.md @@ -115,11 +115,18 @@ non-interactively. `--help` on any command lists its subcommands and flags, and Global flags, available on every command: -| Flag | Purpose | -| ---------- | ------------------------------------------------------------------------------------------------------------------- | -| `--region` | AWS region. Resolved from this flag, then `AWS_REGION`, `AWS_DEFAULT_REGION`, the active profile, then `us-east-1`. | -| `--json` | Emit JSON instead of opening the TUI. | -| `--debug` | Debug logging. | +| Flag | Purpose | +| ----------- | ------------------------------------------------------------------------------------------------------------------- | +| `--region` | AWS region. Resolved from this flag, then `AWS_REGION`, `AWS_DEFAULT_REGION`, the active profile, then `us-east-1`. | +| `--profile` | Named AWS shared configuration profile. Overrides `AWS_PROFILE` and ambient static credentials for this command. | +| `--json` | Emit JSON instead of opening the TUI. | +| `--debug` | Debug logging. | + +For example, `agentcore --profile development harness list --json` uses the `development` +profile's credentials and region. `--region` and the region environment variables still take +precedence over the profile's region. Profile selection also applies to deployment and child +tools; it does not modify your shell environment. Without `--profile`, the usual AWS credential +chain is unchanged. Configure and sign in to SSO profiles with the AWS CLI before using them. ## Extending the CDK app diff --git a/bun.lock b/bun.lock index f34b3d35c7..a657fa4bdd 100644 --- a/bun.lock +++ b/bun.lock @@ -17,6 +17,7 @@ "@aws-sdk/client-s3": "^3.1121.0", "@aws-sdk/client-sts": "^3.1092.0", "@aws-sdk/client-xray": "^3.1092.0", + "@aws-sdk/credential-providers": "^3.1121.0", "@aws-sdk/util-endpoints": "^3.996.44", "@aws/agent-inspector": "0.6.2", "@opentelemetry/api": "^1.9.1", diff --git a/command.md b/command.md index 12412353ec..01e85f9018 100644 --- a/command.md +++ b/command.md @@ -226,6 +226,7 @@ the platform for production AI agents **Options** - `--region `: AWS region +- `--profile `: AWS shared configuration profile - `--debug`: debug logging (default: false) - `--json`: JSON output (default: false) - `-V, --version`: display the CLI version diff --git a/package.json b/package.json index b26d3a8aed..37b59d1649 100644 --- a/package.json +++ b/package.json @@ -79,6 +79,7 @@ "@aws-sdk/client-s3": "^3.1121.0", "@aws-sdk/client-sts": "^3.1092.0", "@aws-sdk/client-xray": "^3.1092.0", + "@aws-sdk/credential-providers": "^3.1121.0", "@aws-sdk/util-endpoints": "^3.996.44", "@aws/agent-inspector": "0.6.2", "@opentelemetry/api": "^1.9.1", diff --git a/src/awsProfile.test.ts b/src/awsProfile.test.ts new file mode 100644 index 0000000000..ff0b8c0350 --- /dev/null +++ b/src/awsProfile.test.ts @@ -0,0 +1,192 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { awsProfileEnvironment, selectedAwsProfile, withAwsProfile } from "./awsProfile"; +import { CoreClient } from "./core"; +import * as factories from "./core/factories"; +import { createSilentLogger } from "./testing"; +import { runProcess } from "./io"; +import { createCdkIoHost, resolveCdkCredentials } from "./core/project/backends/cdk/toolkit"; +import { resolveAwsAccount } from "./core/project/backends/cdk/environment"; + +const keys = [ + "AWS_CONFIG_FILE", + "AWS_SHARED_CREDENTIALS_FILE", + "AWS_PROFILE", + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", +] as const; +let saved: NodeJS.ProcessEnv; +let directory: string; +beforeEach(async () => { + saved = Object.fromEntries(keys.map((key) => [key, process.env[key]])); + directory = await mkdtemp(join(tmpdir(), "agentcore-profile-")); + process.env.AWS_CONFIG_FILE = join(directory, "config"); + process.env.AWS_SHARED_CREDENTIALS_FILE = join(directory, "credentials"); + process.env.AWS_PROFILE = "ambient"; + process.env.AWS_ACCESS_KEY_ID = "environment-key"; + process.env.AWS_SECRET_ACCESS_KEY = "environment-secret"; + await writeFile( + process.env.AWS_CONFIG_FILE, + "[profile first]\nregion = eu-west-1\n[profile second]\nregion = ap-south-1\n", + ); + await writeFile( + process.env.AWS_SHARED_CREDENTIALS_FILE, + "[first]\naws_access_key_id = first-key\naws_secret_access_key = first-secret\n" + + "[second]\naws_access_key_id = second-key\naws_secret_access_key = second-secret\n", + ); +}); +afterEach(async () => { + for (const key of keys) { + if (saved[key] === undefined) delete process.env[key]; + else process.env[key] = saved[key]; + } + await rm(directory, { recursive: true, force: true }); +}); + +describe("explicit AWS profiles", () => { + test("selects profile credentials over ambient profile and static environment keys", async () => { + await withAwsProfile("first", async () => { + expect((await selectedAwsProfile()!.credentials()).accessKeyId).toBe("first-key"); + }); + expect(process.env.AWS_PROFILE).toBe("ambient"); + expect(process.env.AWS_ACCESS_KEY_ID).toBe("environment-key"); + expect(selectedAwsProfile()).toBeUndefined(); + }); + + test("an unknown profile cannot fall back to environment credentials", async () => { + await expect( + withAwsProfile("missing", () => selectedAwsProfile()!.credentials()), + ).rejects.toThrow(); + expect(selectedAwsProfile()).toBeUndefined(); + }); + + test("loads temporary credentials from a credential_process profile", async () => { + const script = join(directory, "credentials.cjs"); + await writeFile( + script, + `process.stdout.write(JSON.stringify({ Version: 1, AccessKeyId: "process-key", SecretAccessKey: "process-secret", SessionToken: "process-token", Expiration: "2099-01-01T00:00:00Z" }));`, + ); + await writeFile( + process.env.AWS_CONFIG_FILE!, + `[profile process]\ncredential_process = "${process.execPath}" "${script}"\n`, + ); + await withAwsProfile("process", async () => { + expect(await selectedAwsProfile()!.credentials()).toMatchObject({ + accessKeyId: "process-key", + secretAccessKey: "process-secret", + sessionToken: "process-token", + }); + }); + }); + + test("STS account lookup signs with the selected profile and honors explicit credentials", async () => { + const authorization: string[] = []; + const server = Bun.serve({ + port: 0, + fetch(request) { + authorization.push(request.headers.get("authorization") ?? ""); + return new Response( + '123456789012', + { headers: { "content-type": "text/xml" } }, + ); + }, + }); + const savedEndpoint = process.env.AWS_ENDPOINT_URL_STS; + process.env.AWS_ENDPOINT_URL_STS = `http://localhost:${server.port}`; + try { + await withAwsProfile("first", async () => { + expect(await resolveAwsAccount("us-east-1")).toBe("123456789012"); + expect( + await resolveAwsAccount("us-east-1", async () => ({ + accessKeyId: "target-key", + secretAccessKey: "target-secret", + })), + ).toBe("123456789012"); + }); + expect(authorization[0]).toContain("Credential=first-key/"); + expect(authorization[1]).toContain("Credential=target-key/"); + } finally { + if (savedEndpoint === undefined) delete process.env.AWS_ENDPOINT_URL_STS; + else process.env.AWS_ENDPOINT_URL_STS = savedEndpoint; + server.stop(true); + } + }); + + test("nested and concurrent commands retain their own selection", async () => { + await withAwsProfile("first", async () => { + await expect( + withAwsProfile("second", async () => { + expect(selectedAwsProfile()!.name).toBe("second"); + throw new Error("command failed"); + }), + ).rejects.toThrow("command failed"); + expect(selectedAwsProfile()!.name).toBe("first"); + }); + const credentials = await Promise.all( + ["first", "second"].map((name) => + withAwsProfile(name, async () => { + await Promise.resolve(); + expect(selectedAwsProfile()!.name).toBe(name); + return (await selectedAwsProfile()!.credentials()).accessKeyId; + }), + ), + ); + expect(credentials).toEqual(["first-key", "second-key"]); + expect(selectedAwsProfile()).toBeUndefined(); + }); + + test("SDK caches distinguish profiles and preserve explicit target credentials", async () => { + const core = new CoreClient({ ...factories, logger: createSilentLogger() }); + const first = await withAwsProfile("first", async () => { + const client = core.control({ region: "us-east-1" }); + expect(core.control({ region: "us-east-1" })).toBe(client); + expect((await client.config.credentials()).accessKeyId).toBe("first-key"); + return client; + }); + await withAwsProfile("second", async () => { + const client = core.control({ region: "us-east-1" }); + expect(client).not.toBe(first); + expect((await client.config.credentials()).accessKeyId).toBe("second-key"); + const target = core.control({ + region: "us-east-1", + credentials: { accessKeyId: "target-key", secretAccessKey: "target-secret" }, + }); + expect((await target.config.credentials()).accessKeyId).toBe("target-key"); + const deploymentProvider = await resolveCdkCredentials( + createCdkIoHost(createSilentLogger()), + "us-east-1", + ); + expect(deploymentProvider).toBe(selectedAwsProfile()!.credentials); + expect((await deploymentProvider()).accessKeyId).toBe("second-key"); + }); + }); + + test("child tools receive the profile without inherited static credentials", async () => { + let output = ""; + await withAwsProfile("first", () => + runProcess( + [ + process.execPath, + "-e", + "process.stdout.write(JSON.stringify({profile:process.env.AWS_PROFILE,defaultProfile:process.env.AWS_DEFAULT_PROFILE,key:process.env.AWS_ACCESS_KEY_ID,token:process.env.AWS_SESSION_TOKEN,marker:process.env.MARKER}))", + ], + { + cwd: directory, + env: { ...process.env, AWS_SESSION_TOKEN: "ambient-token", MARKER: "retained" }, + onOutput: (chunk) => { + output += chunk; + }, + }, + ), + ); + expect(JSON.parse(output)).toEqual({ + profile: "first", + defaultProfile: "first", + marker: "retained", + }); + const ambient = { AWS_PROFILE: "ambient", AWS_ACCESS_KEY_ID: "ambient-key" }; + expect(awsProfileEnvironment(ambient)).toBe(ambient); + }); +}); diff --git a/src/awsProfile.ts b/src/awsProfile.ts new file mode 100644 index 0000000000..d0281957b8 --- /dev/null +++ b/src/awsProfile.ts @@ -0,0 +1,38 @@ +import { AsyncLocalStorage } from "node:async_hooks"; +import { fromIni } from "@aws-sdk/credential-providers"; +import type { AwsCredentialIdentityProvider } from "@smithy/types"; + +type SelectedProfile = { name: string; credentials: AwsCredentialIdentityProvider }; +const profiles = new AsyncLocalStorage(); + +/** Keep profile selection local to a command, including its asynchronous work. */ +export function withAwsProfile(name: string, action: () => T, region?: string): T { + return profiles.run( + { name, credentials: fromIni({ profile: name, clientConfig: { region } }) }, + action, + ); +} + +export function selectedAwsProfile(): SelectedProfile | undefined { + return profiles.getStore(); +} + +/** Child tools use the selected profile instead of inherited static credentials. */ +export function awsProfileEnvironment(env: NodeJS.ProcessEnv = process.env): NodeJS.ProcessEnv { + const profile = selectedAwsProfile(); + if (!profile) return env; + const selected: NodeJS.ProcessEnv = { + ...env, + AWS_PROFILE: profile.name, + AWS_DEFAULT_PROFILE: profile.name, + }; + for (const key of [ + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_SECURITY_TOKEN", + ]) { + delete selected[key]; + } + return selected; +} diff --git a/src/core/dev/container.test.ts b/src/core/dev/container.test.ts index c26250c289..18957920b4 100644 --- a/src/core/dev/container.test.ts +++ b/src/core/dev/container.test.ts @@ -1,4 +1,5 @@ import { afterEach, describe, expect, test } from "bun:test"; +import { withAwsProfile } from "../../awsProfile"; import { createHash } from "node:crypto"; import { mkdir, mkdtemp, readFile, rm, stat, symlink, writeFile } from "node:fs/promises"; import { tmpdir } from "node:os"; @@ -268,6 +269,69 @@ describe("ContainerDevRunner", () => { expect(missing.calls).toHaveLength(0); }); + test("an explicit profile mounts custom shared files outside the AWS directory", async () => { + const projectRuntime = runtime(); + const root = await projectRoot(projectRuntime); + const config = join(root, "custom-config"); + const credentials = join(root, "custom-credentials"); + await writeFile(config, "[profile selected]\nregion = eu-west-1\n"); + await writeFile( + credentials, + "[selected]\naws_access_key_id = selected-key\naws_secret_access_key = selected-secret\n", + ); + const { calls, runner } = harness({ + processEnv: { AWS_CONFIG_FILE: config, AWS_SHARED_CREDENTIALS_FILE: credentials }, + }); + await withAwsProfile("selected", () => collect(runner.run(input(root, projectRuntime)))); + const run = calls.find((call) => call.command[1] === "run")!; + expect(run.command).toContain(`${config}:/aws-profile/config:ro`); + expect(run.command).toContain(`${credentials}:/aws-profile/credentials:ro`); + expect(parseEnv(run.envFile!.contents)).toMatchObject({ + AWS_PROFILE: "selected", + AWS_CONFIG_FILE: "/aws-profile/config", + AWS_SHARED_CREDENTIALS_FILE: "/aws-profile/credentials", + }); + await expect( + withAwsProfile("selected", () => + collect( + harness({ processEnv: { AWS_CONFIG_FILE: join(root, "missing") } }).runner.run( + input(root, projectRuntime), + ), + ), + ), + ).rejects.toThrow("AWS_CONFIG_FILE does not reference"); + }); + + test("an explicit profile overrides ambient keys and profile in the container", async () => { + const projectRuntime = runtime(); + const root = await projectRoot(projectRuntime); + const awsDirectory = join(root, ".aws"); + await mkdir(awsDirectory); + const { calls, runner } = harness({ + awsDirectory, + processEnv: { + AWS_PROFILE: "ambient", + AWS_ACCESS_KEY_ID: "ambient-key", + AWS_SECRET_ACCESS_KEY: "ambient-secret", + }, + }); + await withAwsProfile("selected", () => collect(runner.run(input(root, projectRuntime)))); + expect(parseEnv(commandCall(calls, "run").envFile!.contents)).toMatchObject({ + AWS_PROFILE: "selected", + }); + expect(parseEnv(commandCall(calls, "run").envFile!.contents)).not.toHaveProperty( + "AWS_ACCESS_KEY_ID", + ); + expect(parseEnv(commandCall(calls, "run").envFile!.contents)).not.toHaveProperty( + "AWS_SECRET_ACCESS_KEY", + ); + const missing = harness({ awsDirectory: join(root, "missing") }); + await expect( + withAwsProfile("selected", () => collect(missing.runner.run(input(root, projectRuntime)))), + ).rejects.toBeInstanceOf(InvalidEnvironmentError); + expect(missing.calls).toHaveLength(0); + }); + test("preserves an existing build context .dockerignore", async () => { const projectRuntime = runtime({ buildContextPath: "." }); const root = await projectRoot(projectRuntime); diff --git a/src/core/dev/container.ts b/src/core/dev/container.ts index 8f3d329771..91abe333a2 100644 --- a/src/core/dev/container.ts +++ b/src/core/dev/container.ts @@ -14,6 +14,7 @@ import { } from "../../io"; import { isDirectory, isFile, resolvePathWithinProject } from "./path"; import { DEV_PORTS } from "./port"; +import { awsProfileEnvironment, selectedAwsProfile } from "../../awsProfile"; const CONTAINER_TOOLS = ["docker", "podman", "finch"] as const; const AWS_ENV_KEYS = [ @@ -91,12 +92,28 @@ export class ContainerDevRunner implements DevRunner { throw new ResourceNotFoundError(`no container Dockerfile exists at ${dockerfilePath}`); } + const effectiveEnv = awsProfileEnvironment({ ...this.processEnv, ...input.env }); const hasAwsCredentials = Boolean( - (input.env?.AWS_ACCESS_KEY_ID ?? this.processEnv.AWS_ACCESS_KEY_ID) && - (input.env?.AWS_SECRET_ACCESS_KEY ?? this.processEnv.AWS_SECRET_ACCESS_KEY), + effectiveEnv.AWS_ACCESS_KEY_ID && effectiveEnv.AWS_SECRET_ACCESS_KEY, ); const hasAwsConfig = existsSync(this.awsDirectory); - if (!hasAwsCredentials && !hasAwsConfig) { + const profileFiles: { key: string; path: string; containerPath: string }[] = []; + if (selectedAwsProfile()) { + for (const [key, name] of [ + ["AWS_CONFIG_FILE", "config"], + ["AWS_SHARED_CREDENTIALS_FILE", "credentials"], + ] as const) { + const file = effectiveEnv[key]; + if (!file) continue; + if (!isFile(file)) { + throw new InvalidEnvironmentError( + `${key} does not reference a readable AWS configuration file: ${file}`, + ); + } + profileFiles.push({ key, path: resolve(file), containerPath: `/aws-profile/${name}` }); + } + } + if (!hasAwsCredentials && !hasAwsConfig && profileFiles.length === 0) { throw new InvalidEnvironmentError( "Unable to resolve AWS credentials for the container. Configure AWS credentials " + "or set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY, then retry.", @@ -157,6 +174,13 @@ export class ContainerDevRunner implements DevRunner { PORT: String(containerPort), LOCAL_DEV: "1", }); + const selectedEnv = awsProfileEnvironment(forwardedEnv); + if (selectedEnv !== forwardedEnv) { + for (const key of Object.keys(forwardedEnv)) delete forwardedEnv[key]; + for (const [key, value] of Object.entries(selectedEnv)) { + if (value !== undefined) forwardedEnv[key] = value; + } + } if (input.runtime.protocol === "MCP") { forwardedEnv.FASTMCP_PORT = String(containerPort); } @@ -165,6 +189,10 @@ export class ContainerDevRunner implements DevRunner { forwardedEnv.AWS_CONFIG_FILE = "/aws-config/config"; forwardedEnv.AWS_SHARED_CREDENTIALS_FILE = "/aws-config/credentials"; } + for (const file of profileFiles) { + awsMount.push("-v", `${file.path}:${file.containerPath}:ro`); + forwardedEnv[file.key] = file.containerPath; + } const envFile = join(tmpdir(), `agentcore-dev-${randomUUID()}.env`); try { await writeFile(envFile, serializeEnvironment(forwardedEnv), { diff --git a/src/core/index.tsx b/src/core/index.tsx index f2999d0df3..81fe742024 100644 --- a/src/core/index.tsx +++ b/src/core/index.tsx @@ -33,6 +33,7 @@ import type { ProjectManager } from "../handlers/project/types"; import { FsProjectManager } from "./project"; import { BedrockAgentImporter, type CoreBedrockAgentImporter } from "./project/bedrockAgentImport"; import { FsProjectMigrator } from "./project/migrate"; +import { selectedAwsProfile } from "../awsProfile"; import type { ProjectMigrator } from "../handlers/project/migrate/types"; export type { @@ -209,6 +210,8 @@ class ClientCache { private readonly byCredentials = new WeakMap>(); get(config: ClientConfig, create: (config: ClientConfig) => T): T { + const profile = selectedAwsProfile(); + if (!config.credentials && profile) config = { ...config, credentials: profile.credentials }; const clients = this.forCredentials(config.credentials); const key = configKey(config); let client = clients.get(key); diff --git a/src/core/project/backends/cdk/environment.ts b/src/core/project/backends/cdk/environment.ts index 050c240f56..f21c059fd2 100644 --- a/src/core/project/backends/cdk/environment.ts +++ b/src/core/project/backends/cdk/environment.ts @@ -6,6 +6,7 @@ import { import { MalformedServiceResponseError, ProjectStateError } from "../../../../errors/errors"; import type { CreateCloudFormationClient } from "../../../types"; import type { CdkCredentialProvider } from "./toolkit"; +import { selectedAwsProfile } from "../../../../awsProfile"; const BOOTSTRAP_STACK_NAME = "CDKToolkit"; const BOOTSTRAP_VERSION_OUTPUT = "BootstrapVersion"; @@ -122,13 +123,16 @@ export async function probeBootstrap( } } -/** Omitting `credentials` resolves through the default AWS SDK provider chain. */ +/** Use explicit target credentials, the selected profile, or the default AWS SDK chain. */ export const resolveAwsAccount = async ( region: string, credentials?: CdkCredentialProvider, ): Promise => { const { GetCallerIdentityCommand, STSClient } = await import("@aws-sdk/client-sts"); - const client = new STSClient({ credentials, region }); + const client = new STSClient({ + credentials: credentials ?? selectedAwsProfile()?.credentials, + region, + }); try { const { Account } = await client.send(new GetCallerIdentityCommand({})); if (!Account) { diff --git a/src/core/project/backends/cdk/toolkit.ts b/src/core/project/backends/cdk/toolkit.ts index db567459c8..83dfb436df 100644 --- a/src/core/project/backends/cdk/toolkit.ts +++ b/src/core/project/backends/cdk/toolkit.ts @@ -10,6 +10,7 @@ import type { } from "@aws-cdk/toolkit-lib"; import { AgentCoreCLIError } from "../../../../errors"; import type { Logger } from "../../../../logging"; +import { selectedAwsProfile } from "../../../../awsProfile"; export type CdkOperation = | { kind: "bootstrap"; environments: string[]; templateFile?: string } @@ -160,6 +161,8 @@ export async function resolveCdkCredentials( ioHost: IIoHost, region: string, ): Promise { + const profile = selectedAwsProfile(); + if (profile) return profile.credentials; const { BaseCredentials } = await import("@aws-cdk/toolkit-lib"); const config = await BaseCredentials.awsCliCompatible({ defaultRegion: region, diff --git a/src/core/project/bedrockAgentImport/loader.ts b/src/core/project/bedrockAgentImport/loader.ts index eeb76c1887..e81730a586 100644 --- a/src/core/project/bedrockAgentImport/loader.ts +++ b/src/core/project/bedrockAgentImport/loader.ts @@ -10,6 +10,7 @@ import { type AgentActionGroup, type AgentVersion, } from "@aws-sdk/client-bedrock-agent"; +import { selectedAwsProfile } from "../../../awsProfile"; import { InputValidationError, MalformedServiceResponseError, @@ -31,7 +32,7 @@ const ALIAS_ARN_PATTERN = /^arn:[^:]+:bedrock:[^:]+:[^:]+:agent-alias\/([^/]+)\/ export class BedrockAgentSnapshotLoader { constructor( private readonly createClient: CreateBedrockAgentClient = (region) => - new BedrockAgentClient({ region }), + new BedrockAgentClient({ region, credentials: selectedAwsProfile()?.credentials }), ) {} async load(input: { diff --git a/src/core/region.test.ts b/src/core/region.test.ts index 65ce10b3a4..851e5ba690 100644 --- a/src/core/region.test.ts +++ b/src/core/region.test.ts @@ -1,5 +1,16 @@ import { describe, expect, test } from "bun:test"; -import { regionFlagFromArgv } from "./region"; +import { profileFlagFromArgv, regionFlagFromArgv } from "./region"; + +test.each([ + { argv: ["--profile", "development"], expected: "development" }, + { argv: ["--profile=development"], expected: "development" }, + { argv: ["--profile=first", "--profile", "last"], expected: "last" }, + { argv: ["--", "--profile", "argument"], expected: undefined }, + { argv: ["--profile"], expected: undefined }, + { argv: [], expected: undefined }, +])("extracts the effective profile before command routing: %j", ({ argv, expected }) => { + expect(profileFlagFromArgv(["node", "agentcore", ...argv])).toBe(expected); +}); describe("regionFlagFromArgv", () => { test("extracts --region ", () => { diff --git a/src/core/region.ts b/src/core/region.ts index f7c2ba415d..c226cce968 100644 --- a/src/core/region.ts +++ b/src/core/region.ts @@ -7,11 +7,11 @@ export const DEFAULT_REGION = "us-east-1"; /** * Reads the `region` setting for the active profile from the shared AWS config - * file (~/.aws/config, or $AWS_CONFIG_FILE). The active profile is $AWS_PROFILE, - * else "default". Returns undefined if the file, profile, or setting is absent — + * file (~/.aws/config, or $AWS_CONFIG_FILE). An explicit profile takes precedence + * over $AWS_PROFILE, then "default". Returns undefined if the file, profile, or setting is absent — * never throws. */ -export async function regionFromConfigFile(): Promise { +export async function regionFromConfigFile(selectedProfile?: string): Promise { const path = process.env.AWS_CONFIG_FILE || join(homedir(), ".aws", "config"); let text: string; try { @@ -22,7 +22,7 @@ export async function regionFromConfigFile(): Promise { // In ~/.aws/config the default profile is "[default]" and named profiles are // "[profile name]". - const profile = process.env.AWS_PROFILE || "default"; + const profile = selectedProfile ?? (process.env.AWS_PROFILE || "default"); const wanted = profile === "default" ? "default" : `profile ${profile}`; let inSection = false; @@ -45,12 +45,15 @@ export async function regionFromConfigFile(): Promise { * AWS_REGION / AWS_DEFAULT_REGION, then the shared config file, finally * DEFAULT_REGION. It always resolves to a value. */ -export async function resolveRegion(flagRegion: string | undefined): Promise { +export async function resolveRegion( + flagRegion: string | undefined, + profile?: string, +): Promise { return ( flagRegion || process.env.AWS_REGION || process.env.AWS_DEFAULT_REGION || - (await regionFromConfigFile()) || + (await regionFromConfigFile(profile)) || DEFAULT_REGION ); } @@ -71,3 +74,15 @@ export function regionFlagFromArgv(argv: readonly string[]): string | undefined } return undefined; } + +/** Profile selection for region-dependent work that happens before command routing. */ +export function profileFlagFromArgv(argv: readonly string[]): string | undefined { + let profile: string | undefined; + for (let i = 0; i < argv.length; i++) { + const arg = argv[i]!; + if (arg === "--") break; + if (arg === "--profile") profile = argv[++i]; + else if (arg.startsWith("--profile=")) profile = arg.slice("--profile=".length); + } + return profile; +} diff --git a/src/handlers/index.tsx b/src/handlers/index.tsx index ee174b930b..54e4d5dc6a 100644 --- a/src/handlers/index.tsx +++ b/src/handlers/index.tsx @@ -7,13 +7,14 @@ import { createIdentityHandler } from "./identity/index.tsx"; import { createMemoryHandler } from "./memory/index.tsx"; import { createPaymentHandler } from "./payment/index.tsx"; import { createRuntimeHandler } from "./runtime/index.tsx"; -import { DebugKey, JsonKey, RegionKey } from "./keys.tsx"; +import { DebugKey, JsonKey, ProfileKey, RegionKey } from "./keys.tsx"; import { createConfigHandler } from "./config/"; import { createProjectHandlers } from "./project/index.ts"; import { createUpdateHandler } from "./update/index.tsx"; import { renderTui } from "../tui"; import { withRegion, + withProfile, withJsonRenderer, withLogging, withGlobalConfigAccessor, @@ -68,12 +69,15 @@ export function createRootHandler(core: Core, config: RootHandlerConfig): Router // --endpoint-url intentionally omitted: temporarily disabled (inconsistent // override behavior). EndpointKey + its plumbing stay defined so re-enabling // is adding it back here. - root.groupFlags(RegionKey, DebugKey, JsonKey); + root.groupFlags(RegionKey, ProfileKey, DebugKey, JsonKey); // Resolve the effective AWS region (flag -> env -> config file) and pin it on // the context for every command beneath the root. root.use(withRegion()); + // Credential providers such as STS assume-role also use the effective region. + root.use(withProfile()); + // Pin a JSON renderer wired to the configured stdout so leaf handlers can emit // machine-readable output without touching the process streams directly. root.use(withJsonRenderer(io)); diff --git a/src/handlers/keys.tsx b/src/handlers/keys.tsx index ff97400761..2b515dd7aa 100644 --- a/src/handlers/keys.tsx +++ b/src/handlers/keys.tsx @@ -11,6 +11,15 @@ import { contextKey, globalFlag } from "../router"; // withRegion middleware can fall back to AWS_REGION / the shared config file. export const RegionKey = globalFlag("region", "AWS region", z.string().optional()); +export const ProfileKey = globalFlag( + "profile", + "AWS shared configuration profile", + z + .string() + .refine((value) => value.trim().length > 0, "Profile name must not be empty") + .optional(), +); + // withRegion overwrites RegionKey with the fully-resolved region (flag -> env -> // config -> DEFAULT_REGION) before any handler runs, so by the time a leaf reads // ctx.value(RegionKey) it's never undefined. This key preserves whether --region diff --git a/src/io/exec.ts b/src/io/exec.ts index 11ef4e32a3..d98e09a984 100644 --- a/src/io/exec.ts +++ b/src/io/exec.ts @@ -5,6 +5,7 @@ import { constants as osConstants } from "node:os"; import { createInterface } from "node:readline"; import crossSpawn from "cross-spawn"; import { AgentCoreCLIError, ERROR_SOURCE } from "../errors"; +import { awsProfileEnvironment, selectedAwsProfile } from "../awsProfile"; const isWindows = process.platform === "win32"; const KILL_GRACE_MS = 2000; @@ -33,7 +34,10 @@ export class ProcessFailedError extends AgentCoreCLIError { function spawnCommand(executable: string, args: string[], options: SpawnOptions): ChildProcess { // Keep arguments separate. cross-spawn resolves Windows wrappers and applies // the extra escaping required when a node_modules/.bin shim forwards them. - return crossSpawn(executable, args, options); + return crossSpawn(executable, args, { + ...options, + env: selectedAwsProfile() ? awsProfileEnvironment(options.env) : options.env, + }); } /** Returns true if running `tool` with probeArgs (`--version` by default) exits 0. */ diff --git a/src/middleware/index.tsx b/src/middleware/index.tsx index fff26a14d6..21a3f26865 100644 --- a/src/middleware/index.tsx +++ b/src/middleware/index.tsx @@ -1,4 +1,5 @@ export { withRegion } from "./withRegion"; +export { withProfile } from "./withProfile"; export { withTuiOnEmptyFlagsAndArgs, withTuiWhenInteractive } from "./withTuiOnEmptyFlagsAndArgs"; export { withJsonRenderer } from "./withJsonRenderer"; export { withLogging } from "./withLogging"; diff --git a/src/middleware/withProfile.test.ts b/src/middleware/withProfile.test.ts new file mode 100644 index 0000000000..3c340be5d8 --- /dev/null +++ b/src/middleware/withProfile.test.ts @@ -0,0 +1,127 @@ +import { afterEach, beforeEach, expect, test } from "bun:test"; +import { mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { createRootHandler } from "../handlers"; +import { createSilentLogger, TestCoreClient, TestGlobalConfigAccessor, testIO } from "../testing"; +import type { CoreOptions } from "../core/types"; + +const keys = [ + "AWS_CONFIG_FILE", + "AWS_SHARED_CREDENTIALS_FILE", + "AWS_PROFILE", + "AWS_REGION", + "AWS_DEFAULT_REGION", +] as const; +let saved: NodeJS.ProcessEnv; +let directory: string; +beforeEach(async () => { + saved = Object.fromEntries(keys.map((key) => [key, process.env[key]])); + directory = await mkdtemp(join(tmpdir(), "agentcore-profile-route-")); + process.env.AWS_CONFIG_FILE = join(directory, "config"); + process.env.AWS_SHARED_CREDENTIALS_FILE = join(directory, "credentials"); + process.env.AWS_PROFILE = "ambient"; + delete process.env.AWS_REGION; + delete process.env.AWS_DEFAULT_REGION; + await writeFile( + process.env.AWS_CONFIG_FILE, + "[profile chosen]\nregion = eu-west-1\n[profile ambient]\nregion = ap-south-1\n", + ); + await writeFile( + process.env.AWS_SHARED_CREDENTIALS_FILE, + "[chosen]\naws_access_key_id = selected-key\naws_secret_access_key = selected-secret\n", + ); +}); +afterEach(async () => { + for (const key of keys) { + if (saved[key] === undefined) delete process.env[key]; + else process.env[key] = saved[key]; + } + await rm(directory, { recursive: true, force: true }); +}); + +async function route(args: readonly string[]): Promise { + const core = new TestCoreClient(); + const root = createRootHandler(core, { + io: testIO().io, + logger: createSilentLogger(), + globalConfigAccessor: new TestGlobalConfigAccessor(), + }); + await root.route(["node", "agentcore", ...args]); + return core.harness.calls.at(-1)!.args[2] as CoreOptions; +} + +test.each([ + { args: ["--profile", "chosen", "harness", "list", "--json"] }, + { args: ["harness", "list", "--json", "--profile", "chosen"] }, + { args: ["harness", "list", "--json", "--profile=chosen"] }, +])("global profile works before or after the command: %j", async ({ args }) => { + const options = await route(args); + expect(options.region).toBe("eu-west-1"); + expect( + (await (options.credentials as () => Promise<{ accessKeyId: string }>)()).accessKeyId, + ).toBe("selected-key"); + expect(process.env.AWS_PROFILE).toBe("ambient"); +}); + +test("region flags and environment retain precedence over the chosen profile", async () => { + process.env.AWS_REGION = "us-west-2"; + expect((await route(["harness", "list", "--json", "--profile", "chosen"])).region).toBe( + "us-west-2", + ); + expect( + (await route(["harness", "list", "--json", "--profile", "chosen", "--region", "ca-central-1"])) + .region, + ).toBe("ca-central-1"); +}); + +test("assume-role profiles use the effective region for STS", async () => { + await writeFile( + process.env.AWS_CONFIG_FILE!, + "[profile role]\nregion = eu-west-1\nrole_arn = arn:aws:iam::123456789012:role/test\nsource_profile = chosen\n", + ); + let authorization = ""; + const server = Bun.serve({ + port: 0, + fetch(request) { + authorization = request.headers.get("authorization") ?? ""; + return new Response( + 'role-keyrole-secretrole-token2099-01-01T00:00:00Z', + { headers: { "content-type": "text/xml" } }, + ); + }, + }); + const savedEndpoint = process.env.AWS_ENDPOINT_URL_STS; + process.env.AWS_ENDPOINT_URL_STS = `http://localhost:${server.port}`; + try { + const options = await route([ + "harness", + "list", + "--json", + "--profile", + "role", + "--region", + "ca-central-1", + ]); + expect( + (await (options.credentials as () => Promise<{ accessKeyId: string }>)()).accessKeyId, + ).toBe("role-key"); + expect(authorization).toContain("/ca-central-1/sts/aws4_request"); + } finally { + if (savedEndpoint === undefined) delete process.env.AWS_ENDPOINT_URL_STS; + else process.env.AWS_ENDPOINT_URL_STS = savedEndpoint; + server.stop(true); + } +}); + +test.each(["", " "])("rejects an empty profile: %j", async (profile) => { + await expect(route(["harness", "list", "--json", "--profile", profile])).rejects.toThrow( + "Profile name must not be empty", + ); +}); + +test("missing --profile leaves ambient credentials and region behavior unchanged", async () => { + const options = await route(["harness", "list", "--json"]); + expect(options.region).toBe("ap-south-1"); + expect(options.credentials).toBeUndefined(); +}); diff --git a/src/middleware/withProfile.ts b/src/middleware/withProfile.ts new file mode 100644 index 0000000000..71477412db --- /dev/null +++ b/src/middleware/withProfile.ts @@ -0,0 +1,28 @@ +import { AwsCredentialProviderKey, ProfileKey, RegionKey } from "../handlers/keys"; +import { selectedAwsProfile, withAwsProfile } from "../awsProfile"; +import type { Middleware } from "../router"; + +export function withProfile(): Middleware { + return (handler) => ({ + name: () => handler.name(), + description: () => handler.description(), + flags: () => handler.flags(), + arguments: () => handler.arguments(), + doesSupportTui: () => handler.doesSupportTui(), + children: () => handler.children(), + handle: async (ctx, flags, args) => { + const profile = ctx.value(ProfileKey); + if (profile === undefined) return handler.handle(ctx, flags, args); + await withAwsProfile( + profile, + () => + handler.handle( + ctx.withValue(AwsCredentialProviderKey, selectedAwsProfile()!.credentials), + flags, + args, + ), + ctx.value(RegionKey), + ); + }, + }); +} diff --git a/src/middleware/withRegion.tsx b/src/middleware/withRegion.tsx index e63a75480f..3a7a0b759d 100644 --- a/src/middleware/withRegion.tsx +++ b/src/middleware/withRegion.tsx @@ -1,4 +1,4 @@ -import { RegionExplicitKey, RegionKey } from "../handlers/keys"; +import { ProfileKey, RegionExplicitKey, RegionKey } from "../handlers/keys"; import { resolveRegion } from "../core/region"; import { type Middleware } from "../router"; @@ -17,7 +17,7 @@ export function withRegion(): Middleware { children: () => h.children(), handle: async (ctx, flags, args) => { const explicit = ctx.value(RegionKey) !== undefined; - const region = await resolveRegion(ctx.value(RegionKey)); + const region = await resolveRegion(ctx.value(RegionKey), ctx.value(ProfileKey)); await h.handle( ctx.withValue(RegionKey, region).withValue(RegionExplicitKey, explicit), flags, diff --git a/src/telemetry/client.test.tsx b/src/telemetry/client.test.tsx index cfe0eb858a..6de42e418b 100644 --- a/src/telemetry/client.test.tsx +++ b/src/telemetry/client.test.tsx @@ -1,6 +1,6 @@ import { test, describe, beforeEach, afterEach, expect } from "bun:test"; import { join } from "node:path"; -import { mkdtemp, readFile, rm } from "node:fs/promises"; +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; import os, { tmpdir } from "node:os"; import { DefaultTelemetryClient } from "./client"; import { createFileLogger, type Logger } from "../logging"; @@ -479,6 +479,52 @@ describe("OtelHistogramSink", () => { await rm(auditFilePath, { force: true }); }); + test.each([ + { flags: ["--profile", "china"] }, + { flags: ["--profile=china"] }, + { flags: ["--profile=ambient", "--profile", "china"] }, + ])("a selected China profile suppresses network export: %j", async ({ flags }) => { + const directory = await mkdtemp(join(tmpdir(), "telemetry-profile-")); + const keys = ["AWS_CONFIG_FILE", "AWS_REGION", "AWS_DEFAULT_REGION"] as const; + const saved = Object.fromEntries(keys.map((key) => [key, process.env[key]])); + process.env.AWS_CONFIG_FILE = join(directory, "config"); + delete process.env.AWS_REGION; + delete process.env.AWS_DEFAULT_REGION; + await writeFile(process.env.AWS_CONFIG_FILE, "[profile china]\nregion = cn-northwest-1\n"); + const client = new DefaultTelemetryClient({ + logger, + sessionId: "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee", + globalConfigAccessor: new TestGlobalConfigAccessor({ + initialConfigData: { + ...DEFAULT_GLOBAL_CONFIG, + telemetry: { + enabled: true, + audit: false, + endpoint: `http://localhost:${testCollector.port}`, + }, + }, + }), + argv: ["node", "agentcore", "status", ...flags], + currentVersion: PACKAGE_VERSION, + }); + try { + await client + .createMetricEvent("cli.command_run", { + exit_reason: "success", + command_path: "/agentcore", + }) + .emit(100); + await client.shutdown(); + expect(receivedBodies).toHaveLength(0); + } finally { + for (const key of keys) { + if (saved[key] === undefined) delete process.env[key]; + else process.env[key] = saved[key]; + } + await rm(directory, { recursive: true, force: true }); + } + }); + test("a China region suppresses network export even when telemetry is enabled", async () => { const auditFilePath = join(tmpdir(), `cn-disabled-audit-${crypto.randomUUID()}.jsonl`); const globalConfigAccessor = new TestGlobalConfigAccessor({ diff --git a/src/telemetry/client.tsx b/src/telemetry/client.tsx index 498259d40d..0a888a159b 100644 --- a/src/telemetry/client.tsx +++ b/src/telemetry/client.tsx @@ -15,7 +15,7 @@ import { FileSystemSink } from "./fileSystemSink"; import path from "path"; import { OtelHistogramSink } from "./otelSink"; import { isChinaContext } from "../core/partition"; -import { regionFlagFromArgv, resolveRegion } from "../core/region"; +import { profileFlagFromArgv, regionFlagFromArgv, resolveRegion } from "../core/region"; export type DefaultTelemetryClientConfig = { logger: Logger; @@ -101,7 +101,10 @@ export class DefaultTelemetryClient implements TelemetryClient { // resolved through the same chain the CLI itself uses (--region flag, env // vars, shared config file). The audit file sink above is unaffected (it // only writes locally). - const resolvedRegion = await resolveRegion(regionFlagFromArgv(this.argv)); + const resolvedRegion = await resolveRegion( + regionFlagFromArgv(this.argv), + profileFlagFromArgv(this.argv), + ); if ( globalConfig.telemetry.enabled && !telemetryDisabledByEnv() &&