Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/core/project/backends/cdk.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -861,12 +861,14 @@ describe("CdkBackend.resolveDeployedResources", () => {
name: "checkout_agent",
id: "checkout_agent-AbCdEf1234",
target: TARGET,
credentialProvider: subject.credentials,
},
{
resourceType: "harness",
name: "support_agent",
id: "support_agent-AbCdEf1234",
target: TARGET,
credentialProvider: subject.credentials,
},
]);
expect(subject.stackReads).toHaveLength(1);
Expand Down Expand Up @@ -903,6 +905,7 @@ describe("CdkBackend.resolveDeployedResources", () => {
name: "support",
id: "support-AbCdEf1234",
target: TARGET,
credentialProvider: subject.credentials,
},
]);
expect(subject.stackReads[0]?.stackName).toBe("AgentCore-example-default");
Expand Down
2 changes: 1 addition & 1 deletion src/core/project/backends/cdk.ts
Original file line number Diff line number Diff line change
Expand Up @@ -431,7 +431,7 @@ export class CdkBackend implements ProjectBackend {
];
return resources.flatMap((resource) => {
const id = findDeployedResourceId(stack, resource);
return id ? [{ ...resource, id, target }] : [];
return id ? [{ ...resource, id, target, credentialProvider: credentials }] : [];
});
}

Expand Down
5 changes: 3 additions & 2 deletions src/core/project/manager.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -974,8 +974,9 @@ export class FsProjectManager implements ProjectManager {
({ resourceType, name }) => resourceType === input.resourceType && name === input.name,
);
// The declared target wins over the copy on the item: the manager resolved it
// from aws-targets.json, and both invoke handlers pin the AWS region off this
// value, so trusting a backend's echo would let it redirect the call.
// from aws-targets.json, and both invoke handlers pin the AWS region from this
// value while reusing the backend's verified credential provider. Trusting a
// backend's target echo would let it redirect the call.
if (resource) return { ...resource, target: resolved.target };

const label = input.resourceType === "runtime" ? "Runtime" : "Harness";
Expand Down
1 change: 1 addition & 0 deletions src/core/types.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import type { AwsCredentialIdentity, AwsCredentialIdentityProvider } from "@smit
// default credential chain leave it unset. Every v3 client accepts this same shape,
// so it comes from the shared Smithy types rather than any one client's config.
export type AwsCredentials = AwsCredentialIdentity | AwsCredentialIdentityProvider;
export type AwsCredentialProvider = AwsCredentialIdentityProvider;

// CoreOptions is the standard trailing argument for Core operations. It carries
// the per-call settings a handler resolves from context (the AWS region and an
Expand Down
6 changes: 5 additions & 1 deletion src/handlers/keys.tsx
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import z from "zod";
import { globalFlag } from "../router";
import type { AwsCredentialProvider } from "../core/types";
import { contextKey, globalFlag } from "../router";

// These keys are group-level flags declared on the root router. Because a
// GlobalFlag is also a typed ContextKey, handlers read its validated value back
Expand All @@ -18,3 +19,6 @@ export const EndpointKey = globalFlag(
"endpoint URL override",
z.string().optional(),
);

/** Explicit credential provider pinned by project target resolution. */
export const AwsCredentialProviderKey = contextKey<AwsCredentialProvider>("aws.credentialProvider");
6 changes: 4 additions & 2 deletions src/handlers/project/invoke/harness.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import { InputValidationError } from "../../../errors";
import type { AppIO } from "../../../io";
import { createHandler, flag, ProjectKey } from "../../../router";
import { JsonRendererKey, renderTuiAt } from "../../../tui";
import { JsonKey, RegionKey } from "../../keys";
import { AwsCredentialProviderKey, JsonKey, RegionKey } from "../../keys";
import { invokeHarnessTurn } from "../../harness/invoke/operation";
import type { Core } from "../../types";
import { coreOptsFromCtx } from "../../utils";
Expand Down Expand Up @@ -40,7 +40,9 @@ export const createProjectInvokeHarnessHandler = (
resourceType: "harness",
name,
});
const invokeCtx = ctx.withValue(RegionKey, deployed.target.region);
const invokeCtx = ctx
.withValue(RegionKey, deployed.target.region)
.withValue(AwsCredentialProviderKey, deployed.credentialProvider);

