diff --git a/command.md b/command.md index 7e24899b5..64054bc46 100644 --- a/command.md +++ b/command.md @@ -1045,6 +1045,7 @@ invoke a harness - `--user-id `: the end-user ID for user-scoped AgentCore Identity credentials - `--session-id `: the Runtime session ID to continue (33-100 characters) - `--qualifier `: the harness endpoint qualifier to invoke (default DEFAULT) +- `--bearer-token `: the CUSTOM\_JWT bearer token, from stdin ('-') or a file ('file://<path>') #### `agentcore harness exec` diff --git a/src/core/core.test.ts b/src/core/core.test.ts index b4135e9e5..125e96c98 100644 --- a/src/core/core.test.ts +++ b/src/core/core.test.ts @@ -23,7 +23,7 @@ import { import type { RuntimeInvokeRequest } from "../handlers/runtime/types"; import type { Logger, LoggerBindings } from "../logging"; import { CoreClient } from "./index"; -import type { ClientConfig, CoreFetch } from "./types"; +import type { ClientConfig, CoreFetch, DataClientConfig } from "./types"; import { toClientConfig } from "./utils"; import { createSilentLogger } from "../testing"; @@ -216,6 +216,31 @@ test("data() caches independently of control()", () => { expect(dataBuilt).toBe(1); }); +test("data() never caches a bearer-token client or hands one to a SigV4 caller", () => { + const configs: DataClientConfig[] = []; + const core = new CoreClient({ + createControlClient: fakeControl, + createDataClient: (config) => { + configs.push(config); + return fakeData(config); + }, + createIamClient: fakeIam, + createLogsClient: fakeLogs, + createXrayClient: fakeXray, + createApplicationSignalsClient: fakeApplicationSignals, + logger: createSilentLogger(), + }); + + const b1 = core.data({ region: "us-east-1", bearerToken: "jwt" }); + const b2 = core.data({ region: "us-east-1", bearerToken: "jwt" }); + const sigv4 = core.data({ region: "us-east-1" }); + + expect(b1).not.toBe(b2); + expect(sigv4).not.toBe(b1); + expect(core.data({ region: "us-east-1" })).toBe(sigv4); + expect(configs.map((config) => config.bearerToken)).toEqual(["jwt", "jwt", undefined]); +}); + test("exposes feature sub-clients", () => { const core = new CoreClient({ createControlClient: fakeControl, diff --git a/src/core/factories.test.ts b/src/core/factories.test.ts new file mode 100644 index 000000000..f835a9b70 --- /dev/null +++ b/src/core/factories.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, test } from "bun:test"; +import { InvokeHarnessCommand } from "@aws-sdk/client-bedrock-agentcore"; +import { createDataClient } from "./factories"; +import type { DataClientConfig } from "./types"; + +const HARNESS_ARN = "arn:aws:bedrock-agentcore:us-east-1:123456789012:harness/MyHarness-abc123"; + +// signedHeaders sends an InvokeHarness through a real data client and returns the +// headers of the request as it would leave for the wire, after signing. +async function signedHeaders(config: DataClientConfig): Promise> { + const client = createDataClient(config); + let headers: Record | undefined; + client.middlewareStack.add( + () => async (args) => { + headers = (args.request as { headers: Record }).headers; + throw new Error("request captured"); + }, + { step: "deserialize", priority: "low" }, + ); + await expect( + client.send( + new InvokeHarnessCommand({ + harnessArn: HARNESS_ARN, + runtimeSessionId: "s".repeat(33), + messages: [{ role: "user", content: [{ text: "hi" }] }], + }), + ), + ).rejects.toThrow("request captured"); + return headers!; +} + +describe("createDataClient", () => { + test("a bearer token authenticates with Authorization: Bearer and never resolves AWS credentials", async () => { + const headers = await signedHeaders({ + region: "us-east-1", + bearerToken: "jwt-123", + credentials: () => { + throw new Error("AWS credentials must not be resolved"); + }, + }); + + expect(headers["Authorization"] ?? headers["authorization"]).toBe("Bearer jwt-123"); + expect(Object.keys(headers).map((name) => name.toLowerCase())).not.toContain("x-amz-date"); + }); + + test("without a bearer token the request is SigV4-signed", async () => { + const headers = await signedHeaders({ + region: "us-east-1", + credentials: { accessKeyId: "AKIDEXAMPLE", secretAccessKey: "secret" }, + }); + + expect(headers["authorization"]).toStartWith("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/"); + }); +}); diff --git a/src/core/factories.tsx b/src/core/factories.tsx index ccdda4b8d..c57fa0e7b 100644 --- a/src/core/factories.tsx +++ b/src/core/factories.tsx @@ -5,6 +5,8 @@ import { CloudWatchLogsClient } from "@aws-sdk/client-cloudwatch-logs"; import { CloudFormationClient } from "@aws-sdk/client-cloudformation"; import { XRayClient } from "@aws-sdk/client-xray"; import { ApplicationSignalsClient } from "@aws-sdk/client-application-signals"; +import { HttpBearerAuthSigner } from "@smithy/core"; +import type { TokenIdentity } from "@smithy/types"; import type { CreateApplicationSignalsClient, CreateCloudFormationClient, @@ -23,8 +25,22 @@ import type { export const createControlClient: CreateControlClient = (config) => new BedrockAgentCoreControlClient({ ...config }); -export const createDataClient: CreateDataClient = (config) => - new BedrockAgentCoreClient({ ...config }); +// A bearer token swaps SigV4 for Smithy's HTTP bearer auth scheme, so the client +// sends `Authorization: Bearer ` and never resolves AWS credentials. +export const createDataClient: CreateDataClient = ({ bearerToken, ...config }) => + new BedrockAgentCoreClient({ + ...config, + ...(bearerToken !== undefined && { + httpAuthSchemes: [ + { + schemeId: "smithy.api#httpBearerAuth", + identityProvider: () => async (): Promise => ({ token: bearerToken }), + signer: new HttpBearerAuthSigner(), + }, + ], + httpAuthSchemeProvider: () => [{ schemeId: "smithy.api#httpBearerAuth" }], + }), + }); export const createIamClient: CreateIamClient = (config) => new IAMClient({ ...config }); diff --git a/src/core/harness.test.ts b/src/core/harness.test.ts index da354171c..5030c9217 100644 --- a/src/core/harness.test.ts +++ b/src/core/harness.test.ts @@ -3,8 +3,9 @@ import { GetHarnessCommand, type GetHarnessResponse, } from "@aws-sdk/client-bedrock-agentcore-control"; +import { InvokeHarnessCommand } from "@aws-sdk/client-bedrock-agentcore"; import { InputValidationError } from "../errors"; -import type { AwsClients, ClientConfig } from "./types"; +import type { AwsClients, ClientConfig, DataClientConfig } from "./types"; import { HarnessClient, harnessRuntimeFromResponse } from "./harness"; const RESPONSE = { @@ -57,3 +58,45 @@ describe("HarnessClient.resolveRuntime", () => { expect(configs).toEqual([{ region: "us-west-2", endpoint: "https://control.test" }]); }); }); + +describe("HarnessClient.invokeHarness", () => { + const REQUEST = { + harnessArn: "arn:aws:bedrock-agentcore:us-west-2:123456789012:harness/MyHarness-abc123", + runtimeSessionId: "s".repeat(33), + messages: [{ role: "user" as const, content: [{ text: "hi" }] }], + }; + + async function invoke(bearerToken?: string) { + const configs: DataClientConfig[] = []; + const inputs: unknown[] = []; + const clients = { + data: (config: DataClientConfig) => { + configs.push(config); + return { + send: async (command: unknown) => { + expect(command).toBeInstanceOf(InvokeHarnessCommand); + inputs.push((command as InvokeHarnessCommand).input); + return {}; + }, + }; + }, + } as unknown as AwsClients; + await new HarnessClient(clients).invokeHarness( + { ...REQUEST, ...(bearerToken !== undefined && { bearerToken }) }, + { region: "us-west-2" }, + ); + return { configs, inputs }; + } + + test("routes a bearer token to the data client config, not the request", async () => { + const { configs, inputs } = await invoke("jwt-123"); + expect(configs).toEqual([{ region: "us-west-2", bearerToken: "jwt-123" }]); + expect(inputs).toEqual([REQUEST]); + }); + + test("uses the SigV4 client without a bearer token", async () => { + const { configs, inputs } = await invoke(); + expect(configs).toEqual([{ region: "us-west-2" }]); + expect(inputs).toEqual([REQUEST]); + }); +}); diff --git a/src/core/harness.tsx b/src/core/harness.tsx index ef3f9d598..def41c9f9 100644 --- a/src/core/harness.tsx +++ b/src/core/harness.tsx @@ -32,12 +32,12 @@ import { InvokeHarnessCommand, type InvokeAgentRuntimeCommandRequest, type InvokeAgentRuntimeCommandResponse, - type InvokeHarnessRequest, type InvokeHarnessResponse, } from "@aws-sdk/client-bedrock-agentcore"; import type { CoreHarnessClient, CreateHarnessInput, + HarnessInvokeRequest, ResolvedHarnessRuntime, HarnessShellRequest, } from "../handlers/harness/types"; @@ -206,12 +206,12 @@ export class HarnessClient implements CoreHarnessClient { } async invokeHarness( - request: InvokeHarnessRequest, + { bearerToken, ...request }: HarnessInvokeRequest, options: CoreOptions, abortSignal?: AbortSignal, ): Promise { const response = await this.clients - .data(toClientConfig(options)) + .data({ ...toClientConfig(options), ...(bearerToken !== undefined && { bearerToken }) }) .send(new InvokeHarnessCommand(request), { abortSignal }); if (!response.stream || !abortSignal) return response; // The SDK honors the abort signal while the request is being established, diff --git a/src/core/index.tsx b/src/core/index.tsx index f2999d0df..0293c4fbc 100644 --- a/src/core/index.tsx +++ b/src/core/index.tsx @@ -27,6 +27,7 @@ import type { CreateIamClient, CreateLogsClient, CreateXrayClient, + DataClientConfig, } from "./types"; import type { Logger } from "../logging"; import type { ProjectManager } from "../handlers/project/types"; @@ -46,6 +47,7 @@ export type { CreateIamClient, CreateLogsClient, CreateXrayClient, + DataClientConfig, } from "./types"; type CoreClientConfig = { @@ -166,8 +168,10 @@ export class CoreClient implements AwsClients { } // data returns the data-plane client for `config`, creating and caching it on - // first use. - data(config: ClientConfig): BedrockAgentCoreClient { + // first use. A bearer-token client is built fresh each time: caching it would + // keep the token alive and risk reusing it for a SigV4 caller. + data(config: DataClientConfig): BedrockAgentCoreClient { + if (config.bearerToken !== undefined) return this.createDataClient(config); return this.dataClients.get(config, this.createDataClient); } diff --git a/src/core/types.tsx b/src/core/types.tsx index f7368d235..f1e125db7 100644 --- a/src/core/types.tsx +++ b/src/core/types.tsx @@ -39,6 +39,13 @@ export interface ClientConfig { credentials?: AwsCredentials; } +// DataClientConfig is a ClientConfig for the data plane, which alone can +// authenticate with a CUSTOM_JWT bearer token instead of SigV4. A bearer client +// never resolves AWS credentials, so callers holding only a JWT can use it. +export interface DataClientConfig extends ClientConfig { + bearerToken?: string; +} + export type CredentialedClientConfig = { region: string; credentials: AwsCredentials; @@ -48,7 +55,7 @@ export type CredentialedClientConfig = { // than the clients themselves) lets CoreClient create/cache one client per config // while keeping construction swappable for unit tests. export type CreateControlClient = (config: ClientConfig) => BedrockAgentCoreControlClient; -export type CreateDataClient = (config: ClientConfig) => BedrockAgentCoreClient; +export type CreateDataClient = (config: DataClientConfig) => BedrockAgentCoreClient; export type CreateIamClient = (config: ClientConfig) => IAMClient; export type CreateLogsClient = (config: ClientConfig) => CloudWatchLogsClient; export type CreateXrayClient = (config: ClientConfig) => XRayClient; @@ -65,7 +72,7 @@ export type CoreFetch = ( // endpoint, ...). export interface AwsClients { control(config: ClientConfig): BedrockAgentCoreControlClient; - data(config: ClientConfig): BedrockAgentCoreClient; + data(config: DataClientConfig): BedrockAgentCoreClient; iam(config: ClientConfig): IAMClient; // logs reads the CloudWatch Logs streams AgentCore writes batch-evaluation // results to. CloudWatch is a distinct service from the AgentCore data plane, diff --git a/src/handlers/harness/invoke/index.tsx b/src/handlers/harness/invoke/index.tsx index 1c091eb0e..ecd3184f8 100644 --- a/src/handlers/harness/invoke/index.tsx +++ b/src/handlers/harness/invoke/index.tsx @@ -1,14 +1,14 @@ import z from "zod"; import { createHandler, flag, PathKey } from "../../../router"; -import type { AppIO } from "../../../io"; +import { SourceResolver, type AppIO } from "../../../io"; import type { Core } from "../../types.tsx"; import { coreOptsFromCtx } from "../../utils.tsx"; -import { JsonKey } from "../../keys.tsx"; +import { AwsCredentialProviderKey, JsonKey } from "../../keys.tsx"; import { JsonRendererKey, renderTuiAt } from "../../../tui"; import { runWithProgress } from "../../../tui/progress"; -import { InputValidationError } from "../../../errors"; +import { InputValidationError, SourceResolutionError } from "../../../errors"; import { invokeHarnessTurn } from "./operation.ts"; -import { HarnessInvokeLaunchContextKey } from "./launchContext"; +import { HarnessInvokeArnKey, HarnessInvokeLaunchContextKey } from "./launchContext"; export const invokeHarnessFlags = [ flag("id", "the ID of the harness", z.string().min(1).max(51)), @@ -28,6 +28,12 @@ export const invokeHarnessFlags = [ "the harness endpoint qualifier to invoke (default DEFAULT)", z.string().optional(), ), + flag( + "bearer-token", + "the CUSTOM_JWT bearer token, from stdin ('-') or a file ('file://')", + z.string().optional(), + { sensitive: true }, + ), ] as const; export const createInvokeHarnessHandler = (core: Core, io: AppIO) => @@ -44,6 +50,9 @@ export const createInvokeHarnessHandler = (core: Core, io: AppIO) => if (ctx.require(JsonKey)) { throw new InputValidationError("required option '--prompt ' not specified"); } + if (flags["bearer-token"] !== undefined) { + throw new InputValidationError("--bearer-token requires --prompt"); + } let path = `${ctx.require(PathKey)}/${flags["id"]}`; if (flags["session-id"]) path += `/${flags["session-id"]}`; if (flags["qualifier"]) path += `?qualifier=${encodeURIComponent(flags["qualifier"])}`; @@ -59,17 +68,33 @@ export const createInvokeHarnessHandler = (core: Core, io: AppIO) => return; } + // Without the ARN the invoke must look the harness up with GetHarness. A + // project harness gets that from its target's credentials; otherwise it + // would need the default AWS credentials a bearer-token caller lacks. + const harnessArn = ctx.value(HarnessInvokeArnKey); + if ( + flags["bearer-token"] !== undefined && + harnessArn === undefined && + ctx.value(AwsCredentialProviderKey) === undefined + ) { + throw new InputValidationError( + "--bearer-token needs the harness ARN; use 'agentcore invoke --harness --bearer-token ...'", + ); + } const opts = coreOptsFromCtx(ctx); const prompt = flags["prompt"]; + const bearerToken = await resolveBearerToken(flags["bearer-token"], io.stdin); const invoke = () => invokeHarnessTurn( core.harness, { harnessId: flags["id"], + harnessArn, prompt, qualifier: flags["qualifier"] ?? "DEFAULT", sessionId: flags["session-id"], userId: flags["user-id"], + bearerToken, }, opts, ); @@ -82,4 +107,18 @@ export const createInvokeHarnessHandler = (core: Core, io: AppIO) => }, }); +async function resolveBearerToken( + source: string | undefined, + stdin: NodeJS.ReadStream, +): Promise { + try { + return await new SourceResolver({ stdin }).resolveSecret("bearer-token", source); + } catch (error) { + if (error instanceof SourceResolutionError) { + throw new InputValidationError(error.message, { cause: error }); + } + throw error; + } +} + export { HarnessInvokeScreen } from "./screen.tsx"; diff --git a/src/handlers/harness/invoke/invoke.test.tsx b/src/handlers/harness/invoke/invoke.test.tsx index 2d8e3cae5..db1fc707b 100644 --- a/src/handlers/harness/invoke/invoke.test.tsx +++ b/src/handlers/harness/invoke/invoke.test.tsx @@ -1,4 +1,6 @@ import { test, expect, describe, spyOn } from "bun:test"; +import { writeFile } from "node:fs/promises"; +import { join } from "node:path"; import type { InvokeHarnessRequest, InvokeHarnessStreamOutput, @@ -8,6 +10,7 @@ import { createRootHandler } from "../../index"; import { createSilentLogger, expectError, + inTempDirectory, TestCoreClient, TestGlobalConfigAccessor, testIO, @@ -17,6 +20,7 @@ import { } from "../../../testing"; import { InputValidationError } from "../../../errors"; import * as tui from "../../../tui"; +import type { HarnessInvokeRequest } from "../types"; import { HarnessInvokeLaunchContextKey } from "./launchContext"; // Command-flow tests for `harness invoke`, driven through the real root handler @@ -224,6 +228,86 @@ describe("harness invoke", () => { ]); }); + // A bearer-token invoke needs the ARN, which `agentcore invoke --harness ` + // hands to this handler; `harness invoke --id` alone cannot supply it. + const BEARER_INVOKE = ["invoke", "--harness", ARN, "--prompt", "hi"]; + + test("--bearer-token - reads the JWT from stdin and invokes the ARN without GetHarness", async () => { + const { core } = await run( + [...BEARER_INVOKE, "--bearer-token", "-"], + undefined, + testIO({ stdin: "jwt-123\n" }), + ); + + expect(core.harness.calls.map((c) => c.method)).toEqual(["invokeHarness"]); + const request = core.harness.calls[0]!.args[0] as HarnessInvokeRequest; + expect(request.bearerToken).toBe("jwt-123"); + expect(request.harnessArn).toBe(ARN); + }); + + test("--bearer-token file:// reads the JWT from a file", async () => { + const { path, cleanup } = await inTempDirectory("agentcore-harness-invoke-"); + try { + const tokenFile = join(path, "token"); + await writeFile(tokenFile, "jwt-from-file\n"); + const { core } = await run([...BEARER_INVOKE, "--bearer-token", `file://${tokenFile}`]); + + const invoke = core.harness.calls.find((c) => c.method === "invokeHarness")!; + expect((invoke.args[0] as HarnessInvokeRequest).bearerToken).toBe("jwt-from-file"); + } finally { + await cleanup(); + } + }); + + test("a non-DEFAULT --qualifier reaches InvokeHarness with --bearer-token", async () => { + const { core } = await run( + [...BEARER_INVOKE, "--qualifier", "prod", "--bearer-token", "-"], + undefined, + testIO({ stdin: "jwt-123\n" }), + ); + + const invoke = core.harness.calls.find((c) => c.method === "invokeHarness")!; + const request = invoke.args[0] as HarnessInvokeRequest; + expect(request.qualifier).toBe("prod"); + expect(request.bearerToken).toBe("jwt-123"); + }); + + test("an inline --bearer-token is rejected", async () => { + await expectError( + run([...BEARER_INVOKE, "--bearer-token", "jwt"]), + /inline secret values are not accepted/, + InputValidationError, + ); + }); + + test("--bearer-token with only --id is rejected up front, pointing to invoke --harness ", async () => { + let core: TestCoreClient | undefined; + await expectError( + run( + ["harness", "invoke", "--id", "MyHarness-abc123", "--prompt", "hi", "--bearer-token", "-"], + (c) => (core = c), + testIO({ stdin: "jwt-123\n" }), + ), + /agentcore invoke --harness --bearer-token/, + InputValidationError, + ); + expect(core!.harness.calls).toEqual([]); + }); + + test("--bearer-token without --prompt is rejected instead of opening the TUI", async () => { + const render = spyOn(tui, "renderTuiAt").mockResolvedValue(undefined); + try { + await expectError( + run(["harness", "invoke", "--id", "MyHarness-abc123", "--bearer-token", "-"]), + /--bearer-token requires --prompt/, + InputValidationError, + ); + expect(render).not.toHaveBeenCalled(); + } finally { + render.mockRestore(); + } + }); + test("errors when --id is omitted", async () => { await expectError(run(["harness", "invoke", "--prompt", "hi"]), /--id/, InputValidationError); }); diff --git a/src/handlers/harness/invoke/launchContext.ts b/src/handlers/harness/invoke/launchContext.ts index 44dc56d9f..8795489dd 100644 --- a/src/handlers/harness/invoke/launchContext.ts +++ b/src/handlers/harness/invoke/launchContext.ts @@ -7,3 +7,7 @@ export type HarnessInvokeLaunchContext = { export const HarnessInvokeLaunchContextKey = contextKey("harness.invoke.launch"); + +// HarnessInvokeArnKey carries the full ARN when `agentcore invoke --harness ` +// delegates here, so the invoke can skip GetHarness. `--id` takes only an ID. +export const HarnessInvokeArnKey = contextKey("harness.invoke.arn"); diff --git a/src/handlers/harness/invoke/operation.ts b/src/handlers/harness/invoke/operation.ts index 334b66bed..e97c3c577 100644 --- a/src/handlers/harness/invoke/operation.ts +++ b/src/handlers/harness/invoke/operation.ts @@ -14,23 +14,29 @@ export async function invokeHarnessTurn( client: CoreHarnessClient, input: { harnessId: string; + // A caller that already holds the ARN skips the GetHarness lookup, which is + // what lets a bearer-token invoke run without AWS credentials. + harnessArn?: string; prompt: string; qualifier?: string; sessionId?: string; userId?: string; + bearerToken?: string; }, options: CoreOptions, signal?: AbortSignal, ): Promise { - const detail = await client.getHarness(input.harnessId, options); + const harnessArn = + input.harnessArn ?? (await client.getHarness(input.harnessId, options)).harness?.arn; const sessionId = input.sessionId ?? newSessionId(); const response = await client.invokeHarness( { - harnessArn: detail.harness?.arn, + harnessArn, qualifier: input.qualifier ?? "DEFAULT", runtimeSessionId: sessionId, runtimeUserId: input.userId, messages: [{ role: "user", content: [{ text: input.prompt }] }], + ...(input.bearerToken !== undefined && { bearerToken: input.bearerToken }), }, options, signal, diff --git a/src/handlers/harness/types.tsx b/src/handlers/harness/types.tsx index 043ce5df8..6559754fd 100644 --- a/src/handlers/harness/types.tsx +++ b/src/handlers/harness/types.tsx @@ -35,6 +35,11 @@ export type HarnessShellRequest = { onReconnect?: () => void | Promise; }; +// HarnessInvokeRequest is InvokeHarnessRequest plus an optional bearer token for +// a CUSTOM_JWT harness. With a token, the invoke authenticates with +// `Authorization: Bearer` instead of SigV4 and needs no AWS credentials. +export type HarnessInvokeRequest = InvokeHarnessRequest & { bearerToken?: string }; + // CreateHarnessInput is CreateHarnessRequest with the execution role made // optional: when omitted, Core provisions the default execution role in IAM and // creates the harness with it. @@ -103,7 +108,7 @@ export interface CoreHarnessClient { options: CoreOptions, ): Promise; invokeHarness( - request: InvokeHarnessRequest, + request: HarnessInvokeRequest, options: CoreOptions, abortSignal?: AbortSignal, ): Promise; diff --git a/src/handlers/project/invoke/index.test.tsx b/src/handlers/project/invoke/index.test.tsx index 2daa77dcf..e46143a33 100644 --- a/src/handlers/project/invoke/index.test.tsx +++ b/src/handlers/project/invoke/index.test.tsx @@ -24,6 +24,7 @@ import { waitFor, } from "../../../testing"; import * as tui from "../../../tui"; +import type { HarnessInvokeRequest } from "../../harness/types"; import { createRootHandler } from "../../index"; import { RegionKey } from "../../keys"; @@ -112,7 +113,7 @@ function backend() { async function routedCommand( args: readonly string[], resources: Resources | undefined, - options: { writeTargets?: boolean; isTTY?: boolean } = {}, + options: { writeTargets?: boolean; isTTY?: boolean; stdin?: string } = {}, ) { if (resources) { await inProject(resources, options); @@ -149,7 +150,7 @@ async function routedCommand( yield Buffer.from("gateway response"); })(), }); - const io = testIO({ isTTY: options.isTTY }); + const io = testIO({ isTTY: options.isTTY, stdin: options.stdin }); const root = createRootHandler(core, { io: io.io, logger: createSilentLogger(), @@ -220,6 +221,46 @@ describe("invoke", () => { }, ); + test("invokes a harness ARN with a bearer token and no GetHarness lookup", async () => { + const harnessArn = arn(`harness/${HARNESS_ID}`); + const subject = await routedCommand( + ["--harness", harnessArn, "--bearer-token", "-", "--prompt", "hi"], + undefined, + { stdin: "jwt-123\n" }, + ); + await subject.route(); + + expect(subject.core.harness.calls.map((call) => call.method)).toEqual(["invokeHarness"]); + const [request, options] = subject.core.harness.calls[0]!.args as [ + HarnessInvokeRequest, + { region: string }, + ]; + expect(request.harnessArn).toBe(harnessArn); + expect(request.bearerToken).toBe("jwt-123"); + expect(options.region).toBe(TARGET.region); + }); + + test("invokes a project harness by name with a bearer token, looked up with target credentials", async () => { + const subject = await routedCommand( + ["--harness", "support", "--bearer-token", "-", "--prompt", "hi"], + { harnesses: [HARNESS] }, + { stdin: "jwt-123\n" }, + ); + await subject.route(); + + expect(subject.core.harness.calls.map((call) => call.method)).toEqual([ + "getHarness", + "invokeHarness", + ]); + const [request, options] = subject.core.harness.calls[1]!.args as [ + HarnessInvokeRequest, + { credentials: unknown }, + ]; + expect(request.harnessArn).toBe(arn(`harness/${HARNESS_ID}`)); + expect(request.bearerToken).toBe("jwt-123"); + expect(options.credentials).toBe(TARGET_CREDENTIALS); + }); + test.each([ [["--runtime", "checkout"], `/agentcore/runtime/invoke/${RUNTIME_ID}`], [ diff --git a/src/handlers/project/invoke/index.tsx b/src/handlers/project/invoke/index.tsx index 7cdd7b9a2..4686a8301 100644 --- a/src/handlers/project/invoke/index.tsx +++ b/src/handlers/project/invoke/index.tsx @@ -1,4 +1,5 @@ import z from "zod"; +import { parseArn } from "../../../core/arn"; import { InputValidationError } from "../../../errors"; import type { AppIO } from "../../../io"; import { withProject } from "../../../middleware"; @@ -15,6 +16,7 @@ import { attributeName, parseFlags } from "../../../router/flags"; import { renderTuiAt } from "../../../tui"; import { createInvokeGatewayHandler, invokeGatewayFlags } from "../../gateway/invoke"; import { createInvokeHarnessHandler, invokeHarnessFlags } from "../../harness/invoke"; +import { HarnessInvokeArnKey } from "../../harness/invoke/launchContext"; import { AwsCredentialProviderKey, JsonKey, RegionKey } from "../../keys"; import type { Core } from "../../types"; import { assertMutuallyExclusiveFlags, resolveResource } from "../../utils"; @@ -104,6 +106,9 @@ type ResourceSelection = { resourceType: ProjectInvokableResource; identifier: s const isSet = (value: unknown) => value !== undefined && value !== false; +const isHarnessArn = (identifier: string) => + parseArn(identifier)?.resource.startsWith("harness/") === true; + function flaggedResource(flags: InvokeFlags, headless: boolean): ResourceSelection | undefined { assertMutuallyExclusiveFlags(flags, RESOURCE_TYPES, { exactlyOne: headless }); const resourceType = RESOURCE_TYPES.find((type) => flags[type] !== undefined); @@ -182,6 +187,9 @@ export function createProjectInvokeHandler(core: Core, io: AppIO) { if (resource.credentials) { handlerCtx = handlerCtx.withValue(AwsCredentialProviderKey, resource.credentials); } + if (resourceType === "harness" && !isProjectName && isHarnessArn(identifier)) { + handlerCtx = handlerCtx.withValue(HarnessInvokeArnKey, identifier); + } await handler.handle(handlerCtx, parsedHandlerFlags, {}); }, }); diff --git a/src/testing/fixtures.tsx b/src/testing/fixtures.tsx index 5b166ae89..5dcbef78d 100644 --- a/src/testing/fixtures.tsx +++ b/src/testing/fixtures.tsx @@ -261,7 +261,7 @@ export function fixtureFactories(dir: string): { send: makeRecordingSend(real, dir), } as unknown as BedrockAgentCoreControlClient; }, - createDataClient: (config: ClientConfig) => { + createDataClient: (config) => { const real = createDataClient(config); return { send: makeRecordingSend(real, dir),