@aws-blocks/bb-agent 0.3.5 → 0.4.0

package/src/index.cdk.test.ts

@@ -2,74 +2,166 @@
 // SPDX-License-Identifier: Apache-2.0
 
 /**
- * CDK-side regression tests for the Agent's internal AsyncJob event source.
+ * CDK-side tests for the Agent BB.
  *
- * `stream()` submits one job per interactive turn (and a second on HITL resume),
- * so the caller is blocked on that job starting. AsyncJob's defaults
- * (batchSize 10 / maxBatchingWindowSeconds 5) would add up to 5s of SQS
- * batching delay to that human-blocking path, so the Agent opts out at both
- * construction sites. These tests pin the opt-out to the synthesized template:
- * if the defaults are ever inherited again, the latency regression fails here
- * instead of surfacing as a slow agent in production.
- *
- * Must run under `--conditions=cdk`; otherwise the internal BBs resolve to
- * their mock implementations and no CloudFormation resources are produced.
+ * Pin that the Agent provisions an AgentCore Runtime for the streaming loop that runs AS the shared
+ * Blocks execution role (the same principal as the handler) — not a bespoke per-runtime role. That
+ * shared role already carries everything the loop touches (Realtime publish via the handler wiring,
+ * Bedrock, the conversation/message tables, the session bucket); the Agent adds only the
+ * AgentCore-specific bits — the scoped `bedrock-agentcore` assume-role trust, Bedrock model access,
+ * and the shared handler's permission to INVOKE the runtime — and injects the config location so the
+ * container loads the same app config as the handler.
  */
-import { test, afterEach } from 'node:test';
+import { test, before, after } from 'node:test';
 import assert from 'node:assert';
+import { mkdtempSync, rmSync, writeFileSync } from 'node:fs';
+import { dirname, join } from 'node:path';
+import { fileURLToPath } from 'node:url';
 import * as cdk from 'aws-cdk-lib';
-import type { Construct } from 'constructs';
 import { Template } from 'aws-cdk-lib/assertions';
-import { Scope, DEFAULT_NODE_RUNTIME } from '@aws-blocks/core/cdk';
+import { BlocksStack, BlocksPresets } from '@aws-blocks/core/cdk';
+import type { DefaultComputeFactory } from '@aws-blocks/core/cdk/internal';
+import { LambdaCompute } from '@aws-blocks/bb-lambda-compute/cdk';
 import { Agent } from './index.cdk.js';
 
-class StubBlocksStack extends cdk.Stack {
-	public readonly handler: cdk.aws_lambda.Function;
-	public readonly executionRole: cdk.aws_iam.IRole;
-	public readonly id: string;
-	constructor(scope: Construct, id: string) {
-		super(scope, id);
-		this.id = id;
-		(globalThis as any).CURRENT_BLOCKS_STACK = this;
-		this.executionRole = new cdk.aws_iam.Role(this, 'BlocksRole', {
-			assumedBy: new cdk.aws_iam.ServicePrincipal('lambda.amazonaws.com'),
-		});
-		this.handler = new cdk.aws_lambda.Function(this, 'StubHandler', {
-			runtime: DEFAULT_NODE_RUNTIME,
-			handler: 'index.handler',
-			code: cdk.aws_lambda.Code.fromInline('exports.handler = async () => {};'),
-			role: this.executionRole,
-		});
-	}
-}
+/**
+ * Inject LambdaCompute as the stack's default compute, the same way
+ * `@aws-blocks/blocks` does for real apps. `main` (PR #459) made the Agent
+ * resolve the stack's `defaultCompute` at construction (the Realtime/AgentCore
+ * wiring reads it), so a bare handler-only stub stack (no default compute) can
+ * no longer synthesize it — `BlocksStack.create` initializes the default
+ * compute the getter resolves to. (`root as never` is core's existing plumbing
+ * pattern for the factory signature.)
+ */
+const lambdaFactory: DefaultComputeFactory = (root) => new LambdaCompute(root as never, 'DefaultCompute');
 