if (!flags.prompt) {
if (invokeCtx.require(JsonKey)) {
Expand Down
20 changes: 17 additions & 3 deletions src/handlers/project/invoke/index.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import {
inTempDirectory,
} from "../../../testing";
import { createRootHandler } from "../../index";
import { JsonKey, RegionKey } from "../../keys";
import { AwsCredentialProviderKey, JsonKey, RegionKey } from "../../keys";
import { RuntimeInvokeLaunchContextKey } from "../../runtime/invoke/launchContext";
import type { RuntimeInvokeRequest } from "../../runtime/types";
import type { Project } from "../types";
Expand All @@ -35,6 +35,10 @@ const TARGET = {
account: "111122223333",
region: "eu-west-1",
} as const;
const TARGET_CREDENTIALS = async () => ({
accessKeyId: "target-access-key",
secretAccessKey: "target-secret-key",
});
const RUNTIME_ID = "checkout-AbCdEf1234";
const RUNTIME_ARN = `arn:aws:bedrock-agentcore:${TARGET.region}:${TARGET.account}:runtime/${RUNTIME_ID}`;
const HARNESS_ID = "support-AbCdEf1234";
Expand Down Expand Up @@ -96,12 +100,14 @@ function backend() {
name,
id: RUNTIME_ID,
target: input.target,
credentialProvider: TARGET_CREDENTIALS,
})),
...project.spec.harnesses.map(({ name }) => ({
resourceType: "harness" as const,
name,
id: HARNESS_ID,
target: input.target,
credentialProvider: TARGET_CREDENTIALS,
})),
];
},
Expand Down Expand Up @@ -382,7 +388,10 @@ describe("project invoke", () => {
expect(new TextDecoder().decode(request.payload)).toBe(payload);
expect(request.contentType).toBe("application/custom+json");
expect(request.runtimeUserId).toBe("default");
expect(core.runtime.calls.at(-1)!.args[1]).toEqual({ region: TARGET.region });
expect(core.runtime.calls.at(-1)!.args[1]).toEqual({
region: TARGET.region,
credentials: TARGET_CREDENTIALS,
});
expect(io.stdout()).toBe("runtime response");
expect(resolved.calls).toEqual([{ target: TARGET }]);
});
Expand All @@ -399,7 +408,10 @@ describe("project invoke", () => {
qualifier: "DEFAULT",
messages: [{ role: "user", content: [{ text: "hello" }] }],
});
expect(core.harness.calls.at(-1)!.args[1]).toEqual({ region: TARGET.region });
expect(core.harness.calls.at(-1)!.args[1]).toEqual({
region: TARGET.region,
credentials: TARGET_CREDENTIALS,
});
expect(JSON.parse(io.stdout()).transcript).toContainEqual({
kind: "text",
text: "harness response",
Expand Down Expand Up @@ -453,6 +465,7 @@ describe("project invoke", () => {

expect(launches[0]!.path).toBe(`/agentcore/runtime/invoke/${RUNTIME_ID}`);
expect(launches[0]!.context.require(RegionKey)).toBe(TARGET.region);
expect(launches[0]!.context.require(AwsCredentialProviderKey)).toBe(TARGET_CREDENTIALS);
expect(launches[0]!.context.require(RuntimeInvokeLaunchContextKey)).toMatchObject({
runtimeId: RUNTIME_ID,
});
Expand Down Expand Up @@ -498,6 +511,7 @@ describe("project invoke", () => {

expect(launches[0]!.path).toBe(`/agentcore/harness/invoke/${HARNESS_ID}?qualifier=prod`);
expect(launches[0]!.context.require(RegionKey)).toBe(TARGET.region);
expect(launches[0]!.context.require(AwsCredentialProviderKey)).toBe(TARGET_CREDENTIALS);
});

test("bare project invoke opens the project resource picker", async () => {
Expand Down
49 changes: 44 additions & 5 deletions src/handlers/project/invoke/invoke.screen.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -55,10 +55,26 @@ function endpoint(name: string): AgentRuntimeEndpoint {
}

const TARGET = { name: "default", account: "111122223333", region: "eu-west-1" } as const;
const TARGET_CREDENTIALS = async () => ({
accessKeyId: "target-access-key",
secretAccessKey: "target-secret-key",
});

const DEPLOYED_RESOURCES: ResolvedDeployedResource[] = [
{ resourceType: "runtime", name: "checkout", id: "runtime-123", target: TARGET },
{ resourceType: "harness", name: "support", id: "harness-123", target: TARGET },
{
resourceType: "runtime",
name: "checkout",
id: "runtime-123",
target: TARGET,
credentialProvider: TARGET_CREDENTIALS,
},
{
resourceType: "harness",
name: "support",
id: "harness-123",
target: TARGET,
credentialProvider: TARGET_CREDENTIALS,
},
];

function core(resources: ResolvedDeployedResource[] = DEPLOYED_RESOURCES): TestCoreClient {
Expand All @@ -68,6 +84,7 @@ function core(resources: ResolvedDeployedResource[] = DEPLOYED_RESOURCES): TestC
name: input.name,
id: input.resourceType === "runtime" ? "runtime-123" : "harness-123",
target: TARGET,
credentialProvider: TARGET_CREDENTIALS,
});
value.projectManager.resolveDeployedResources = async () => ({ resources, target: TARGET });
value.runtime
Expand All @@ -88,7 +105,15 @@ function core(resources: ResolvedDeployedResource[] = DEPLOYED_RESOURCES): TestC
describe("project invoke picker", () => {
test("lists only resources present in the deployed target", async () => {
const screen = renderScreen("/agentcore/project/invoke", {
core: core([{ resourceType: "harness", name: "support", id: "harness-123", target: TARGET }]),
core: core([
{
resourceType: "harness",
name: "support",
id: "harness-123",
target: TARGET,
credentialProvider: TARGET_CREDENTIALS,
},
]),
withContext: (ctx) => ctx.withValue(ProjectKey, project),
});

Expand Down Expand Up @@ -164,8 +189,9 @@ describe("project invoke picker", () => {
});

test("opens the selected Harness chat in the same TUI", async () => {
const value = core();
const screen = renderScreen("/agentcore/project/invoke", {
core: core(),
core: value,
withContext: (ctx) => ctx.withValue(ProjectKey, project),
});

Expand All @@ -174,11 +200,17 @@ describe("project invoke picker", () => {
await screen.press("return");
await waitForText(screen.lastFrame, "send a message…");
expect(screen.lastFrame()).toContain("harness-123");
expect(value.harness.calls.find(({ method }) => method === "getHarness")?.args[1]).toEqual({
region: TARGET.region,
endpointUrl: undefined,
credentials: TARGET_CREDENTIALS,
});
});

test("uses the existing Runtime endpoint picker before its JSON console", async () => {
const value = core();
const screen = renderScreen("/agentcore/project/invoke", {
core: core(),
core: value,
withContext: (ctx) => ctx.withValue(ProjectKey, project),
});

Expand All @@ -188,5 +220,12 @@ describe("project invoke picker", () => {
await screen.press("return");
await waitForText(screen.lastFrame, "Enter JSON payload");
expect(screen.lastFrame()).not.toContain("Enter prompt");
expect(
value.runtime.calls.find(({ method }) => method === "listRuntimeEndpoints")?.args[3],
).toEqual({
region: TARGET.region,
endpointUrl: undefined,
credentials: TARGET_CREDENTIALS,
});
});
});
6 changes: 4 additions & 2 deletions src/handlers/project/invoke/runtime.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ import type { AppIO } from "../../../io";
import { ExitCode, withUserCancellation } from "../../../runnable";
import { createHandler, flag, ProjectKey } from "../../../router";
import { renderTuiAt } from "../../../tui";
import { JsonKey, RegionKey } from "../../keys";
import { AwsCredentialProviderKey, JsonKey, RegionKey } from "../../keys";
import { RuntimeInvokeLaunchContextKey } from "../../runtime/invoke/launchContext";
import { invokeRuntimeTarget } from "../../runtime/invoke/operation";
import {
Expand Down Expand Up @@ -136,7 +136,9 @@ export const createProjectInvokeRuntimeHandler = (
resourceType: "runtime",
name,
});
const invokeCtx = ctx.withValue(RegionKey, deployed.target.region);
const invokeCtx = ctx
.withValue(RegionKey, deployed.target.region)
.withValue(AwsCredentialProviderKey, deployed.credentialProvider);

if (flags.payload === undefined) {
const hasHeadlessOnlyFlag = Object.entries(flags).some(
Expand Down
6 changes: 4 additions & 2 deletions src/handlers/project/invoke/screen.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import { Spinner } from "../../../components/ui/spinner";
import { glyphs } from "../../../components/ui/_core.js";
import { ProjectKey, type Context } from "../../../router";
import { HarnessChat } from "../../harness/invoke/screen";
import { RegionKey } from "../../keys";
import { AwsCredentialProviderKey, RegionKey } from "../../keys";
import { RuntimeInvokeConsole } from "../../runtime/invoke/screen";
import type { ScreenProps } from "../../types";
import type { Project, ResolvedDeployedResources } from "../types";
Expand Down Expand Up @@ -110,7 +110,9 @@ function ProjectInvokePicker({
setDestination({
resourceType: row.resourceType,
id: row.id,
ctx: ctx.withValue(RegionKey, deployed.target.region),
ctx: ctx
.withValue(RegionKey, deployed.target.region)
.withValue(AwsCredentialProviderKey, row.credentialProvider),
});
};

Expand Down
17 changes: 15 additions & 2 deletions src/handlers/project/log/runtime.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,10 @@ const PRODUCTION_TARGET = {
account: "111122223333",
region: "ap-southeast-2",
} as const;
const TARGET_CREDENTIALS = async () => ({
accessKeyId: "target-access-key",
secretAccessKey: "target-secret-key",
});
const RUNTIMES = [
{
name: "checkout",
Expand Down Expand Up @@ -78,6 +82,7 @@ function backend(options: { deployed?: boolean } = {}) {
name,
id: `${name}-AbCdEf1234`,
target: input.target,
credentialProvider: TARGET_CREDENTIALS,
}));
},
async resolveProjectResources() {
Expand Down Expand Up @@ -131,7 +136,11 @@ describe("project log runtime", () => {
logGroupName: "/aws/bedrock-agentcore/runtimes/checkout-AbCdEf1234-DEFAULT",
});
expect(call.args[1]).toEqual({ filterPattern: undefined });
expect(call.args[2]).toEqual({ region: DEFAULT_TARGET.region, endpointUrl: undefined });
expect(call.args[2]).toEqual({
region: DEFAULT_TARGET.region,
endpointUrl: undefined,
credentials: TARGET_CREDENTIALS,
});
expect(subject.io.stderr()).toContain(
"Streaming logs for Runtime 'checkout' on target 'default'... (Ctrl+C to stop)",
);
Expand Down Expand Up @@ -163,7 +172,11 @@ describe("project log runtime", () => {
logGroupName: "/aws/bedrock-agentcore/runtimes/inventory-AbCdEf1234-BLUE",
});
expect(call.args[1]).toMatchObject({ limit: 25 });
expect(call.args[2]).toEqual({ region: PRODUCTION_TARGET.region, endpointUrl: undefined });
expect(call.args[2]).toEqual({
region: PRODUCTION_TARGET.region,
endpointUrl: undefined,
credentials: TARGET_CREDENTIALS,
});
});

test("requires --name when the project declares several Runtimes", async () => {
Expand Down
1 change: 1 addition & 0 deletions src/handlers/project/log/runtime.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ export const createProjectRuntimeLogHandler = (core: Core, io: AppIO) =>
const options = {
...coreOptsFromCtx(ctx),
region: deployed.target.region,
credentials: deployed.credentialProvider,
};
const source = {
logGroupName: runtimeLogGroup(deployed.id, flags.qualifier ?? DEFAULT_ENDPOINT_QUALIFIER),
Expand Down
5 changes: 4 additions & 1 deletion src/handlers/project/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import type { AgentCoreGateway, AgentCoreGatewayTarget } from "../../projectSche
import type { PolicyEngineSchema, PolicySchema } from "../../projectSchemas/policy";
import type { AwsDeploymentTarget } from "../../projectSchemas/aws-targets";
import type { ProgressEvent } from "../../tui/progress";
import type { AwsCredentialProvider } from "../../core/types";

type CreateProjectInputBase = {
/** The name of the project; also the directory it is scaffolded into. */
Expand Down Expand Up @@ -213,6 +214,8 @@ export type ResolvedDeployedResource = {
name: string;
id: string;
target: AwsDeploymentTarget;
/** Credential provider used to resolve and access this target. */
credentialProvider: AwsCredentialProvider;
};

export type ResolvedDeployedResources = {
Expand Down Expand Up @@ -459,7 +462,7 @@ export interface ProjectManager {
/** Locate an existing AgentCore project. Returns undefined if no project can be found. */
resolve(input: ResolveProjectInput): Promise<Project | undefined>;

/** Resolve a logical project resource to its deployed physical ID and target. */
/** Resolve a logical project resource to its deployed physical ID, target, and credential provider. */
resolveDeployedResource(
project: Project,
input: ResolveDeployedResourceInput,
Expand Down
Loading
Loading