Skip to content
Open
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
7 changes: 5 additions & 2 deletions src/assets/templates/bedrock-agent-proxy-python/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,10 @@ and invoked through AgentCore without changing it.
baked in at import time and can be overridden with the `BEDROCK_AGENT_ID`,
`BEDROCK_AGENT_ALIAS_ID`, and `BEDROCK_AGENT_REGION` environment variables.
- `bedrock-agent-policy.json` — grants the runtime's execution role
`bedrock:InvokeAgent` on the imported agent's alias. It is wired in through
the runtime's `additionalPolicies` entry in `agentcore/agentcore.json`.
`bedrock:InvokeAgent` on the imported agent's alias.
{{#if usesExistingExecutionRole}}This project uses a caller-owned execution
role, so attach `bedrock-agent-policy.json` to that role before deploying.
AgentCore CDK does not modify existing roles.{{else}}It is wired in through
the runtime's `additionalPolicies` entry in `agentcore/agentcore.json`.{{/if}}

Invoke it with a JSON payload like `{"prompt": "hello"}`.
43 changes: 38 additions & 5 deletions src/assets/templates/bedrock-agent-proxy-python/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,10 @@
# this runtime are forwarded to the Bedrock Agent, and its reply is streamed
# back — edit or replace this file to take ownership of the behavior.

import asyncio
import hashlib
import os
import re
import uuid

import boto3
Expand All @@ -12,30 +15,60 @@
AGENT_ID = os.environ.get("BEDROCK_AGENT_ID", "{{agentId}}")
AGENT_ALIAS_ID = os.environ.get("BEDROCK_AGENT_ALIAS_ID", "{{agentAliasId}}")
AGENT_REGION = os.environ.get("BEDROCK_AGENT_REGION", "{{agentRegion}}")
BEDROCK_SESSION_ID_PATTERN = re.compile(r"^[0-9A-Za-z._:-]{2,100}$")
END_OF_COMPLETION = object()

app = BedrockAgentCoreApp()
client = boto3.client("bedrock-agent-runtime", region_name=AGENT_REGION)


def normalize_session_id(value):
"""Return a stable Bedrock-compatible session id."""
if isinstance(value, str) and BEDROCK_SESSION_ID_PATTERN.fullmatch(value):
return value
if isinstance(value, str) and value:
return hashlib.sha256(value.encode("utf-8")).hexdigest()
return uuid.uuid4().hex


def next_completion_event(completion):
"""Read one streaming event without leaking StopIteration into asyncio."""
try:
return next(completion)
except StopIteration:
return END_OF_COMPLETION


@app.entrypoint
async def invoke(payload, context):
"""Forward the prompt to the Bedrock Agent and stream its completion."""
if not isinstance(payload, dict):
yield "Invalid payload; expected a JSON object with a non-empty 'prompt' field."
return

prompt = payload.get("prompt", "")
if not isinstance(prompt, str) or not prompt:
yield "No query provided; include a 'prompt' field in the payload."
return

# Bedrock Agent sessions require ids of 2+ chars; reuse the runtime session
# so multi-turn conversations keep the agent's own memory of the exchange.
session_id = context.session_id or payload.get("sessionId") or uuid.uuid4().hex
# Preserve compatible ids and hash longer/unsupported AgentCore ids so
# multi-turn conversations retain a stable Bedrock Agent session.
session_id = normalize_session_id(
getattr(context, "session_id", None) or payload.get("sessionId")
)

response = client.invoke_agent(
response = await asyncio.to_thread(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hah it looks like the AI didn't actually import this feature properly. It just proxies to Bedrock Agents. 😆

We'll have to implement this properly. I think we already have the code. We just need to port it over.

client.invoke_agent,
agentId=AGENT_ID,
agentAliasId=AGENT_ALIAS_ID,
sessionId=session_id,
inputText=prompt,
)
for event in response["completion"]:
completion = iter(response["completion"])
while True:
event = await asyncio.to_thread(next_completion_event, completion)
if event is END_OF_COMPLETION:
break
chunk = event.get("chunk")
if chunk and "bytes" in chunk:
yield chunk["bytes"].decode("utf-8")
Expand Down
138 changes: 138 additions & 0 deletions src/core/project/bedrockAgent.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,138 @@
import { describe, expect, test } from "bun:test";
import type {
GetAgentAliasCommandOutput,
GetAgentCommandOutput,
} from "@aws-sdk/client-bedrock-agent";
import { InputValidationError, MalformedServiceResponseError } from "../../errors";
import { createDescribeBedrockAgent, type BedrockAgentControlClient } from "./bedrockAgent";

const agent = {
agentId: "A1B2C3D4E5",
agentName: "SupportAgent",
agentArn: "arn:aws:bedrock:us-east-1:111122223333:agent/A1B2C3D4E5",
agentVersion: "DRAFT",
agentStatus: "PREPARED",
idleSessionTTLInSeconds: 600,
agentResourceRoleArn: "arn:aws:iam::111122223333:role/BedrockAgentRole",
createdAt: new Date(0),
updatedAt: new Date(0),
foundationModel: "us.amazon.nova-lite-v1:0",
};

const agentAlias = {
agentId: agent.agentId,
agentAliasId: "TSTALIASID",
agentAliasName: "live",
agentAliasArn: `arn:aws:bedrock:us-east-1:111122223333:agent-alias/${agent.agentId}/TSTALIASID`,
routingConfiguration: [{ agentVersion: "1" }],
createdAt: new Date(0),
updatedAt: new Date(0),
agentAliasStatus: "PREPARED",
};

class TestBedrockAgentControlClient implements BedrockAgentControlClient {
readonly calls: string[] = [];
agentOutput: GetAgentCommandOutput = { agent } as GetAgentCommandOutput;
aliasOutput: GetAgentAliasCommandOutput = {
agentAlias,
} as GetAgentAliasCommandOutput;
agentError?: Error;
aliasError?: Error;

async getAgent(): Promise<GetAgentCommandOutput> {
this.calls.push("getAgent");
if (this.agentError) throw this.agentError;
return this.agentOutput;
}

async getAgentAlias(): Promise<GetAgentAliasCommandOutput> {
this.calls.push("getAgentAlias");
if (this.aliasError) throw this.aliasError;
return this.aliasOutput;
}
}

function resourceNotFound(): Error {
return Object.assign(new Error("not found"), { name: "ResourceNotFoundException" });
}

describe("describeBedrockAgent", () => {
test("returns metadata from the requested agent and alias", async () => {
const client = new TestBedrockAgentControlClient();
const describeAgent = createDescribeBedrockAgent(() => client);

await expect(
describeAgent({
region: "us-east-1",
agentId: agent.agentId,
agentAliasId: agentAlias.agentAliasId,
}),
).resolves.toEqual({
agentName: agent.agentName,
agentStatus: agent.agentStatus,
agentAliasArn: agentAlias.agentAliasArn,
agentAliasName: agentAlias.agentAliasName,
agentAliasStatus: agentAlias.agentAliasStatus,
foundationModel: agent.foundationModel,
description: undefined,
});
expect(client.calls).toEqual(["getAgent", "getAgentAlias"]);
});

test("rejects an incomplete agent response before requesting the alias", async () => {
const client = new TestBedrockAgentControlClient();
client.agentOutput = {} as GetAgentCommandOutput;
const describeAgent = createDescribeBedrockAgent(() => client);

await expect(
describeAgent({
region: "us-east-1",
agentId: agent.agentId,
agentAliasId: agentAlias.agentAliasId,
}),
).rejects.toBeInstanceOf(MalformedServiceResponseError);
expect(client.calls).toEqual(["getAgent"]);
});

test("rejects an alias response for a different agent", async () => {
const client = new TestBedrockAgentControlClient();
client.aliasOutput = {
agentAlias: { ...agentAlias, agentId: "OTHERAGENT" },
} as GetAgentAliasCommandOutput;
const describeAgent = createDescribeBedrockAgent(() => client);

await expect(
describeAgent({
region: "us-east-1",
agentId: agent.agentId,
agentAliasId: agentAlias.agentAliasId,
}),
).rejects.toBeInstanceOf(MalformedServiceResponseError);
});

test("maps agent and alias not-found errors independently", async () => {
const missingAgentClient = new TestBedrockAgentControlClient();
missingAgentClient.agentError = resourceNotFound();
const describeMissingAgent = createDescribeBedrockAgent(() => missingAgentClient);
await expect(
describeMissingAgent({
region: "us-east-1",
agentId: agent.agentId,
agentAliasId: agentAlias.agentAliasId,
}),
).rejects.toBeInstanceOf(InputValidationError);
expect(missingAgentClient.calls).toEqual(["getAgent"]);

const missingAliasClient = new TestBedrockAgentControlClient();
missingAliasClient.aliasError = resourceNotFound();
const describeMissingAlias = createDescribeBedrockAgent(() => missingAliasClient);
await expect(
describeMissingAlias({
region: "us-east-1",
agentId: agent.agentId,
agentAliasId: agentAlias.agentAliasId,
}),
).rejects.toBeInstanceOf(InputValidationError);
expect(missingAliasClient.calls).toEqual(["getAgent", "getAgentAlias"]);
});
});
138 changes: 93 additions & 45 deletions src/core/project/bedrockAgent.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,10 @@
import {
BedrockAgentClient,
GetAgentAliasCommand,
GetAgentCommand,
type GetAgentAliasCommandOutput,
type GetAgentCommandOutput,
} from "@aws-sdk/client-bedrock-agent";
import { InputValidationError, MalformedServiceResponseError } from "../../errors";

/**
Expand All @@ -8,8 +15,12 @@ export const BEDROCK_AGENT_IMPORT_REGIONS = [
"us-east-1",
"us-west-2",
"eu-west-1",
"eu-west-2",
"eu-west-3",
"eu-central-1",
"eu-central-2",
"ap-southeast-1",
"ap-southeast-2",
"ap-northeast-1",
"ap-south-1",
"ca-central-1",
Expand Down Expand Up @@ -40,6 +51,27 @@ export type DescribeBedrockAgent = (
input: DescribeBedrockAgentInput,
) => Promise<BedrockAgentMetadata>;

export interface BedrockAgentControlClient {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the upshot of this interface?

getAgent(agentId: string): Promise<GetAgentCommandOutput>;
getAgentAlias(agentId: string, agentAliasId: string): Promise<GetAgentAliasCommandOutput>;
}

class AwsBedrockAgentControlClient implements BedrockAgentControlClient {
private readonly client: BedrockAgentClient;

constructor(region: string) {
this.client = new BedrockAgentClient({ region });
}

getAgent(agentId: string): Promise<GetAgentCommandOutput> {
return this.client.send(new GetAgentCommand({ agentId }));
}

getAgentAlias(agentId: string, agentAliasId: string): Promise<GetAgentAliasCommandOutput> {
return this.client.send(new GetAgentAliasCommand({ agentId, agentAliasId }));
}
}

function isNamedError(error: unknown, name: string): boolean {
return error instanceof Error && error.name === name;
}
Expand All @@ -49,55 +81,71 @@ function isNamedError(error: unknown, name: string): boolean {
* both to fail fast on a nonexistent agent/alias and to capture the metadata
* the scaffolded proxy embeds.
*/
export const describeBedrockAgent: DescribeBedrockAgent = async (input) => {
const { BedrockAgentClient, GetAgentCommand, GetAgentAliasCommand } =
await import("@aws-sdk/client-bedrock-agent");
const client = new BedrockAgentClient({ region: input.region });

let agent;
try {
({ agent } = await client.send(new GetAgentCommand({ agentId: input.agentId })));
} catch (error) {
if (isNamedError(error, "ResourceNotFoundException")) {
throw new InputValidationError(
`no Bedrock Agent with id '${input.agentId}' exists in ${input.region}; ` +
`check --agent-id and --region`,
{ cause: error },
);
export function createDescribeBedrockAgent(
createClient: (region: string) => BedrockAgentControlClient = (region) =>
new AwsBedrockAgentControlClient(region),
): DescribeBedrockAgent {
return async (input) => {
const client = createClient(input.region);

let agent;
try {
({ agent } = await client.getAgent(input.agentId));
} catch (error) {
if (isNamedError(error, "ResourceNotFoundException")) {
throw new InputValidationError(
`no Bedrock Agent with id '${input.agentId}' exists in ${input.region}; ` +
`check --agent-id and --region`,
{ cause: error },
);
}
throw error;
}
throw error;
}

let agentAlias;
try {
({ agentAlias } = await client.send(
new GetAgentAliasCommand({ agentId: input.agentId, agentAliasId: input.agentAliasId }),
));
} catch (error) {
if (isNamedError(error, "ResourceNotFoundException")) {
throw new InputValidationError(
`Bedrock Agent '${input.agentId}' has no alias with id '${input.agentAliasId}' in ` +
`${input.region}; check --agent-alias-id`,
{ cause: error },
if (!agent?.agentId || agent.agentId !== input.agentId || !agent.agentName) {
throw new MalformedServiceResponseError(
`the Bedrock Agent service returned an incomplete description for agent '${input.agentId}'`,
);
}
throw error;
}

if (!agent?.agentName || !agentAlias?.agentAliasArn || !agentAlias.agentAliasName) {
throw new MalformedServiceResponseError(
`the Bedrock Agent service returned an incomplete description for agent ` +
`'${input.agentId}' / alias '${input.agentAliasId}'`,
);
}
let agentAlias;
try {
({ agentAlias } = await client.getAgentAlias(input.agentId, input.agentAliasId));
} catch (error) {
if (isNamedError(error, "ResourceNotFoundException")) {
throw new InputValidationError(
`Bedrock Agent '${input.agentId}' has no alias with id '${input.agentAliasId}' in ` +
`${input.region}; check --agent-alias-id`,
{ cause: error },
);
}
throw error;
}

return {
agentName: agent.agentName,
agentStatus: agent.agentStatus ?? "UNKNOWN",
agentAliasArn: agentAlias.agentAliasArn,
agentAliasName: agentAlias.agentAliasName,
agentAliasStatus: agentAlias.agentAliasStatus ?? "UNKNOWN",
foundationModel: agent.foundationModel,
description: agent.description,
if (
!agentAlias?.agentId ||
agentAlias.agentId !== input.agentId ||
!agentAlias.agentAliasId ||
agentAlias.agentAliasId !== input.agentAliasId ||
!agentAlias.agentAliasArn ||
!agentAlias.agentAliasName
) {
throw new MalformedServiceResponseError(
`the Bedrock Agent service returned an incomplete description for agent ` +
`'${input.agentId}' / alias '${input.agentAliasId}'`,
);
}

return {
agentName: agent.agentName,
agentStatus: agent.agentStatus ?? "UNKNOWN",
agentAliasArn: agentAlias.agentAliasArn,
agentAliasName: agentAlias.agentAliasName,
agentAliasStatus: agentAlias.agentAliasStatus ?? "UNKNOWN",
foundationModel: agent.foundationModel,
description: agent.description,
};
};
};
}

export const describeBedrockAgent = createDescribeBedrockAgent();
Loading
Loading