feat(bedrock-agentcore): add grantInvokeWithWebSocketStream method to Runtime

packages/@aws-cdk/aws-bedrock-agentcore-alpha/lib/runtime/perms.ts

@@ -12,6 +12,12 @@ export const RUNTIME_INVOKE_PERMS = ['bedrock-agentcore:InvokeAgentRuntime'];
  */
 export const RUNTIME_INVOKE_USER_PERMS = ['bedrock-agentcore:InvokeAgentRuntimeForUser'];
 
+/**
+ * Permissions to invoke the agent runtime via WebSocket stream
+ * Required when signing WebSocket connection requests to the agent runtime
+ */
+export const RUNTIME_INVOKE_WEBSOCKET_STREAM_PERMS = ['bedrock-agentcore:InvokeAgentRuntimeWithWebSocketStream'];
+
 /******************************************************************************
  *                         Control Plane Permissions
  *****************************************************************************/

packages/@aws-cdk/aws-bedrock-agentcore-alpha/lib/runtime/runtime-base.ts

@@ -28,7 +28,7 @@ import * as iam from 'aws-cdk-lib/aws-iam';
 import { ValidationError } from 'aws-cdk-lib/core/lib/errors';
 import { lit } from 'aws-cdk-lib/core/lib/helpers-internal';
 import type { Construct } from 'constructs';
-import { RUNTIME_INVOKE_PERMS, RUNTIME_INVOKE_USER_PERMS } from './perms';
+import { RUNTIME_INVOKE_PERMS, RUNTIME_INVOKE_USER_PERMS, RUNTIME_INVOKE_WEBSOCKET_STREAM_PERMS } from './perms';
 
 /******************************************************************************
  *                                Interface
@@ -185,6 +185,13 @@ export interface IBedrockAgentRuntime extends IResource, iam.IGrantable, ec2.ICo
    * @param grantee The principal to grant access to
    */
   grantInvoke(grantee: iam.IGrantable): iam.Grant;
+
+  /**
+   * Permits an IAM principal to invoke this runtime via WebSocket stream
+   * Grants the bedrock-agentcore:InvokeAgentRuntimeWithWebSocketStream permission
+   * @param grantee The principal to grant access to
+   */
+  grantInvokeWithWebSocketStream(grantee: iam.IGrantable): iam.Grant;
 }
 
 /******************************************************************************
@@ -315,6 +322,22 @@ export abstract class RuntimeBase extends Resource implements IBedrockAgentRunti
     });
   }
 
+  /**
+   * Permits an IAM principal to invoke this runtime via WebSocket stream
+   * Grants the bedrock-agentcore:InvokeAgentRuntimeWithWebSocketStream permission
+   *
+   * [disable-awslint:no-grants]
+   *
+   * @param grantee The principal to grant access to
+   */
+  public grantInvokeWithWebSocketStream(grantee: iam.IGrantable): iam.Grant {
+    return iam.Grant.addToPrincipal({
+      grantee,
+      actions: RUNTIME_INVOKE_WEBSOCKET_STREAM_PERMS,
+      resourceArns: [this.runtimeRef.agentRuntimeArn, `${this.runtimeRef.agentRuntimeArn}/*`],
+    });
+  }
+
   // ------------------------------------------------------
   // Metrics
   // ------------------------------------------------------

packages/@aws-cdk/aws-bedrock-agentcore-alpha/test/agentcore/runtime/runtime.test.ts

@@ -2088,6 +2088,82 @@ describe('Runtime grantInvokeRuntime permission tests', () => {
       ],
     });
   });
+
+  test('Should grant InvokeAgentRuntimeWithWebSocketStream permission with grantInvokeWithWebSocketStream', () => {
+    runtime.grantInvokeWithWebSocketStream(granteeRole);
+
+    const template = Template.fromStack(stack);
+
+    template.hasResourceProperties('AWS::IAM::Policy', {
+      PolicyDocument: {
+        Statement: Match.arrayWith([
+          Match.objectLike({
+            Action: 'bedrock-agentcore:InvokeAgentRuntimeWithWebSocketStream',
+            Effect: 'Allow',
+            Resource: [
+              {
+                'Fn::GetAtt': [
+                  Match.stringLikeRegexp('testruntime.*'),
+                  'AgentRuntimeArn',
+                ],
+              },
+              {
+                'Fn::Join': [
+                  '',
+                  [
+                    {
+                      'Fn::GetAtt': [
+                        Match.stringLikeRegexp('testruntime.*'),
+                        'AgentRuntimeArn',
+                      ],
+                    },
+                    '/*',
+                  ],
+                ],
+              },
+            ],
+          }),
+        ]),
+      },
+    });
+  });
+
+  test('Should grant WebSocket stream permission on imported runtime', () => {
+    const importedRuntime = Runtime.fromAgentRuntimeAttributes(stack, 'ImportedRuntimeWs', {
+      agentRuntimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/test-runtime-id',
+      agentRuntimeId: 'test-runtime-id',
+      agentRuntimeName: 'test-runtime',
+      roleArn: 'arn:aws:iam::123456789012:role/test-role',
+      agentRuntimeVersion: '1',
+    });
+
+    importedRuntime.grantInvokeWithWebSocketStream(granteeRole);
+
+    const template = Template.fromStack(stack);
+
+    template.hasResourceProperties('AWS::IAM::Policy', {
+      PolicyDocument: {
+        Statement: Match.arrayWith([
+          Match.objectLike({
+            Action: 'bedrock-agentcore:InvokeAgentRuntimeWithWebSocketStream',
+            Effect: 'Allow',
+            Resource: [
+              'arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/test-runtime-id',
+              'arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/test-runtime-id/*',
+            ],
+          }),
+        ]),
+      },
+    });
+  });
+
+  test('Should return Grant object with success for grantInvokeWithWebSocketStream', () => {
+    const grant = runtime.grantInvokeWithWebSocketStream(granteeRole);
+
+    expect(grant).toBeDefined();
+    expect(grant.success).toBe(true);
+    expect(grant.principalStatement).toBeDefined();
+  });
 });
 
 describe('Runtime role validation tests', () => {