-// synthAgent() installs its own stack as the ambient CURRENT_BLOCKS_STACK; clear it
-// afterwards so no test observes a stack left behind by the previous one (node
-// --test isolates files by default, so this is hygiene against future sharing).
-afterEach(() => {
-	delete (globalThis as any).CURRENT_BLOCKS_STACK;
+/** A real directory to hand fromCodeAsset (any existing dir works for a synth-only test). */
+const ASSET_DIR = dirname(fileURLToPath(import.meta.url));
+
+let tmpDir: string;
+let handlerPath: string;
+let backendPath: string;
+
+before(() => {
+	// Building Blocks resolve to their mock entry points unless `--conditions=cdk`
+	// is active; keep it set so `BlocksStack.create` produces real CloudFormation.
+	process.env.NODE_OPTIONS = `${process.env.NODE_OPTIONS ?? ''} --conditions=cdk`;
+	// `BlocksStack.create` needs a real backend handler + backend module on disk
+	// (it imports the backend module and points the handler compute at the file).
+	tmpDir = mkdtempSync(join(ASSET_DIR, 'tmp-agent-cdk-'));
+	handlerPath = join(tmpDir, 'handler.mjs');
+	writeFileSync(handlerPath, "export const handler = async () => ({ statusCode: 200, body: '{}' });\n");
+	backendPath = join(tmpDir, 'backend.mjs');
+	writeFileSync(backendPath, 'export default () => {};\n');
 });
 
-function synthAgent(): any {
-	const app = new cdk.App();
-	const stack = new StubBlocksStack(app, 'teststack');
-	const parent = new Scope('app');
-	new Agent(parent, 'agent', { inferenceOnly: true });
+after(() => {
+	rmSync(tmpDir, { recursive: true, force: true });
+});
 
-	const mappings = Template.fromStack(stack).findResources('AWS::Lambda::EventSourceMapping');
-	const keys = Object.keys(mappings);
-	assert.strictEqual(keys.length, 1, 'exactly one event source mapping expected for the agent job');
-	return mappings[keys[0]].Properties;
+async function synth(): Promise<Template> {
+	const app = new cdk.App();
+	// Real BlocksStack (with the shared executionRole/BlocksRole + a LambdaCompute default compute)
+	// so the Agent can resolve its default compute. `id` === 'teststack' so BLOCKS_STACK_NAME (derived
+	// from the owning root id) resolves to 'teststack', which the injection assertion below matches.
+	const stack = await BlocksStack.create(app, 'teststack', {
+		backendHandlerPath: handlerPath,
+		backendCDKPath: backendPath,
+		defaults: BlocksPresets.production,
+		defaultComputeFactory: lambdaFactory,
+	});
+	// agentcoreAssetPath bypasses the synth-time co-bundle (which needs a BlocksStack backend path).
+	new Agent(stack, 'agent', { systemPrompt: 'You are a test agent.', agentcoreAssetPath: ASSET_DIR });
+	return Template.fromStack(stack);
 }
 
-test('CDK: the Agent job takes one message per invocation (no batching on the interactive path)', () => {
-	assert.strictEqual(synthAgent().BatchSize, 1);
+test('CDK: Agent provisions an AgentCore Runtime for the loop', async () => {
+	const template = await synth();
+	assert.ok(
+		Object.keys(template.findResources('AWS::BedrockAgentCore::Runtime')).length >= 1,
+		'expected an AWS::BedrockAgentCore::Runtime resource',
+	);
+});
+
+test('CDK: the loop runs as the shared execution role, which carries everything it touches', async () => {
+	const template = await synth();
+	const json = JSON.stringify(template.toJSON());
+	// Realtime publish (both halves) — load-bearing for streaming from the container. Granted to the
+	// shared role by the Realtime BB's handler wiring (not by the Agent), and inherited because the
+	// loop runs AS that role.
+	assert.ok(json.includes('execute-api:ManageConnections'), 'Realtime postToConnection');
+	assert.ok(json.includes('dynamodb:Query'), 'connections-table + history query');
+	// Model + storage the loop uses (Bedrock granted here; S3/DynamoDB via the child BBs).
+	assert.ok(json.includes('bedrock:InvokeModel'), 'Bedrock invoke');
+	assert.ok(json.includes('s3:GetObject'), 'session-bucket access (via FileBucket)');
+	// The RPC handler (also the shared role) can start the loop.
+	assert.ok(json.includes('bedrock-agentcore:InvokeAgentRuntime'), 'InvokeAgentRuntime');
+	// The runtime runs AS the shared BlocksRole — not a bespoke per-runtime role. Its RoleArn
+	// should reference BlocksRole, and there should be no `RuntimeRole` construct anywhere.
+	const runtime = Object.values(template.findResources('AWS::BedrockAgentCore::Runtime'))[0] as { Properties?: Record<string, unknown> };
+	assert.ok(JSON.stringify(runtime.Properties ?? {}).includes('BlocksRole'), 'runtime executionRole should be the shared BlocksRole');
+	assert.ok(!json.includes('RuntimeRole'), 'no bespoke per-runtime role should be created');
+	// The container is given the config location so loadConfigToProcessEnv() loads the same full app
+	// config as the handler (delivers BLOCKS_RT_CALLBACK_URL + any config-backed BB values a tool needs).
+	const runtimeJson = JSON.stringify(runtime.Properties ?? {});
+	assert.ok(runtimeJson.includes('BLOCKS_CONFIG_BUCKET'), 'runtime must be injected BLOCKS_CONFIG_BUCKET');
+	assert.ok(runtimeJson.includes('BLOCKS_CONFIG_KEY'), 'runtime must be injected BLOCKS_CONFIG_KEY');
+	// BLOCKS_STACK_NAME must be the owning root id (`backendStackName`) — the SAME value the handler
+	// uses to derive resource names — so the container's namespace matches what CDK provisioned. Here
+	// the owning BlocksStack's id is 'teststack', so backendStackName resolves to 'teststack'.
+	assert.ok(runtimeJson.includes('BLOCKS_STACK_NAME'), 'runtime must be injected BLOCKS_STACK_NAME');
+	assert.ok(runtimeJson.includes('teststack'), 'BLOCKS_STACK_NAME resolves to the owning root id');
 });
 
-test('CDK: the Agent job has no SQS batching window (no added latency for the caller)', () => {
-	assert.strictEqual(synthAgent().MaximumBatchingWindowInSeconds, 0);
+test('CDK: the Agent adds the bedrock-agentcore assume-role trust to the shared role (scoped)', async () => {
+	// Core is BB-agnostic — it does not trust bedrock-agentcore. The Agent adds that trust here, so
+	// the AgentCore Runtime can assume the shared role it runs AS. It must be scoped by
+	// aws:SourceAccount (AWS's recommended AgentCore trust policy), and lambda trust must remain.
+	const template = await synth();
+	const roles = template.findResources('AWS::IAM::Role');
+	const blocksRoleId = Object.keys(roles).find(k => k.includes('BlocksRole'));
+	assert.ok(blocksRoleId, 'expected the shared BlocksRole');
+
+	const statements = roles[blocksRoleId].Properties.AssumeRolePolicyDocument.Statement as Array<{
+		Principal?: { Service?: string | string[] };
+		Condition?: Record<string, Record<string, unknown>>;
+	}>;
+	const servicesOf = (s: (typeof statements)[number]) => {
+		const svc = s.Principal?.Service;
+		return Array.isArray(svc) ? svc : svc ? [svc] : [];
+	};
+	const services = statements.flatMap(servicesOf);
+	assert.ok(services.includes('lambda.amazonaws.com'), 'shared role must stay Lambda-assumable');
+	assert.ok(services.includes('bedrock-agentcore.amazonaws.com'), 'shared role must be assumable by the AgentCore Runtime');
+
+	const agentCoreStmt = statements.find(s => servicesOf(s).includes('bedrock-agentcore.amazonaws.com'));
+	assert.ok(agentCoreStmt?.Condition, 'AgentCore trust statement must carry a scoping Condition');
+	assert.ok(
+		JSON.stringify(agentCoreStmt.Condition).includes('aws:SourceAccount'),
+		'AgentCore trust must be scoped by aws:SourceAccount',
+	);
 });
 
-test('CDK: the Agent job still reports partial batch failures', () => {
-	assert.deepStrictEqual(synthAgent().FunctionResponseTypes, ['ReportBatchItemFailures']);
+test('CDK: multiple agents add the identical shared-role grants ONCE, not per-agent', async () => {
+	// Every agent would otherwise add the SAME bedrock-agentcore trust / Bedrock / InvokeAgentRuntime
+	// statements to the one shared role. They're identical (stack-scoped), so they must be added
+	// exactly once regardless of agent count — keeping the shared role's policy from bloating with
+	// duplicates. (Overflow into managed policies is a normal CDK mechanism and not asserted against.)
+	const app = new cdk.App();
+	const stack = await BlocksStack.create(app, 'teststack', {
+		backendHandlerPath: handlerPath,
+		backendCDKPath: backendPath,
+		defaults: BlocksPresets.production,
+		defaultComputeFactory: lambdaFactory,
+	});
+	for (const id of ['agent1', 'agent2', 'agent3']) {
+		new Agent(stack, id, { systemPrompt: 'You are a test agent.', agentcoreAssetPath: ASSET_DIR });
+	}
+	const template = Template.fromStack(stack);
+	const json = JSON.stringify(template.toJSON());
+
+	// Three runtimes (per-agent), but the shared-role grant statements appear once (deduped).
+	assert.strictEqual(Object.keys(template.findResources('AWS::BedrockAgentCore::Runtime')).length, 3, 'three runtimes');
+	const invokeGrants = json.split('bedrock-agentcore:InvokeAgentRuntime').length - 1;
+	assert.strictEqual(invokeGrants, 1, `InvokeAgentRuntime should be granted once, found ${invokeGrants}`);
 });

package/src/index.cdk.ts

@@ -1,51 +1,44 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { PolicyStatement } from 'aws-cdk-lib/aws-iam';
 import { Scope } from '@aws-blocks/core/cdk';
 import type { ScopeParent } from '@aws-blocks/core';
 import { DistributedTable } from '@aws-blocks/bb-distributed-table';
 import { Realtime } from '@aws-blocks/bb-realtime';
-import { AsyncJob } from '@aws-blocks/bb-async-job';
 import { FileBucket } from '@aws-blocks/bb-file-bucket';
+import { AgentCoreRuntime } from './agentcore-runtime.cdk.js';
 import { messageSchema, conversationSchema, agentStreamChunkSchema } from './schemas.js';
-import { INTERACTIVE_JOB_EVENT_SOURCE } from './job-event-source.js';
-import { z } from 'zod';
+import type { AgentConfig } from './types.js';
 
 export { AgentErrors } from './errors.js';
 export { BedrockModels, OllamaModels } from './models.js';
 
-const jobPayloadSchema = z.object({
-	message: z.string(),
-	conversationId: z.string().optional(),
-});
-
 export class Agent extends Scope {
 	/**
 	 * CDK layer for the Agent BB.
-	 * Mirrors the runtime's BB creation so CDK discovers and provisions all resources.
 	 *
-	 * TODO: scope Bedrock IAM grant to specific modelId from config
-	 * TODO: guardrails CDK provisioning
+	 * Provisions the session FileBucket, the conversation + message DistributedTables, the
+	 * Realtime BB used to stream chunks to the browser, and the AgentCore Runtime that hosts
+	 * the streaming agent loop. All AgentCore-specific provisioning (co-bundle, runtime role
+	 * and its grants, container env, and the handler's invoke permission) lives in the
+	 * self-contained {@link AgentCoreRuntime} so it can later fold into a per-BB compute abstraction.
+	 *
+	 * The loop runs inside the AgentCore Runtime (not the shared handler Lambda), so the shared
+	 * handler no longer needs Bedrock access — the runtime's own role gets it (see AgentCoreRuntime).
 	 */
-	constructor(scope: ScopeParent, id: string, config?: any) {
+	constructor(scope: ScopeParent, id: string, config?: AgentConfig) {
 		super(id, { parent: scope });
 
-		this.executionRole.addToPrincipalPolicy(new PolicyStatement({
-			actions: ['bedrock:InvokeModel', 'bedrock:InvokeModelWithResponseStream', 'bedrock:GetFoundationModel', 'bedrock:ListFoundationModels', 'bedrock:GetInferenceProfile'],
-			resources: [
-				'arn:aws:bedrock:*::foundation-model/*',
-				'arn:aws:bedrock:*:*:inference-profile/*',
-			],
-		}));
-
-		// Propagate `removalPolicy` to the sessions bucket so customers can
-		// opt sandbox stacks into clean teardown. Without it, CDK's RETAIN
-		// default applies (production-safe) and `cdk destroy` will fail on
-		// a non-empty bucket — same pattern as FileBucket / KnowledgeBase.
-		// ID shortened to keep S3 bucket names within the 63-char limit
+		// Session-snapshot bucket. Provisioned here (and granted to the shared execution role that the
+		// AgentCore Runtime runs as); the deployed loop re-derives its name from this bucket's `fullId`
+		// in-process — the same `'sn'` id → same fullId → same physical bucket — so no name needs to be
+		// injected into the container. Propagate `removalPolicy` so customers can opt sandbox stacks into
+		// clean teardown (without it, CDK's RETAIN default applies). ID shortened to keep the S3 bucket
+		// name within the 63-char limit.
 		new FileBucket(this, 'sn', { removalPolicy: config?.removalPolicy });
 
+		// Conversation metadata + message history. These grant read/write to the shared execution
+		// role, which the AgentCore Runtime then runs as — so the loop can persist history.
 		if (!config?.inferenceOnly) {
 			new DistributedTable(this, 'convos', {
 				schema: conversationSchema,
@@ -59,14 +52,19 @@ export class Agent extends Scope {
 
 		new Realtime(this, 'rt', {
 			namespaces: {
-				chunks: { schema: agentStreamChunkSchema },
+				chunks: Realtime.namespace(agentStreamChunkSchema),
 			},
 		});
 
-		new AsyncJob(this, 'job', {
-			schema: jobPayloadSchema,
-			...INTERACTIVE_JOB_EVENT_SOURCE,
-			handler: async () => {},
+		// The agent loop runs on the AgentCore Runtime (as the shared Blocks execution role) and
+		// streams to the browser over Realtime. AgentCoreRuntime co-bundles the app backend,
+		// provisions the runtime, adds the bedrock-agentcore trust + Bedrock to the shared role, and
+		// injects the config location so the container loads the same app config as the handler
+		// (the session bucket/tables and Realtime publish are already granted to that role). Kept
+		// self-contained so it can later fold into a per-BB compute abstraction.
+		new AgentCoreRuntime(this, 'runtime', {
+			agentFullId: this.fullId,
+			agentcoreAssetPath: config?.agentcoreAssetPath,
 		});
 	}
 }

package/src/index.mock.ts

@@ -2,6 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 export { Agent } from './agent.mock.js';
+// Exported so api-extractor can resolve the (protected, @internal) dispatchTurn signature; the type
+// itself is @internal — not part of the public API (customers use stream()/resume()).
+export type { AgentTurnPayload } from './agent.js';
 export { AgentErrors, InterruptError } from './errors.js';
 export { BedrockModels, OllamaModels } from './models.js';
 export type { AgentConfig, AgentResult, AgentStreamChunk, AgentStreamResult, ToolDefinition, AgentTool, ToolFactory, ToolsConfig, ToolHandlerArgs, DefaultToolContext, InterruptResponse, ToolCallRecord, ModelConfig, StreamOptions, Message, Conversation, JSONValue, TokenUsage } from './types.js';