From d504493c98cf4c90b0c29144b1cfd5c07952d1a8 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Tue, 3 Sep 2024 12:42:05 -0700
Subject: [PATCH 01/18] fix(stepfunctions-tasks): fix bedrock input/output path
 in step-functions tasks

---
 ...efaultTestDeployAssert9C0D2DFC.assets.json |  2 +-
 ...sks-bedrock-invoke-model-integ.assets.json |  6 ++--
 ...s-bedrock-invoke-model-integ.template.json | 14 +++++++++-
 .../integ.invoke-model.js.snapshot/cdk.out    |  2 +-
 .../integ.invoke-model.js.snapshot/integ.json |  2 +-
 .../manifest.json                             |  4 +--
 .../integ.invoke-model.js.snapshot/tree.json  | 22 ++++++++++++++-
 .../test/bedrock/integ.invoke-model.ts        | 23 +++++++++++++--
 .../lib/bedrock/invoke-model.ts               | 28 +++++++++++++++----
 9 files changed, 86 insertions(+), 17 deletions(-)

diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json
index 62c1eb1585be8..81419156e1160 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/InvokeModelDefaultTestDeployAssert9C0D2DFC.assets.json
@@ -1,5 +1,5 @@
 {
-  "version": "36.0.0",
+  "version": "36.0.5",
   "files": {
     "21fbb51d7b23f6a6c262b46a9caee79d744a3ac019fd45422d988b96d44b2a22": {
       "source": {
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
index ea5387a06b82c..0541b64d40527 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
@@ -1,7 +1,7 @@
 {
-  "version": "36.0.0",
+  "version": "36.0.5",
   "files": {
-    "e9d946bbac52fb88d3fc7c9ea4f26da0e6a6965417c8ae9fb5464e5342269444": {
+    "b69306f91f17e8d5e4b1d4e770e821de6110faa1c6811a445ce6a256e197e329": {
       "source": {
         "path": "aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json",
         "packaging": "file"
@@ -9,7 +9,7 @@
       "destinations": {
         "current_account-current_region": {
           "bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}",
-          "objectKey": "e9d946bbac52fb88d3fc7c9ea4f26da0e6a6965417c8ae9fb5464e5342269444.json",
+          "objectKey": "b69306f91f17e8d5e4b1d4e770e821de6110faa1c6811a445ce6a256e197e329.json",
           "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}"
         }
       }
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
index a67f72f4ea158..2a171f897511b 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
@@ -103,7 +103,19 @@
        {
         "Ref": "AWS::Region"
        },
-       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"End\":true,\"Type\":\"Task\",\"InputPath\":\"$.names\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
+       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
+       {
+        "Ref": "AWS::Partition"
+       },
+       ":states:::bedrock:invokeModel\",\"Parameters\":{\"ModelId\":\"arn:",
+       {
+        "Ref": "AWS::Partition"
+       },
+       ":bedrock:",
+       {
+        "Ref": "AWS::Region"
+       },
+       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
        {
         "Ref": "AWS::Partition"
        },
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out
index 1f0068d32659a..bd5311dc372de 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/cdk.out
@@ -1 +1 @@
-{"version":"36.0.0"}
\ No newline at end of file
+{"version":"36.0.5"}
\ No newline at end of file
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json
index 36143be21ab57..6e638249caa25 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/integ.json
@@ -1,5 +1,5 @@
 {
-  "version": "36.0.0",
+  "version": "36.0.5",
   "testCases": {
     "InvokeModel/DefaultTest": {
       "stacks": [
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
index 677092f63b132..dd8726e66ddae 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
@@ -1,5 +1,5 @@
 {
-  "version": "36.0.0",
+  "version": "36.0.5",
   "artifacts": {
     "aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets": {
       "type": "cdk:asset-manifest",
@@ -18,7 +18,7 @@
         "validateOnSynth": false,
         "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-deploy-role-${AWS::AccountId}-${AWS::Region}",
         "cloudFormationExecutionRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-cfn-exec-role-${AWS::AccountId}-${AWS::Region}",
-        "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/e9d946bbac52fb88d3fc7c9ea4f26da0e6a6965417c8ae9fb5464e5342269444.json",
+        "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/b69306f91f17e8d5e4b1d4e770e821de6110faa1c6811a445ce6a256e197e329.json",
         "requiresBootstrapStackVersion": 6,
         "bootstrapStackVersionSsmParameter": "/cdk-bootstrap/hnb659fds/version",
         "additionalDependencies": [
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
index 17b23007671bd..3f39dab1113e9 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
@@ -32,6 +32,14 @@
               "version": "0.0.0"
             }
           },
+          "Prompt4": {
+            "id": "Prompt4",
+            "path": "aws-stepfunctions-tasks-bedrock-invoke-model-integ/Prompt4",
+            "constructInfo": {
+              "fqn": "aws-cdk-lib.aws_stepfunctions_tasks.BedrockInvokeModel",
+              "version": "0.0.0"
+            }
+          },
           "StateMachine": {
             "id": "StateMachine",
             "path": "aws-stepfunctions-tasks-bedrock-invoke-model-integ/StateMachine",
@@ -186,7 +194,19 @@
                           {
                             "Ref": "AWS::Region"
                           },
-                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"End\":true,\"Type\":\"Task\",\"InputPath\":\"$.names\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
+                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
+                          {
+                            "Ref": "AWS::Partition"
+                          },
+                          ":states:::bedrock:invokeModel\",\"Parameters\":{\"ModelId\":\"arn:",
+                          {
+                            "Ref": "AWS::Partition"
+                          },
+                          ":bedrock:",
+                          {
+                            "Ref": "AWS::Region"
+                          },
+                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
                           {
                             "Ref": "AWS::Partition"
                           },
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
index 2d98f804b4475..1611d9e0faa1e 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
@@ -52,13 +52,32 @@ const prompt2 = new BedrockInvokeModel(stack, 'Prompt2', {
   resultPath: '$',
 });
 
+/** Test for Bedrock Output Path */
 const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
   model,
-  inputPath: sfn.JsonPath.stringAt('$.names'),
+  body: sfn.TaskInput.fromObject(
+    {
+      inputText: sfn.JsonPath.format(
+        'Alphabetize this list of first names:\n{}',
+        sfn.JsonPath.stringAt('$.names'),
+      ),
+      textGenerationConfig: {
+        maxTokenCount: 100,
+        temperature: 1,
+      },
+    },
+  ),
   outputPath: sfn.JsonPath.stringAt('$.names'),
 });
 
-const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3);
+/** Test for Bedrock s3 URI Path */
+const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
+  model,
+  s3inputPath: sfn.JsonPath.stringAt('$.names'),
+  s3outputPath: sfn.JsonPath.stringAt('$.names'),
+});
+
+const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4);
 
 new sfn.StateMachine(stack, 'StateMachine', {
   definitionBody: sfn.DefinitionBody.fromChainable(chain),
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
index 3abc7338cab8c..19e31911072cc 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
@@ -108,6 +108,24 @@ export interface BedrockInvokeModelProps extends sfn.TaskStateBaseProps {
    */
   readonly output?: BedrockInvokeModelOutputProps;
 
+  /**
+   * The destination location where the API response is written.
+   *
+   * This field can be used to specify s3 URI in the form of token
+   *
+   * @default - The API response body is returned in the result.
+   */
+  readonly s3inputPath?: string;
+
+  /**
+   * The source location where the API response is written.
+   *
+   * This field can be used to specify s3 URI in the form of token
+   *
+   * @default - The API response body is returned in the result.
+   */
+  readonly s3outputPath?: string;
+
   /**
    * The guardrail is applied to the invocation
    *
@@ -147,7 +165,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
 
     const isBodySpecified = props.body !== undefined;
     //Either specific props.input with bucket name and object key or input s3 path
-    const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.inputPath !== undefined);
+    const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.s3inputPath !== undefined);
 
     if (isBodySpecified && isInputSpecified) {
       throw new Error('Either `body` or `input` must be specified, but not both.');
@@ -173,7 +191,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       }),
     ];
 
-    if (this.props.inputPath !== undefined) {
+    if (this.props.s3inputPath !== undefined) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:GetObject'],
@@ -204,7 +222,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       );
     }
 
-    if (this.props.outputPath !== undefined) {
+    if (this.props.s3outputPath !== undefined) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:PutObject'],
@@ -271,10 +289,10 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
         Body: this.props.body?.value,
         Input: this.props.input?.s3Location ? {
           S3Uri: `s3://${this.props.input.s3Location.bucketName}/${this.props.input.s3Location.objectKey}`,
-        } : this.props.inputPath ? { S3Uri: this.props.inputPath } : undefined,
+        } : this.props.s3inputPath ? { S3Uri: this.props.s3inputPath } : undefined,
         Output: this.props.output?.s3Location ? {
           S3Uri: `s3://${this.props.output.s3Location.bucketName}/${this.props.output.s3Location.objectKey}`,
-        } : this.props.outputPath ? { S3Uri: this.props.outputPath }: undefined,
+        } : this.props.s3outputPath ? { S3Uri: this.props.s3outputPath }: undefined,
         GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
         GuardrailVersion: this.props.guardrail?.guardrailVersion,
         Trace: this.props.traceEnabled === undefined

From 8f0d77fd23ea8cca1f5b8d1e510e4df8263c8e29 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Tue, 3 Sep 2024 13:27:42 -0700
Subject: [PATCH 02/18] fix(stepfunctions-tasks): fix bedrock input/output path
 in step-functions tasks

---
 .../test/bedrock/integ.invoke-model.ts             |  4 ++--
 .../lib/bedrock/invoke-model.ts                    | 14 +++++++-------
 .../test/bedrock/invoke-model.test.ts              | 10 ++++------
 3 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
index 1611d9e0faa1e..2f91d3b3b173e 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
@@ -73,8 +73,8 @@ const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
 /** Test for Bedrock s3 URI Path */
 const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
   model,
-  s3inputPath: sfn.JsonPath.stringAt('$.names'),
-  s3outputPath: sfn.JsonPath.stringAt('$.names'),
+  s3InputUri: sfn.JsonPath.stringAt('$.names'),
+  s3OutputUri: sfn.JsonPath.stringAt('$.names'),
 });
 
 const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4);
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
index 19e31911072cc..20d9692927b38 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
@@ -115,7 +115,7 @@ export interface BedrockInvokeModelProps extends sfn.TaskStateBaseProps {
    *
    * @default - The API response body is returned in the result.
    */
-  readonly s3inputPath?: string;
+  readonly s3InputUri?: string;
 
   /**
    * The source location where the API response is written.
@@ -124,7 +124,7 @@ export interface BedrockInvokeModelProps extends sfn.TaskStateBaseProps {
    *
    * @default - The API response body is returned in the result.
    */
-  readonly s3outputPath?: string;
+  readonly s3OutputUri?: string;
 
   /**
    * The guardrail is applied to the invocation
@@ -165,7 +165,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
 
     const isBodySpecified = props.body !== undefined;
     //Either specific props.input with bucket name and object key or input s3 path
-    const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.s3inputPath !== undefined);
+    const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.s3InputUri !== undefined);
 
     if (isBodySpecified && isInputSpecified) {
       throw new Error('Either `body` or `input` must be specified, but not both.');
@@ -191,7 +191,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       }),
     ];
 
-    if (this.props.s3inputPath !== undefined) {
+    if (this.props.s3InputUri !== undefined) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:GetObject'],
@@ -222,7 +222,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       );
     }
 
-    if (this.props.s3outputPath !== undefined) {
+    if (this.props.s3OutputUri !== undefined) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:PutObject'],
@@ -289,10 +289,10 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
         Body: this.props.body?.value,
         Input: this.props.input?.s3Location ? {
           S3Uri: `s3://${this.props.input.s3Location.bucketName}/${this.props.input.s3Location.objectKey}`,
-        } : this.props.s3inputPath ? { S3Uri: this.props.s3inputPath } : undefined,
+        } : this.props.s3InputUri ? { S3Uri: this.props.s3InputUri } : undefined,
         Output: this.props.output?.s3Location ? {
           S3Uri: `s3://${this.props.output.s3Location.bucketName}/${this.props.output.s3Location.objectKey}`,
-        } : this.props.s3outputPath ? { S3Uri: this.props.s3outputPath }: undefined,
+        } : this.props.s3OutputUri ? { S3Uri: this.props.s3OutputUri }: undefined,
         GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
         GuardrailVersion: this.props.guardrail?.guardrailVersion,
         Trace: this.props.traceEnabled === undefined
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
index 02fa2737bb2b8..cb74b84bfad25 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
@@ -210,8 +210,8 @@ describe('Invoke Model', () => {
 
     const task = new BedrockInvokeModel(stack, 'Invoke', {
       model,
-      inputPath: sfn.JsonPath.stringAt('$.prompt'),
-      outputPath: sfn.JsonPath.stringAt('$.prompt'),
+      s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
+      s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
     });
 
     new sfn.StateMachine(stack, 'StateMachine', {
@@ -234,8 +234,6 @@ describe('Invoke Model', () => {
         ],
       },
       End: true,
-      InputPath: '$.prompt',
-      OutputPath: '$.prompt',
       Parameters: {
         ModelId: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
         Input: {
@@ -330,8 +328,8 @@ describe('Invoke Model', () => {
     // WHEN
     const task = new BedrockInvokeModel(stack, 'Invoke', {
       model,
-      inputPath: sfn.JsonPath.stringAt('$.prompt'),
-      outputPath: sfn.JsonPath.stringAt('$.prompt'),
+      s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
+      s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
     });
 
     new sfn.StateMachine(stack, 'StateMachine', {

From 7682fa2c9216c26998a1188f65cbf137eb21d5ac Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Tue, 3 Sep 2024 13:40:36 -0700
Subject: [PATCH 03/18] update readme

---
 packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md b/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
index a8175c94cc8ff..594b2dd33a7c0 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
@@ -415,8 +415,8 @@ const model = bedrock.FoundationModel.fromFoundationModelId(
 
 const task = new tasks.BedrockInvokeModel(this, 'Prompt Model', {
   model,
-  inputPath: sfn.JsonPath.stringAt('$.prompt'),
-  outputPath: sfn.JsonPath.stringAt('$.prompt'),
+  s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
+  s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
 });
 
 ```

From 6c6ac5fbeda83930de1c053c69593906598e94a9 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Tue, 3 Sep 2024 16:23:09 -0700
Subject: [PATCH 04/18] move input and output s3 uri to bedrock invoke model
 props

---
 .../test/bedrock/integ.invoke-model.ts        |  4 +-
 .../lib/bedrock/invoke-model.ts               | 47 ++++++++++---------
 .../test/bedrock/invoke-model.test.ts         |  8 ++--
 3 files changed, 30 insertions(+), 29 deletions(-)

diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
index 2f91d3b3b173e..d33854ebf90f2 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
@@ -73,8 +73,8 @@ const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
 /** Test for Bedrock s3 URI Path */
 const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
   model,
-  s3InputUri: sfn.JsonPath.stringAt('$.names'),
-  s3OutputUri: sfn.JsonPath.stringAt('$.names'),
+  input: { s3InputUri: sfn.JsonPath.stringAt('$.names') },
+  output: { s3OutputUri: sfn.JsonPath.stringAt('$.names') },
 });
 
 const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4);
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
index 20d9692927b38..909875118d0c8 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
@@ -22,6 +22,15 @@ export interface BedrockInvokeModelInputProps {
    * @default - Input data is retrieved from the `body` field
    */
   readonly s3Location?: s3.Location;
+
+  /**
+   * The source location where the API response is written.
+   *
+   * This field can be used to specify s3 URI in the form of token
+   *
+   * @default - The API response body is returned in the result.
+   */
+  readonly s3InputUri?: string;
 }
 
 /**
@@ -40,6 +49,15 @@ export interface BedrockInvokeModelOutputProps {
    * @default - Response body is returned in the task result
    */
   readonly s3Location?: s3.Location;
+
+  /**
+   * The destination location where the API response is written.
+   *
+   * This field can be used to specify s3 URI in the form of token
+   *
+   * @default - The API response body is returned in the result.
+   */
+  readonly s3OutputUri?: string;
 }
 
 /**
@@ -108,24 +126,6 @@ export interface BedrockInvokeModelProps extends sfn.TaskStateBaseProps {
    */
   readonly output?: BedrockInvokeModelOutputProps;
 
-  /**
-   * The destination location where the API response is written.
-   *
-   * This field can be used to specify s3 URI in the form of token
-   *
-   * @default - The API response body is returned in the result.
-   */
-  readonly s3InputUri?: string;
-
-  /**
-   * The source location where the API response is written.
-   *
-   * This field can be used to specify s3 URI in the form of token
-   *
-   * @default - The API response body is returned in the result.
-   */
-  readonly s3OutputUri?: string;
-
   /**
    * The guardrail is applied to the invocation
    *
@@ -164,8 +164,9 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     validatePatternSupported(this.integrationPattern, BedrockInvokeModel.SUPPORTED_INTEGRATION_PATTERNS);
 
     const isBodySpecified = props.body !== undefined;
+
     //Either specific props.input with bucket name and object key or input s3 path
-    const isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.s3InputUri !== undefined);
+    const isInputSpecified = props.input!==undefined ? props.input?.s3Location !== undefined || props.input?.s3InputUri !== undefined : false;
 
     if (isBodySpecified && isInputSpecified) {
       throw new Error('Either `body` or `input` must be specified, but not both.');
@@ -191,7 +192,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       }),
     ];
 
-    if (this.props.s3InputUri !== undefined) {
+    if (this.props.input?.s3InputUri !== undefined) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:GetObject'],
@@ -222,7 +223,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       );
     }
 
-    if (this.props.s3OutputUri !== undefined) {
+    if (this.props.output?.s3OutputUri !== undefined) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:PutObject'],
@@ -289,10 +290,10 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
         Body: this.props.body?.value,
         Input: this.props.input?.s3Location ? {
           S3Uri: `s3://${this.props.input.s3Location.bucketName}/${this.props.input.s3Location.objectKey}`,
-        } : this.props.s3InputUri ? { S3Uri: this.props.s3InputUri } : undefined,
+        } : this.props.input?.s3InputUri ? { S3Uri: this.props.input?.s3InputUri } : undefined,
         Output: this.props.output?.s3Location ? {
           S3Uri: `s3://${this.props.output.s3Location.bucketName}/${this.props.output.s3Location.objectKey}`,
-        } : this.props.s3OutputUri ? { S3Uri: this.props.s3OutputUri }: undefined,
+        } : this.props.output?.s3OutputUri? { S3Uri: this.props.output.s3OutputUri }: undefined,
         GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
         GuardrailVersion: this.props.guardrail?.guardrailVersion,
         Trace: this.props.traceEnabled === undefined
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
index cb74b84bfad25..b8d5ff27c22a5 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
@@ -210,8 +210,8 @@ describe('Invoke Model', () => {
 
     const task = new BedrockInvokeModel(stack, 'Invoke', {
       model,
-      s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
-      s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
+      input: { s3InputUri: sfn.JsonPath.stringAt('$.prompt') },
+      output: { s3OutputUri: sfn.JsonPath.stringAt('$.prompt') },
     });
 
     new sfn.StateMachine(stack, 'StateMachine', {
@@ -328,8 +328,8 @@ describe('Invoke Model', () => {
     // WHEN
     const task = new BedrockInvokeModel(stack, 'Invoke', {
       model,
-      s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
-      s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
+      input: { s3InputUri: sfn.JsonPath.stringAt('$.prompt') },
+      output: { s3OutputUri: sfn.JsonPath.stringAt('$.prompt') },
     });
 
     new sfn.StateMachine(stack, 'StateMachine', {

From 070af9e17a5eddaef37240a287cb29ae445ec2fc Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Tue, 3 Sep 2024 19:10:24 -0700
Subject: [PATCH 05/18] adding feature flag

---
 .../test/bedrock/integ.invoke-model.ts        |  28 +++-
 .../aws-stepfunctions-tasks/README.md         |  38 ++++-
 .../lib/bedrock/invoke-model.ts               |  64 ++++++--
 .../test/bedrock/invoke-model.test.ts         | 155 +++++++++++++++++-
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md  |  21 ++-
 packages/aws-cdk-lib/cx-api/README.md         |  16 ++
 packages/aws-cdk-lib/cx-api/lib/features.ts   |  17 ++
 7 files changed, 316 insertions(+), 23 deletions(-)

diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
index d33854ebf90f2..1ce8ea14974bf 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
@@ -10,7 +10,11 @@ import { BedrockInvokeModel } from 'aws-cdk-lib/aws-stepfunctions-tasks';
  * * aws stepfunctions describe-execution --execution-arn <exection-arn generated before> : should return status as SUCCEEDED
  * This integ test does not actually verify a Step Functions execution, as not all AWS accounts have Bedrock model access.
  */
-const app = new cdk.App();
+const app = new cdk.App({
+  postCliContext: {
+    '@aws-cdk/aws-cdk.aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask': true,
+  },
+});
 const stack = new cdk.Stack(app, 'aws-stepfunctions-tasks-bedrock-invoke-model-integ');
 
 const model = bedrock.FoundationModel.fromFoundationModelId(stack, 'Model', bedrock.FoundationModelIdentifier.AMAZON_TITAN_TEXT_G1_EXPRESS_V1);
@@ -70,14 +74,32 @@ const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
   outputPath: sfn.JsonPath.stringAt('$.names'),
 });
 
-/** Test for Bedrock s3 URI Path */
+/**Test for Bedrock Input Path */
 const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
+  model,
+  body: sfn.TaskInput.fromObject(
+    {
+      inputText: sfn.JsonPath.format(
+        'Alphabetize this list of first names:\n{}',
+        sfn.JsonPath.stringAt('$.names'),
+      ),
+      textGenerationConfig: {
+        maxTokenCount: 100,
+        temperature: 1,
+      },
+    },
+  ),
+  inputPath: sfn.JsonPath.stringAt('$.names'),
+});
+
+/** Test for Bedrock s3 URI Path */
+const prompt5 = new BedrockInvokeModel(stack, 'Prompt5', {
   model,
   input: { s3InputUri: sfn.JsonPath.stringAt('$.names') },
   output: { s3OutputUri: sfn.JsonPath.stringAt('$.names') },
 });
 
-const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4);
+const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4).next(prompt5);
 
 new sfn.StateMachine(stack, 'StateMachine', {
   definitionBody: sfn.DefinitionBody.fromChainable(chain),
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md b/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
index 594b2dd33a7c0..d81574470dab8 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
@@ -398,11 +398,45 @@ const task = new tasks.BedrockInvokeModel(this, 'Prompt Model', {
     names: sfn.JsonPath.stringAt('$.Body.results[0].outputText'),
   },
 });
+
+```
+### Using Input Path for S3 URI
+
+Provide S3 URI as an input or output path to invoke a model
+
+To specify the S3 URI as JSON path to your input or output fields, use props `s3InputUri` and `s3OutputUri` under BedrockInvokeModelProps and set 
+feature flag `@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask` to true.
+
+If this flag is not set, then the existing behaviour of populating the S3Uri from `InputPath` and `OutputPath` will take effect.
+
+```ts
+
+import * as bedrock from 'aws-cdk-lib/aws-bedrock';
+
+const model = bedrock.FoundationModel.fromFoundationModelId(
+  this,
+  'Model',
+  bedrock.FoundationModelIdentifier.AMAZON_TITAN_TEXT_G1_EXPRESS_V1,
+);
+
+const task = new tasks.BedrockInvokeModel(this, 'Prompt Model', {
+  model,
+  input : { s3InputUri: sfn.JsonPath.stringAt('$.prompt') },
+  output: { s3OutputUri: sfn.JsonPath.stringAt('$.prompt') },
+});
+
 ```
+
 ### Using Input Path
 
 Provide S3 URI as an input or output path to invoke a model
 
+Currently, input and output Path provided in the BedrockInvokeModelProps input is defined as S3URI field under task definition of state machine.
+To modify the existing behaviour, set `@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask` to true. 
+
+If this feature flag is enabled, S3URI fields will be generated from other Props(`s3InputUri` and `s3OutputUri`), and the given inputPath, OutputPath will be rendered as 
+it is in the JSON task definition.
+
 ```ts
 
 import * as bedrock from 'aws-cdk-lib/aws-bedrock';
@@ -415,8 +449,8 @@ const model = bedrock.FoundationModel.fromFoundationModelId(
 
 const task = new tasks.BedrockInvokeModel(this, 'Prompt Model', {
   model,
-  s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
-  s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
+  inputPath: sfn.JsonPath.stringAt('$.prompt'),
+  outputPath: sfn.JsonPath.stringAt('$.prompt'),
 });
 
 ```
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
index 909875118d0c8..ebcda68d8c3e6 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
@@ -4,7 +4,8 @@ import * as bedrock from '../../../aws-bedrock';
 import * as iam from '../../../aws-iam';
 import * as s3 from '../../../aws-s3';
 import * as sfn from '../../../aws-stepfunctions';
-import { Stack } from '../../../core';
+import { Annotations, Stack, FeatureFlags } from '../../../core';
+import * as cxapi from '../../../cx-api';
 import { integrationResourceArn, validatePatternSupported } from '../private/task-utils';
 
 /**
@@ -163,10 +164,17 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
 
     validatePatternSupported(this.integrationPattern, BedrockInvokeModel.SUPPORTED_INTEGRATION_PATTERNS);
 
+    const isFeatureFlagEnabled = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
+
     const isBodySpecified = props.body !== undefined;
 
-    //Either specific props.input with bucket name and object key or input s3 path
-    const isInputSpecified = props.input!==undefined ? props.input?.s3Location !== undefined || props.input?.s3InputUri !== undefined : false;
+    let isInputSpecified: boolean;
+    if (!isFeatureFlagEnabled) {
+      isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.inputPath !== undefined);
+    } else {
+      //Either specific props.input with bucket name and object key or input s3 path
+      isInputSpecified = props.input!==undefined ? props.input?.s3Location !== undefined || props.input?.s3InputUri !== undefined : false;
+    }
 
     if (isBodySpecified && isInputSpecified) {
       throw new Error('Either `body` or `input` must be specified, but not both.');
@@ -181,10 +189,16 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       throw new Error('Output S3 object version is not supported.');
     }
 
-    this.taskPolicies = this.renderPolicyStatements();
+    //Warning to let users know about the newly introduced props
+    if (props.inputPath || props.outputPath && !isFeatureFlagEnabled) {
+      Annotations.of(scope).addWarningV2('aws-cdk-lib/aws-stepfunctions-taks',
+        'These props will set the value of inputPath/outputPath as s3 URI under input/output field in state machine JSON definition. To modify the behaviour set feature flag `@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true` and use props s3InputUri/s3OutputUri');
+    }
+
+    this.taskPolicies = this.renderPolicyStatements(isFeatureFlagEnabled);
   }
 
-  private renderPolicyStatements(): iam.PolicyStatement[] {
+  private renderPolicyStatements(isFeatureFlagEnabled?: boolean): iam.PolicyStatement[] {
     const policyStatements = [
       new iam.PolicyStatement({
         actions: ['bedrock:InvokeModel'],
@@ -192,7 +206,8 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       }),
     ];
 
-    if (this.props.input?.s3InputUri !== undefined) {
+    //For Compatibility with existing behaviour of input path
+    if (this.props.input?.s3InputUri !== undefined || (!isFeatureFlagEnabled && this.props.inputPath !== undefined)) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:GetObject'],
@@ -223,7 +238,8 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       );
     }
 
-    if (this.props.output?.s3OutputUri !== undefined) {
+    //For Compatibility with existing behaviour of output path
+    if (this.props.output?.s3OutputUri !== undefined || (!isFeatureFlagEnabled && this.props.outputPath !== undefined)) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:PutObject'],
@@ -281,6 +297,10 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
    * @internal
    */
   protected _renderTask(): any {
+
+    const isFeatureFlagEnabled = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
+    const inputSource = this.getInputSource(this.props.input, this.props.inputPath, isFeatureFlagEnabled);
+    const outputSource = this.getOutputSource(this.props.output, this.props.outputPath, isFeatureFlagEnabled);
     return {
       Resource: integrationResourceArn('bedrock', 'invokeModel'),
       Parameters: sfn.FieldUtils.renderObject({
@@ -288,12 +308,8 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
         Accept: this.props.accept,
         ContentType: this.props.contentType,
         Body: this.props.body?.value,
-        Input: this.props.input?.s3Location ? {
-          S3Uri: `s3://${this.props.input.s3Location.bucketName}/${this.props.input.s3Location.objectKey}`,
-        } : this.props.input?.s3InputUri ? { S3Uri: this.props.input?.s3InputUri } : undefined,
-        Output: this.props.output?.s3Location ? {
-          S3Uri: `s3://${this.props.output.s3Location.bucketName}/${this.props.output.s3Location.objectKey}`,
-        } : this.props.output?.s3OutputUri? { S3Uri: this.props.output.s3OutputUri }: undefined,
+        Input: this.props.input || (this.props.inputPath && !isFeatureFlagEnabled) ? { S3Uri: inputSource } : undefined,
+        Output: this.props.output || ( this.props.outputPath && !isFeatureFlagEnabled) ? { S3Uri: outputSource } : undefined,
         GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
         GuardrailVersion: this.props.guardrail?.guardrailVersion,
         Trace: this.props.traceEnabled === undefined
@@ -304,5 +320,27 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
       }),
     };
   };
+
+  private getInputSource(props?: BedrockInvokeModelInputProps, inputPath?: string, isFeatureFlagEnabled?: boolean): string | undefined {
+    if (props?.s3Location) {
+      return `s3://${props.s3Location.bucketName}/${props.s3Location.objectKey}`;
+    } else if (isFeatureFlagEnabled && props?.s3InputUri) {
+      return props.s3InputUri;
+    } else if (!isFeatureFlagEnabled && inputPath) {
+      return inputPath;
+    }
+    return undefined;
+  }
+
+  private getOutputSource(props?: BedrockInvokeModelOutputProps, outputPath?: string, isFeatureFlagEnabled?: boolean): string | undefined {
+    if (props?.s3Location) {
+      return `s3://${props.s3Location.bucketName}/${props.s3Location.objectKey}`;
+    } else if (isFeatureFlagEnabled && props?.s3OutputUri) {
+      return props.s3OutputUri;
+    } else if (!isFeatureFlagEnabled && outputPath) {
+      return outputPath;
+    }
+    return undefined;
+  }
 }
 
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
index b8d5ff27c22a5..e66cd7301e7d2 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
@@ -2,6 +2,7 @@ import { Template, Match } from '../../../assertions';
 import * as bedrock from '../../../aws-bedrock';
 import * as sfn from '../../../aws-stepfunctions';
 import * as cdk from '../../../core';
+import * as cxapi from '../../../cx-api';
 import { Guardrail } from '../../lib/bedrock/guardrail';
 import { BedrockInvokeModel } from '../../lib/bedrock/invoke-model';
 
@@ -204,14 +205,157 @@ describe('Invoke Model', () => {
     });
   });
 
-  test('invoke model allows input and output json path', () => {
-    const stack = new cdk.Stack();
+  test('invoke model allows input and output as s3 Uri with feature flag set to true', () => {
+    const app = new cdk.App({ context: { [cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: true } });
+    const stack = new cdk.Stack(app);
+    const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
+
+    const task = new BedrockInvokeModel(stack, 'Invoke', {
+      model,
+      input: { s3InputUri: sfn.JsonPath.stringAt('$.prompt') },
+      output: { s3OutputUri: sfn.JsonPath.stringAt('$.prompt') },
+    });
+
+    new sfn.StateMachine(stack, 'StateMachine', {
+      definitionBody: sfn.DefinitionBody.fromChainable(task),
+    });
+
+    // THEN
+    expect(stack.resolve(task.toStateJson())).toEqual({
+      Type: 'Task',
+      Resource: {
+        'Fn::Join': [
+          '',
+          [
+            'arn:',
+            {
+              Ref: 'AWS::Partition',
+            },
+            ':states:::bedrock:invokeModel',
+          ],
+        ],
+      },
+      End: true,
+      Parameters: {
+        ModelId: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
+        Input: {
+          //Expected key modified from S3Uri to S3Uri.$ as per the State Machine context key field transformation
+          //Reference: https://docs.aws.amazon.com/step-functions/latest/dg/input-output-example.html
+          'S3Uri.$': '$.prompt',
+        },
+        Output: {
+          'S3Uri.$': '$.prompt',
+        },
+      },
+    });
+  });
+
+  test('invoke model renderes input and output path correctly with feature flag set to true', () => {
+    const app = new cdk.App({ context: { [cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: true } });
+    const stack = new cdk.Stack(app);
     const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
 
     const task = new BedrockInvokeModel(stack, 'Invoke', {
       model,
       input: { s3InputUri: sfn.JsonPath.stringAt('$.prompt') },
       output: { s3OutputUri: sfn.JsonPath.stringAt('$.prompt') },
+      inputPath: sfn.JsonPath.stringAt('$.names'),
+      outputPath: sfn.JsonPath.stringAt('$.names'),
+    });
+
+    new sfn.StateMachine(stack, 'StateMachine', {
+      definitionBody: sfn.DefinitionBody.fromChainable(task),
+    });
+
+    // THEN
+    expect(stack.resolve(task.toStateJson())).toEqual({
+      Type: 'Task',
+      Resource: {
+        'Fn::Join': [
+          '',
+          [
+            'arn:',
+            {
+              Ref: 'AWS::Partition',
+            },
+            ':states:::bedrock:invokeModel',
+          ],
+        ],
+      },
+      End: true,
+      InputPath: '$.names',
+      OutputPath: '$.names',
+      Parameters: {
+        ModelId: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
+        Input: {
+          //Expected key modified from S3Uri to S3Uri.$ as per the State Machine context key field transformation
+          //Reference: https://docs.aws.amazon.com/step-functions/latest/dg/input-output-example.html
+          'S3Uri.$': '$.prompt',
+        },
+        Output: {
+          'S3Uri.$': '$.prompt',
+        },
+      },
+    });
+  });
+
+  //Should not throw an error for input path and body if feature flag is set to true
+  test('validation input and body correctly with feature flag set to true', () => {
+    const app = new cdk.App({ context: { [cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: true } });
+    const stack = new cdk.Stack(app);
+    const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
+
+    const task = new BedrockInvokeModel(stack, 'Invoke', {
+      model,
+      body: sfn.TaskInput.fromObject(
+        {
+          prompt: 'Hello world',
+        },
+      ),
+      inputPath: sfn.JsonPath.stringAt('$.names'),
+      outputPath: sfn.JsonPath.stringAt('$.names'),
+    });
+
+    new sfn.StateMachine(stack, 'StateMachine', {
+      definitionBody: sfn.DefinitionBody.fromChainable(task),
+    });
+
+    // THEN
+    expect(stack.resolve(task.toStateJson())).toEqual({
+      Type: 'Task',
+      Resource: {
+        'Fn::Join': [
+          '',
+          [
+            'arn:',
+            {
+              Ref: 'AWS::Partition',
+            },
+            ':states:::bedrock:invokeModel',
+          ],
+        ],
+      },
+      End: true,
+      InputPath: '$.names',
+      OutputPath: '$.names',
+      Parameters: {
+        ModelId: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
+        Body: {
+          prompt: 'Hello world',
+        },
+      },
+    });
+  });
+
+  test('invoke model renders input and output JSON Path as s3 URI with feature flag set to false', () => {
+    const app = new cdk.App({ context: { [cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: false } });
+    const stack = new cdk.Stack(app);
+    const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
+
+    const task = new BedrockInvokeModel(stack, 'Invoke', {
+      model,
+      inputPath: sfn.JsonPath.stringAt('$.prompt'),
+      outputPath: sfn.JsonPath.stringAt('$.prompt'),
     });
 
     new sfn.StateMachine(stack, 'StateMachine', {
@@ -234,6 +378,8 @@ describe('Invoke Model', () => {
         ],
       },
       End: true,
+      InputPath: '$.prompt',
+      OutputPath: '$.prompt',
       Parameters: {
         ModelId: 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123',
         Input: {
@@ -320,9 +466,10 @@ describe('Invoke Model', () => {
     });
   });
 
-  test('S3 permissions are created in generated policy when input and output path are specified', () => {
+  test('S3 permissions are created in generated policy when input/output s3uri is specified', () => {
     // GIVEN
-    const stack = new cdk.Stack();
+    const app = new cdk.App({ context: { [cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: true } });
+    const stack = new cdk.Stack(app);
     const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
 
     // WHEN
diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index d911731f3b9f9..9a786d4eb6060 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -72,6 +72,7 @@ Flags come in three types:
 | [@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm](#aws-cdkaws-ecsremovedefaultdeploymentalarm) | When enabled, remove default deployment alarm settings | 2.143.0 | (default) |
 | [@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault](#aws-cdkcustom-resourceslogapiresponsedatapropertytruedefault) | When enabled, the custom resource used for `AwsCustomResource` will configure the `logApiResponseData` property as true by default | 2.145.0 | (fix) |
 | [@aws-cdk/aws-s3:keepNotificationInImportedBucket](#aws-cdkaws-s3keepnotificationinimportedbucket) | When enabled, Adding notifications to a bucket in the current stack will not remove notification from imported stack. | 2.155.0 | (fix) |
+| [@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | 2.155.1 | (fix) |
 
 <!-- END table -->
 
@@ -133,7 +134,8 @@ The following json shows the current recommended set of flags, as `cdk init` wou
     "@aws-cdk/aws-ec2:ebsDefaultGp3Volume": true,
     "@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm": true,
     "@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault": false,
-    "@aws-cdk/aws-s3:keepNotificationInImportedBucket": false
+    "@aws-cdk/aws-s3:keepNotificationInImportedBucket": false,
+    "@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true
   }
 }
 ```
@@ -1356,4 +1358,21 @@ Other notifications that are not managed by this stack will be kept.
 | 2.155.0 | `false` | `false` |
 
 
+### @aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask
+
+*When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model.* (fix)
+
+Currently, 'inputPath' and 'outputPath' from the TaskStateBase Props is being used under BedrockInvokeModelProps to define S3URI under 'input' and 'output' fields
+of State Machine Task definition.
+
+When this feature flag is enabled, specify newly introduced props 's3InputUri' and 
+'s3OutputUri' to populate S3 uri under input and output fields in state machine task definition for Bedrock invoker model.
+
+
+| Since | Default | Recommended |
+| ----- | ----- | ----- |
+| (not in v1) |  |  |
+| 2.155.1 | `true` | `true` |
+
+
 <!-- END details -->
diff --git a/packages/aws-cdk-lib/cx-api/README.md b/packages/aws-cdk-lib/cx-api/README.md
index 06f364f9922f5..d094fd410d235 100644
--- a/packages/aws-cdk-lib/cx-api/README.md
+++ b/packages/aws-cdk-lib/cx-api/README.md
@@ -375,4 +375,20 @@ _cdk.json_
     "@aws-cdk/aws-stepfunctions-tasks:ecsReduceRunTaskPermissions": true
   }
 }
+```
+
+* `@aws-cdk/aws-stepfunctions-taks:useNewS3UriParametersForBedrockInvokeModelTask`
+
+When enabled, use new props for S3 URI under `input` and `output` fields in task definition of state machine for bedrock invoke model.
+
+When this feature flag is enabled, use newly introduced props `s3InputUri` and `s3OutputUri` to populate S3 uri under input and output fields in state machine task definition for Bedrock invoke model.
+
+_cdk.json_
+
+```json
+{
+  "context": {
+    "@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true
+  }
+}
 ```
\ No newline at end of file
diff --git a/packages/aws-cdk-lib/cx-api/lib/features.ts b/packages/aws-cdk-lib/cx-api/lib/features.ts
index 5ce4c237986d9..9463ca7ca8848 100644
--- a/packages/aws-cdk-lib/cx-api/lib/features.ts
+++ b/packages/aws-cdk-lib/cx-api/lib/features.ts
@@ -106,6 +106,7 @@ export const EBS_DEFAULT_GP3 = '@aws-cdk/aws-ec2:ebsDefaultGp3Volume';
 export const ECS_REMOVE_DEFAULT_DEPLOYMENT_ALARM = '@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm';
 export const LOG_API_RESPONSE_DATA_PROPERTY_TRUE_DEFAULT = '@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault';
 export const S3_KEEP_NOTIFICATION_IN_IMPORTED_BUCKET = '@aws-cdk/aws-s3:keepNotificationInImportedBucket';
+export const USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK = '@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask';
 
 export const FLAGS: Record<string, FlagInfo> = {
   //////////////////////////////////////////////////////////////////////
@@ -1107,6 +1108,22 @@ export const FLAGS: Record<string, FlagInfo> = {
     introducedIn: { v2: '2.155.0' },
     recommendedValue: false,
   },
+
+  //////////////////////////////////////////////////////////////////////
+  [USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: {
+    type: FlagType.BugFix,
+    summary: 'When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model.',
+    detailsMd: `
+    Currently, 'inputPath' and 'outputPath' from the TaskStateBase Props is being used under BedrockInvokeModelProps to define S3URI under 'input' and 'output' fields
+    of State Machine Task definition.
+
+    When this feature flag is enabled, specify newly introduced props 's3InputUri' and 
+    's3OutputUri' to populate S3 uri under input and output fields in state machine task definition for Bedrock invoker model.  
+
+    `,
+    introducedIn: { v2: '2.155.1' },
+    recommendedValue: true,
+  },
 };
 
 const CURRENT_MV = 'v2';

From 33cd3bcb4107eed0e347aaab8733490db0da7d43 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Wed, 4 Sep 2024 19:39:52 -0700
Subject: [PATCH 06/18] addressing comments and last feedback

---
 ...sks-bedrock-invoke-model-integ.assets.json |  4 +-
 ...s-bedrock-invoke-model-integ.template.json |  6 +--
 .../manifest.json                             |  2 +-
 .../integ.invoke-model.js.snapshot/tree.json  |  6 +--
 .../test/bedrock/integ.invoke-model.ts        | 33 ++++----------
 .../lib/bedrock/invoke-model.ts               | 45 +++++++++++--------
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md  |  4 +-
 7 files changed, 45 insertions(+), 55 deletions(-)

diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
index 0541b64d40527..d0ccc30d253bf 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
@@ -1,7 +1,7 @@
 {
   "version": "36.0.5",
   "files": {
-    "b69306f91f17e8d5e4b1d4e770e821de6110faa1c6811a445ce6a256e197e329": {
+    "d96d52e835333e009463f3e0e610457c80f74ce8aae54a095631544b72592209": {
       "source": {
         "path": "aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json",
         "packaging": "file"
@@ -9,7 +9,7 @@
       "destinations": {
         "current_account-current_region": {
           "bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}",
-          "objectKey": "b69306f91f17e8d5e4b1d4e770e821de6110faa1c6811a445ce6a256e197e329.json",
+          "objectKey": "d96d52e835333e009463f3e0e610457c80f74ce8aae54a095631544b72592209.json",
           "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}"
         }
       }
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
index 2a171f897511b..836e301d90e2b 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
@@ -103,7 +103,7 @@
        {
         "Ref": "AWS::Region"
        },
-       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
+       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names: {}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.Body.results[0].outputText\",\"Resource\":\"arn:",
        {
         "Ref": "AWS::Partition"
        },
@@ -115,7 +115,7 @@
        {
         "Ref": "AWS::Region"
        },
-       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
+       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Echo list of first names: {}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
        {
         "Ref": "AWS::Partition"
        },
@@ -127,7 +127,7 @@
        {
         "Ref": "AWS::Region"
        },
-       "::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri.$\":\"$.names\"},\"Output\":{\"S3Uri.$\":\"$.names\"}}}},\"TimeoutSeconds\":30}"
+       "::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri\":\"$.names\"},\"Output\":{\"S3Uri\":\"$.names\"}}}},\"TimeoutSeconds\":30}"
       ]
      ]
     },
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
index dd8726e66ddae..a0ff1e0d1ccd0 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
@@ -18,7 +18,7 @@
         "validateOnSynth": false,
         "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-deploy-role-${AWS::AccountId}-${AWS::Region}",
         "cloudFormationExecutionRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-cfn-exec-role-${AWS::AccountId}-${AWS::Region}",
-        "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/b69306f91f17e8d5e4b1d4e770e821de6110faa1c6811a445ce6a256e197e329.json",
+        "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/d96d52e835333e009463f3e0e610457c80f74ce8aae54a095631544b72592209.json",
         "requiresBootstrapStackVersion": 6,
         "bootstrapStackVersionSsmParameter": "/cdk-bootstrap/hnb659fds/version",
         "additionalDependencies": [
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
index 3f39dab1113e9..4bdd15996f1ae 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
@@ -194,7 +194,7 @@
                           {
                             "Ref": "AWS::Region"
                           },
-                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.names\",\"Resource\":\"arn:",
+                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names: {}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.Body.results[0].outputText\",\"Resource\":\"arn:",
                           {
                             "Ref": "AWS::Partition"
                           },
@@ -206,7 +206,7 @@
                           {
                             "Ref": "AWS::Region"
                           },
-                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:\\n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
+                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Echo list of first names: {}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt4\":{\"End\":true,\"Type\":\"Task\",\"Resource\":\"arn:",
                           {
                             "Ref": "AWS::Partition"
                           },
@@ -218,7 +218,7 @@
                           {
                             "Ref": "AWS::Region"
                           },
-                          "::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri.$\":\"$.names\"},\"Output\":{\"S3Uri.$\":\"$.names\"}}}},\"TimeoutSeconds\":30}"
+                          "::foundation-model/amazon.titan-text-express-v1\",\"Input\":{\"S3Uri\":\"$.names\"},\"Output\":{\"S3Uri\":\"$.names\"}}}},\"TimeoutSeconds\":30}"
                         ]
                       ]
                     },
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
index 1ce8ea14974bf..184b3f37794ef 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
@@ -41,7 +41,7 @@ const prompt2 = new BedrockInvokeModel(stack, 'Prompt2', {
   body: sfn.TaskInput.fromObject(
     {
       inputText: sfn.JsonPath.format(
-        'Alphabetize this list of first names:\n{}',
+        'Alphabetize this list of first names: {}',
         sfn.JsonPath.stringAt('$.names'),
       ),
       textGenerationConfig: {
@@ -62,7 +62,7 @@ const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
   body: sfn.TaskInput.fromObject(
     {
       inputText: sfn.JsonPath.format(
-        'Alphabetize this list of first names:\n{}',
+        'Echo list of first names: {}',
         sfn.JsonPath.stringAt('$.names'),
       ),
       textGenerationConfig: {
@@ -71,35 +71,18 @@ const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
       },
     },
   ),
-  outputPath: sfn.JsonPath.stringAt('$.names'),
-});
-
-/**Test for Bedrock Input Path */
-const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
-  model,
-  body: sfn.TaskInput.fromObject(
-    {
-      inputText: sfn.JsonPath.format(
-        'Alphabetize this list of first names:\n{}',
-        sfn.JsonPath.stringAt('$.names'),
-      ),
-      textGenerationConfig: {
-        maxTokenCount: 100,
-        temperature: 1,
-      },
-    },
-  ),
-  inputPath: sfn.JsonPath.stringAt('$.names'),
+  outputPath: '$.Body.results[0].outputText',
 });
 
 /** Test for Bedrock s3 URI Path */
-const prompt5 = new BedrockInvokeModel(stack, 'Prompt5', {
+//Execution will fail for the following input as it expects a valid s3 URI from previous prompt
+const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
   model,
-  input: { s3InputUri: sfn.JsonPath.stringAt('$.names') },
-  output: { s3OutputUri: sfn.JsonPath.stringAt('$.names') },
+  input: { s3InputUri: '$.names' },
+  output: { s3OutputUri: '$.names' },
 });
 
-const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4).next(prompt5);
+const chain = sfn.Chain.start(prompt1).next(prompt2).next(prompt3).next(prompt4);
 
 new sfn.StateMachine(stack, 'StateMachine', {
   definitionBody: sfn.DefinitionBody.fromChainable(chain),
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
index ebcda68d8c3e6..03a7837560676 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
@@ -164,12 +164,12 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
 
     validatePatternSupported(this.integrationPattern, BedrockInvokeModel.SUPPORTED_INTEGRATION_PATTERNS);
 
-    const isFeatureFlagEnabled = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
+    const useNewS3UriParamsForTask = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
 
     const isBodySpecified = props.body !== undefined;
 
     let isInputSpecified: boolean;
-    if (!isFeatureFlagEnabled) {
+    if (!useNewS3UriParamsForTask) {
       isInputSpecified = (props.input !== undefined && props.input.s3Location !== undefined) || (props.inputPath !== undefined);
     } else {
       //Either specific props.input with bucket name and object key or input s3 path
@@ -188,17 +188,24 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     if (props.output?.s3Location?.objectVersion !== undefined) {
       throw new Error('Output S3 object version is not supported.');
     }
+    if (props.input?.s3InputUri && props.input.s3Location) {
+      throw new Error('Cannot specify both S3 InputUri and S3 location');
+    }
+    if (props.input?.s3InputUri === '') {
+      throw new Error('S3 InputUri cannot be an empty string');
+    }
 
     //Warning to let users know about the newly introduced props
-    if (props.inputPath || props.outputPath && !isFeatureFlagEnabled) {
+    if (props.inputPath || props.outputPath && !useNewS3UriParamsForTask) {
       Annotations.of(scope).addWarningV2('aws-cdk-lib/aws-stepfunctions-taks',
-        'These props will set the value of inputPath/outputPath as s3 URI under input/output field in state machine JSON definition. To modify the behaviour set feature flag `@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true` and use props s3InputUri/s3OutputUri');
+        'These props will set the value of inputPath/outputPath as s3 URI under input/output field in state machine JSON definition. To modify the behaviour set feature flag `@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true` and use props input.s3InputUri/output.s3OutputUri');
     }
 
-    this.taskPolicies = this.renderPolicyStatements(isFeatureFlagEnabled);
+    this.taskPolicies = this.renderPolicyStatements();
   }
 
-  private renderPolicyStatements(isFeatureFlagEnabled?: boolean): iam.PolicyStatement[] {
+  private renderPolicyStatements(): iam.PolicyStatement[] {
+    const useNewS3UriParamsForTask = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
     const policyStatements = [
       new iam.PolicyStatement({
         actions: ['bedrock:InvokeModel'],
@@ -207,7 +214,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     ];
 
     //For Compatibility with existing behaviour of input path
-    if (this.props.input?.s3InputUri !== undefined || (!isFeatureFlagEnabled && this.props.inputPath !== undefined)) {
+    if (this.props.input?.s3InputUri !== undefined || (!useNewS3UriParamsForTask && this.props.inputPath !== undefined)) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:GetObject'],
@@ -239,7 +246,7 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     }
 
     //For Compatibility with existing behaviour of output path
-    if (this.props.output?.s3OutputUri !== undefined || (!isFeatureFlagEnabled && this.props.outputPath !== undefined)) {
+    if (this.props.output?.s3OutputUri !== undefined || (!useNewS3UriParamsForTask && this.props.outputPath !== undefined)) {
       policyStatements.push(
         new iam.PolicyStatement({
           actions: ['s3:PutObject'],
@@ -298,9 +305,9 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
    */
   protected _renderTask(): any {
 
-    const isFeatureFlagEnabled = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
-    const inputSource = this.getInputSource(this.props.input, this.props.inputPath, isFeatureFlagEnabled);
-    const outputSource = this.getOutputSource(this.props.output, this.props.outputPath, isFeatureFlagEnabled);
+    const useNewS3UriParamsForTask = FeatureFlags.of(this).isEnabled(cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK);
+    const inputSource = this.getInputSource(this.props.input, this.props.inputPath, useNewS3UriParamsForTask);
+    const outputSource = this.getOutputSource(this.props.output, this.props.outputPath, useNewS3UriParamsForTask);
     return {
       Resource: integrationResourceArn('bedrock', 'invokeModel'),
       Parameters: sfn.FieldUtils.renderObject({
@@ -308,8 +315,8 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
         Accept: this.props.accept,
         ContentType: this.props.contentType,
         Body: this.props.body?.value,
-        Input: this.props.input || (this.props.inputPath && !isFeatureFlagEnabled) ? { S3Uri: inputSource } : undefined,
-        Output: this.props.output || ( this.props.outputPath && !isFeatureFlagEnabled) ? { S3Uri: outputSource } : undefined,
+        Input: inputSource ? { S3Uri: inputSource } : undefined,
+        Output: outputSource ? { S3Uri: outputSource } : undefined,
         GuardrailIdentifier: this.props.guardrail?.guardrailIdentifier,
         GuardrailVersion: this.props.guardrail?.guardrailVersion,
         Trace: this.props.traceEnabled === undefined
@@ -321,23 +328,23 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     };
   };
 
-  private getInputSource(props?: BedrockInvokeModelInputProps, inputPath?: string, isFeatureFlagEnabled?: boolean): string | undefined {
+  private getInputSource(props?: BedrockInvokeModelInputProps, inputPath?: string, useNewS3UriParamsForTask?: boolean): string | undefined {
     if (props?.s3Location) {
       return `s3://${props.s3Location.bucketName}/${props.s3Location.objectKey}`;
-    } else if (isFeatureFlagEnabled && props?.s3InputUri) {
+    } else if (useNewS3UriParamsForTask && props?.s3InputUri) {
       return props.s3InputUri;
-    } else if (!isFeatureFlagEnabled && inputPath) {
+    } else if (!useNewS3UriParamsForTask && inputPath) {
       return inputPath;
     }
     return undefined;
   }
 
-  private getOutputSource(props?: BedrockInvokeModelOutputProps, outputPath?: string, isFeatureFlagEnabled?: boolean): string | undefined {
+  private getOutputSource(props?: BedrockInvokeModelOutputProps, outputPath?: string, useNewS3UriParamsForTask?: boolean): string | undefined {
     if (props?.s3Location) {
       return `s3://${props.s3Location.bucketName}/${props.s3Location.objectKey}`;
-    } else if (isFeatureFlagEnabled && props?.s3OutputUri) {
+    } else if (useNewS3UriParamsForTask && props?.s3OutputUri) {
       return props.s3OutputUri;
-    } else if (!isFeatureFlagEnabled && outputPath) {
+    } else if (!useNewS3UriParamsForTask && outputPath) {
       return outputPath;
     }
     return undefined;
diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index 9a786d4eb6060..87785d00f7478 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -72,7 +72,7 @@ Flags come in three types:
 | [@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm](#aws-cdkaws-ecsremovedefaultdeploymentalarm) | When enabled, remove default deployment alarm settings | 2.143.0 | (default) |
 | [@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault](#aws-cdkcustom-resourceslogapiresponsedatapropertytruedefault) | When enabled, the custom resource used for `AwsCustomResource` will configure the `logApiResponseData` property as true by default | 2.145.0 | (fix) |
 | [@aws-cdk/aws-s3:keepNotificationInImportedBucket](#aws-cdkaws-s3keepnotificationinimportedbucket) | When enabled, Adding notifications to a bucket in the current stack will not remove notification from imported stack. | 2.155.0 | (fix) |
-| [@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | 2.155.1 | (fix) |
+| [@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | 2.156.0 | (fix) |
 
 <!-- END table -->
 
@@ -1372,7 +1372,7 @@ When this feature flag is enabled, specify newly introduced props 's3InputUri' a
 | Since | Default | Recommended |
 | ----- | ----- | ----- |
 | (not in v1) |  |  |
-| 2.155.1 | `true` | `true` |
+| 2.156.0 | `true` | `true` |
 
 
 <!-- END details -->

From 47d60bd27364813dc563db50d6481734fd08cd92 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 02:23:36 -0700
Subject: [PATCH 07/18] update feature flag

---
 .../lib/bedrock/invoke-model.ts               |  8 ++--
 .../test/bedrock/invoke-model.test.ts         | 48 ++++++++++++++++++-
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md  |  4 +-
 packages/aws-cdk-lib/cx-api/lib/features.ts   |  2 +-
 4 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
index 03a7837560676..6d789ff85a356 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts
@@ -188,11 +188,11 @@ export class BedrockInvokeModel extends sfn.TaskStateBase {
     if (props.output?.s3Location?.objectVersion !== undefined) {
       throw new Error('Output S3 object version is not supported.');
     }
-    if (props.input?.s3InputUri && props.input.s3Location) {
-      throw new Error('Cannot specify both S3 InputUri and S3 location');
+    if (props.input?.s3InputUri && props.input.s3Location || props.output?.s3OutputUri && props.output.s3Location) {
+      throw new Error('Either specify S3 Uri or S3 location, but not both.');
     }
-    if (props.input?.s3InputUri === '') {
-      throw new Error('S3 InputUri cannot be an empty string');
+    if (useNewS3UriParamsForTask && (props.input?.s3InputUri === '' || props.output?.s3OutputUri === '')) {
+      throw new Error('S3 Uri cannot be an empty string');
     }
 
     //Warning to let users know about the newly introduced props
diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
index e66cd7301e7d2..4b5fdb894dba4 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
@@ -394,6 +394,50 @@ describe('Invoke Model', () => {
     });
   });
 
+  test('throws error S3 input Uri is specified as an empty string', () => {
+    const app = new cdk.App({ context: { [cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: true } });
+    const stack = new cdk.Stack(app);
+    const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
+
+    expect(() => {
+      new BedrockInvokeModel(stack, 'Invoke', {
+        model,
+        input: {
+          s3InputUri: '',
+        },
+        output: {
+          s3OutputUri: '',
+        },
+      });
+    }).toThrow('S3 Uri cannot be an empty string');
+  });
+
+  test('cannot specify both s3 uri and s3 bucket', () => {
+    const app = new cdk.App({ context: { [cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: true } });
+    const stack = new cdk.Stack(app);
+    const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
+
+    expect(() => {
+      new BedrockInvokeModel(stack, 'Invoke', {
+        model,
+        input: {
+          s3Location: {
+            bucketName: 'test-bucket',
+            objectKey: 'input-key',
+          },
+          s3InputUri: sfn.JsonPath.stringAt('$.prompt'),
+        },
+        output: {
+          s3Location: {
+            bucketName: 'test-bucket',
+            objectKey: 'output-key',
+          },
+          s3OutputUri: sfn.JsonPath.stringAt('$.prompt'),
+        },
+      });
+    }).toThrow('Either specify S3 Uri or S3 location, but not both.');
+  });
+
   test('S3 permissions are created in generated policy when input and output locations are specified', () => {
     // GIVEN
     const stack = new cdk.Stack();
@@ -475,8 +519,8 @@ describe('Invoke Model', () => {
     // WHEN
     const task = new BedrockInvokeModel(stack, 'Invoke', {
       model,
-      input: { s3InputUri: sfn.JsonPath.stringAt('$.prompt') },
-      output: { s3OutputUri: sfn.JsonPath.stringAt('$.prompt') },
+      input: { s3InputUri: 's3://input-bucket/input-key' },
+      output: { s3OutputUri: 's3://input-bucket/output-key' },
     });
 
     new sfn.StateMachine(stack, 'StateMachine', {
diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index 87785d00f7478..105449326312a 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -72,7 +72,7 @@ Flags come in three types:
 | [@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm](#aws-cdkaws-ecsremovedefaultdeploymentalarm) | When enabled, remove default deployment alarm settings | 2.143.0 | (default) |
 | [@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault](#aws-cdkcustom-resourceslogapiresponsedatapropertytruedefault) | When enabled, the custom resource used for `AwsCustomResource` will configure the `logApiResponseData` property as true by default | 2.145.0 | (fix) |
 | [@aws-cdk/aws-s3:keepNotificationInImportedBucket](#aws-cdkaws-s3keepnotificationinimportedbucket) | When enabled, Adding notifications to a bucket in the current stack will not remove notification from imported stack. | 2.155.0 | (fix) |
-| [@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | 2.156.0 | (fix) |
+| [@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | V2NEXT | (fix) |
 
 <!-- END table -->
 
@@ -1372,7 +1372,7 @@ When this feature flag is enabled, specify newly introduced props 's3InputUri' a
 | Since | Default | Recommended |
 | ----- | ----- | ----- |
 | (not in v1) |  |  |
-| 2.156.0 | `true` | `true` |
+| V2NEXT | `false` | `true` |
 
 
 <!-- END details -->
diff --git a/packages/aws-cdk-lib/cx-api/lib/features.ts b/packages/aws-cdk-lib/cx-api/lib/features.ts
index 9463ca7ca8848..db501b3ea03be 100644
--- a/packages/aws-cdk-lib/cx-api/lib/features.ts
+++ b/packages/aws-cdk-lib/cx-api/lib/features.ts
@@ -1121,7 +1121,7 @@ export const FLAGS: Record<string, FlagInfo> = {
     's3OutputUri' to populate S3 uri under input and output fields in state machine task definition for Bedrock invoker model.  
 
     `,
-    introducedIn: { v2: '2.155.1' },
+    introducedIn: { v2: 'V2NEXT' },
     recommendedValue: true,
   },
 };

From 7589bd40a89d7e30a73b569aae6c8097a9f6be1e Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 09:11:54 -0700
Subject: [PATCH 08/18] fix small nits

---
 .../aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts   | 2 +-
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md                    | 2 +-
 packages/aws-cdk-lib/cx-api/lib/features.ts                     | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
index 4b5fdb894dba4..869be41943937 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/test/bedrock/invoke-model.test.ts
@@ -300,7 +300,7 @@ describe('Invoke Model', () => {
   });
 
   //Should not throw an error for input path and body if feature flag is set to true
-  test('validation input and body correctly with feature flag set to true', () => {
+  test('validation for input and body correctly with feature flag set to true', () => {
     const app = new cdk.App({ context: { [cxapi.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: true } });
     const stack = new cdk.Stack(app);
     const model = bedrock.ProvisionedModel.fromProvisionedModelArn(stack, 'Imported', 'arn:aws:bedrock:us-turbo-2:123456789012:provisioned-model/abc-123');
diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index 105449326312a..55c2a51601acf 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -1366,7 +1366,7 @@ Currently, 'inputPath' and 'outputPath' from the TaskStateBase Props is being us
 of State Machine Task definition.
 
 When this feature flag is enabled, specify newly introduced props 's3InputUri' and 
-'s3OutputUri' to populate S3 uri under input and output fields in state machine task definition for Bedrock invoker model.
+'s3OutputUri' to populate S3 uri under input and output fields in state machine task definition for Bedrock invoke model.
 
 
 | Since | Default | Recommended |
diff --git a/packages/aws-cdk-lib/cx-api/lib/features.ts b/packages/aws-cdk-lib/cx-api/lib/features.ts
index db501b3ea03be..8fa033b658420 100644
--- a/packages/aws-cdk-lib/cx-api/lib/features.ts
+++ b/packages/aws-cdk-lib/cx-api/lib/features.ts
@@ -1118,7 +1118,7 @@ export const FLAGS: Record<string, FlagInfo> = {
     of State Machine Task definition.
 
     When this feature flag is enabled, specify newly introduced props 's3InputUri' and 
-    's3OutputUri' to populate S3 uri under input and output fields in state machine task definition for Bedrock invoker model.  
+    's3OutputUri' to populate S3 uri under input and output fields in state machine task definition for Bedrock invoke model.  
 
     `,
     introducedIn: { v2: 'V2NEXT' },

From 38b80bd05869292f7b9834caf90c71ea5fef8567 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 09:13:49 -0700
Subject: [PATCH 09/18] fix small nits

---
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index 55c2a51601acf..7eeaa4112a2c0 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -1372,7 +1372,7 @@ When this feature flag is enabled, specify newly introduced props 's3InputUri' a
 | Since | Default | Recommended |
 | ----- | ----- | ----- |
 | (not in v1) |  |  |
-| V2NEXT | `false` | `true` |
+| V2NEXT | `true` | `true` |
 
 
 <!-- END details -->

From 211e23f089956f663acd135ccc55d18fbc631301 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 09:26:06 -0700
Subject: [PATCH 10/18] fix README

---
 packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md b/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
index d81574470dab8..2067a10332668 100644
--- a/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
+++ b/packages/aws-cdk-lib/aws-stepfunctions-tasks/README.md
@@ -407,7 +407,8 @@ Provide S3 URI as an input or output path to invoke a model
 To specify the S3 URI as JSON path to your input or output fields, use props `s3InputUri` and `s3OutputUri` under BedrockInvokeModelProps and set 
 feature flag `@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask` to true.
 
-If this flag is not set, then the existing behaviour of populating the S3Uri from `InputPath` and `OutputPath` will take effect.
+
+If this flag is not enabled, the code will populate the S3Uri using `InputPath` and `OutputPath` fields, which is not recommended.
 
 ```ts
 
@@ -437,6 +438,8 @@ To modify the existing behaviour, set `@aws-cdk/aws-stepfunctions-tasks:useNewS3
 If this feature flag is enabled, S3URI fields will be generated from other Props(`s3InputUri` and `s3OutputUri`), and the given inputPath, OutputPath will be rendered as 
 it is in the JSON task definition.
 
+If the feature flag is set to `false`, the behavior will be to populate the S3Uri using the `InputPath` and `OutputPath` fields, which is not recommended.
+
 ```ts
 
 import * as bedrock from 'aws-cdk-lib/aws-bedrock';

From af747c8a489e53b6651111b539cc6b566f94329d Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 09:33:07 -0700
Subject: [PATCH 11/18] set feature flag default to true

---
 packages/aws-cdk-lib/cx-api/lib/features.ts | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/packages/aws-cdk-lib/cx-api/lib/features.ts b/packages/aws-cdk-lib/cx-api/lib/features.ts
index 8fa033b658420..612566a315452 100644
--- a/packages/aws-cdk-lib/cx-api/lib/features.ts
+++ b/packages/aws-cdk-lib/cx-api/lib/features.ts
@@ -1122,7 +1122,9 @@ export const FLAGS: Record<string, FlagInfo> = {
 
     `,
     introducedIn: { v2: 'V2NEXT' },
+    defaults: { v2: true },
     recommendedValue: true,
+    compatibilityWithOldBehaviorMd: 'Disable the feature flag to use input and output path fields for s3 URI',
   },
 };
 

From a52bdcf2f3298725cdce6bfce0897df4733a4b00 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 10:00:33 -0700
Subject: [PATCH 12/18] fix integ

---
 ...ctions-tasks-bedrock-invoke-model-integ.assets.json |  4 ++--
 ...ions-tasks-bedrock-invoke-model-integ.template.json |  2 +-
 .../integ.invoke-model.js.snapshot/manifest.json       |  2 +-
 .../bedrock/integ.invoke-model.js.snapshot/tree.json   |  2 +-
 .../test/bedrock/integ.invoke-model.ts                 | 10 +++-------
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md           |  1 +
 6 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
index d0ccc30d253bf..aab5d8193de11 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.assets.json
@@ -1,7 +1,7 @@
 {
   "version": "36.0.5",
   "files": {
-    "d96d52e835333e009463f3e0e610457c80f74ce8aae54a095631544b72592209": {
+    "b6d22af12502fb22f7fee80b3981c8767eb09e8e9f4ab927a380544cdcdc384b": {
       "source": {
         "path": "aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json",
         "packaging": "file"
@@ -9,7 +9,7 @@
       "destinations": {
         "current_account-current_region": {
           "bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}",
-          "objectKey": "d96d52e835333e009463f3e0e610457c80f74ce8aae54a095631544b72592209.json",
+          "objectKey": "b6d22af12502fb22f7fee80b3981c8767eb09e8e9f4ab927a380544cdcdc384b.json",
           "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}"
         }
       }
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
index 836e301d90e2b..ce041c6ffeaac 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/aws-stepfunctions-tasks-bedrock-invoke-model-integ.template.json
@@ -103,7 +103,7 @@
        {
         "Ref": "AWS::Region"
        },
-       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names: {}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.Body.results[0].outputText\",\"Resource\":\"arn:",
+       "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:/n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.Body.results[0].outputText\",\"Resource\":\"arn:",
        {
         "Ref": "AWS::Partition"
        },
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
index a0ff1e0d1ccd0..085c560bba157 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/manifest.json
@@ -18,7 +18,7 @@
         "validateOnSynth": false,
         "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-deploy-role-${AWS::AccountId}-${AWS::Region}",
         "cloudFormationExecutionRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-cfn-exec-role-${AWS::AccountId}-${AWS::Region}",
-        "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/d96d52e835333e009463f3e0e610457c80f74ce8aae54a095631544b72592209.json",
+        "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/b6d22af12502fb22f7fee80b3981c8767eb09e8e9f4ab927a380544cdcdc384b.json",
         "requiresBootstrapStackVersion": 6,
         "bootstrapStackVersionSsmParameter": "/cdk-bootstrap/hnb659fds/version",
         "additionalDependencies": [
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
index 4bdd15996f1ae..217a79721df41 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.js.snapshot/tree.json
@@ -194,7 +194,7 @@
                           {
                             "Ref": "AWS::Region"
                           },
-                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names: {}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.Body.results[0].outputText\",\"Resource\":\"arn:",
+                          "::foundation-model/amazon.titan-text-express-v1\",\"Body\":{\"inputText.$\":\"States.Format('Alphabetize this list of first names:/n{}', $.names)\",\"textGenerationConfig\":{\"maxTokenCount\":100,\"temperature\":1}}}},\"Prompt3\":{\"Next\":\"Prompt4\",\"Type\":\"Task\",\"OutputPath\":\"$.Body.results[0].outputText\",\"Resource\":\"arn:",
                           {
                             "Ref": "AWS::Partition"
                           },
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
index 184b3f37794ef..7ae58db6d8692 100644
--- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
+++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions-tasks/test/bedrock/integ.invoke-model.ts
@@ -10,11 +10,7 @@ import { BedrockInvokeModel } from 'aws-cdk-lib/aws-stepfunctions-tasks';
  * * aws stepfunctions describe-execution --execution-arn <exection-arn generated before> : should return status as SUCCEEDED
  * This integ test does not actually verify a Step Functions execution, as not all AWS accounts have Bedrock model access.
  */
-const app = new cdk.App({
-  postCliContext: {
-    '@aws-cdk/aws-cdk.aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask': true,
-  },
-});
+const app = new cdk.App();
 const stack = new cdk.Stack(app, 'aws-stepfunctions-tasks-bedrock-invoke-model-integ');
 
 const model = bedrock.FoundationModel.fromFoundationModelId(stack, 'Model', bedrock.FoundationModelIdentifier.AMAZON_TITAN_TEXT_G1_EXPRESS_V1);
@@ -41,7 +37,7 @@ const prompt2 = new BedrockInvokeModel(stack, 'Prompt2', {
   body: sfn.TaskInput.fromObject(
     {
       inputText: sfn.JsonPath.format(
-        'Alphabetize this list of first names: {}',
+        'Alphabetize this list of first names:/n{}',
         sfn.JsonPath.stringAt('$.names'),
       ),
       textGenerationConfig: {
@@ -75,7 +71,7 @@ const prompt3 = new BedrockInvokeModel(stack, 'Prompt3', {
 });
 
 /** Test for Bedrock s3 URI Path */
-//Execution will fail for the following input as it expects a valid s3 URI from previous prompt
+//State Machine Execution will fail for the following input as it expects a valid s3 URI from previous prompt
 const prompt4 = new BedrockInvokeModel(stack, 'Prompt4', {
   model,
   input: { s3InputUri: '$.names' },
diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index 7eeaa4112a2c0..8d7b2e200e425 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -1374,5 +1374,6 @@ When this feature flag is enabled, specify newly introduced props 's3InputUri' a
 | (not in v1) |  |  |
 | V2NEXT | `true` | `true` |
 
+**Compatibility with old behavior:** Disable the feature flag to use input and output path fields for s3 URI
 
 <!-- END details -->

From 451b44272d67b29a1f25b19eabcdb38ac242e47d Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 10:24:43 -0700
Subject: [PATCH 13/18] fix nit

---
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index 8d7b2e200e425..1d92c4582c8f6 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -135,7 +135,7 @@ The following json shows the current recommended set of flags, as `cdk init` wou
     "@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm": true,
     "@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault": false,
     "@aws-cdk/aws-s3:keepNotificationInImportedBucket": false,
-    "@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true
+    "@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true
   }
 }
 ```

From 0b7d4e8bc637cf299db41aabe3ef21c52f7fc0a1 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 10:41:53 -0700
Subject: [PATCH 14/18] fix last nit hopefully

---
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index 1d92c4582c8f6..7fa4b544b9ed4 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -72,7 +72,7 @@ Flags come in three types:
 | [@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm](#aws-cdkaws-ecsremovedefaultdeploymentalarm) | When enabled, remove default deployment alarm settings | 2.143.0 | (default) |
 | [@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault](#aws-cdkcustom-resourceslogapiresponsedatapropertytruedefault) | When enabled, the custom resource used for `AwsCustomResource` will configure the `logApiResponseData` property as true by default | 2.145.0 | (fix) |
 | [@aws-cdk/aws-s3:keepNotificationInImportedBucket](#aws-cdkaws-s3keepnotificationinimportedbucket) | When enabled, Adding notifications to a bucket in the current stack will not remove notification from imported stack. | 2.155.0 | (fix) |
-| [@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | V2NEXT | (fix) |
+| [@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | V2NEXT | (fix) |
 
 <!-- END table -->
 

From 4415b8977186f12d4461d00975d671a6b5725345 Mon Sep 17 00:00:00 2001
From: Shikha Aggarwal <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 11:01:31 -0700
Subject: [PATCH 15/18] Update features.ts

---
 packages/aws-cdk-lib/cx-api/lib/features.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/cx-api/lib/features.ts b/packages/aws-cdk-lib/cx-api/lib/features.ts
index 612566a315452..c4c5f13a0ee28 100644
--- a/packages/aws-cdk-lib/cx-api/lib/features.ts
+++ b/packages/aws-cdk-lib/cx-api/lib/features.ts
@@ -106,7 +106,7 @@ export const EBS_DEFAULT_GP3 = '@aws-cdk/aws-ec2:ebsDefaultGp3Volume';
 export const ECS_REMOVE_DEFAULT_DEPLOYMENT_ALARM = '@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm';
 export const LOG_API_RESPONSE_DATA_PROPERTY_TRUE_DEFAULT = '@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault';
 export const S3_KEEP_NOTIFICATION_IN_IMPORTED_BUCKET = '@aws-cdk/aws-s3:keepNotificationInImportedBucket';
-export const USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK = '@aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask';
+export const USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK = '@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask';
 
 export const FLAGS: Record<string, FlagInfo> = {
   //////////////////////////////////////////////////////////////////////

From 2ce70ffd122ebba93a963c0982a360f4541bb1d8 Mon Sep 17 00:00:00 2001
From: Shikha Aggarwal <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 11:33:25 -0700
Subject: [PATCH 16/18] set flag type to api default

---
 packages/aws-cdk-lib/cx-api/lib/features.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/cx-api/lib/features.ts b/packages/aws-cdk-lib/cx-api/lib/features.ts
index c4c5f13a0ee28..94f9294a9c85f 100644
--- a/packages/aws-cdk-lib/cx-api/lib/features.ts
+++ b/packages/aws-cdk-lib/cx-api/lib/features.ts
@@ -1111,7 +1111,7 @@ export const FLAGS: Record<string, FlagInfo> = {
 
   //////////////////////////////////////////////////////////////////////
   [USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: {
-    type: FlagType.BugFix,
+    type: FlagType.ApiDefault,
     summary: 'When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model.',
     detailsMd: `
     Currently, 'inputPath' and 'outputPath' from the TaskStateBase Props is being used under BedrockInvokeModelProps to define S3URI under 'input' and 'output' fields

From dfa27393fa13d56cdec8154957fb8207018a1013 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 10:44:22 -0700
Subject: [PATCH 17/18] fix nit

---
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index 7fa4b544b9ed4..fdc3e50d23051 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -1358,7 +1358,7 @@ Other notifications that are not managed by this stack will be kept.
 | 2.155.0 | `false` | `false` |
 
 
-### @aws-cdk:aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask
+### @aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask
 
 *When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model.* (fix)
 

From 0a7c25b9a62f51668ff8065c8395cbfab2ce2be4 Mon Sep 17 00:00:00 2001
From: shikha372 <shikhaaggarwal1995@gmail.com>
Date: Thu, 5 Sep 2024 10:46:28 -0700
Subject: [PATCH 18/18] final fix for successful build

---
 packages/@aws-cdk/cx-api/FEATURE_FLAGS.md     | 40 ++++++++++++++++++-
 packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md  |  8 ++--
 packages/aws-cdk-lib/cx-api/lib/features.ts   |  2 +-
 .../aws-cdk-lib/cx-api/test/features.test.ts  |  1 +
 4 files changed, 46 insertions(+), 5 deletions(-)

diff --git a/packages/@aws-cdk/cx-api/FEATURE_FLAGS.md b/packages/@aws-cdk/cx-api/FEATURE_FLAGS.md
index 9e980cc71faae..fdc3e50d23051 100644
--- a/packages/@aws-cdk/cx-api/FEATURE_FLAGS.md
+++ b/packages/@aws-cdk/cx-api/FEATURE_FLAGS.md
@@ -71,6 +71,8 @@ Flags come in three types:
 | [@aws-cdk/pipelines:reduceAssetRoleTrustScope](#aws-cdkpipelinesreduceassetroletrustscope) | Remove the root account principal from PipelineAssetsFileRole trust policy | 2.141.0 | (default) |
 | [@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm](#aws-cdkaws-ecsremovedefaultdeploymentalarm) | When enabled, remove default deployment alarm settings | 2.143.0 | (default) |
 | [@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault](#aws-cdkcustom-resourceslogapiresponsedatapropertytruedefault) | When enabled, the custom resource used for `AwsCustomResource` will configure the `logApiResponseData` property as true by default | 2.145.0 | (fix) |
+| [@aws-cdk/aws-s3:keepNotificationInImportedBucket](#aws-cdkaws-s3keepnotificationinimportedbucket) | When enabled, Adding notifications to a bucket in the current stack will not remove notification from imported stack. | 2.155.0 | (fix) |
+| [@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | V2NEXT | (fix) |
 
 <!-- END table -->
 
@@ -131,7 +133,9 @@ The following json shows the current recommended set of flags, as `cdk init` wou
     "@aws-cdk/aws-eks:nodegroupNameAttribute": true,
     "@aws-cdk/aws-ec2:ebsDefaultGp3Volume": true,
     "@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm": true,
-    "@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault": false
+    "@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault": false,
+    "@aws-cdk/aws-s3:keepNotificationInImportedBucket": false,
+    "@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true
   }
 }
 ```
@@ -1338,4 +1342,38 @@ property from the event object.
 | 2.145.0 | `false` | `false` |
 
 
+### @aws-cdk/aws-s3:keepNotificationInImportedBucket
+
+*When enabled, Adding notifications to a bucket in the current stack will not remove notification from imported stack.* (fix)
+
+Currently, adding notifications to a bucket where it was created by ourselves will override notification added where it is imported.
+
+When this feature flag is enabled, adding notifications to a bucket in the current stack will only update notification defined in this stack.
+Other notifications that are not managed by this stack will be kept.
+
+
+| Since | Default | Recommended |
+| ----- | ----- | ----- |
+| (not in v1) |  |  |
+| 2.155.0 | `false` | `false` |
+
+
+### @aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask
+
+*When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model.* (fix)
+
+Currently, 'inputPath' and 'outputPath' from the TaskStateBase Props is being used under BedrockInvokeModelProps to define S3URI under 'input' and 'output' fields
+of State Machine Task definition.
+
+When this feature flag is enabled, specify newly introduced props 's3InputUri' and 
+'s3OutputUri' to populate S3 uri under input and output fields in state machine task definition for Bedrock invoke model.
+
+
+| Since | Default | Recommended |
+| ----- | ----- | ----- |
+| (not in v1) |  |  |
+| V2NEXT | `true` | `true` |
+
+**Compatibility with old behavior:** Disable the feature flag to use input and output path fields for s3 URI
+
 <!-- END details -->
diff --git a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
index fdc3e50d23051..5d2b104c1f8bf 100644
--- a/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
+++ b/packages/aws-cdk-lib/cx-api/FEATURE_FLAGS.md
@@ -134,8 +134,7 @@ The following json shows the current recommended set of flags, as `cdk init` wou
     "@aws-cdk/aws-ec2:ebsDefaultGp3Volume": true,
     "@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm": true,
     "@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault": false,
-    "@aws-cdk/aws-s3:keepNotificationInImportedBucket": false,
-    "@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": true
+    "@aws-cdk/aws-s3:keepNotificationInImportedBucket": false
   }
 }
 ```
@@ -179,6 +178,7 @@ are migrating a v1 CDK project to v2, explicitly set any of these flags which do
 | [@aws-cdk/aws-lambda:recognizeVersionProps](#aws-cdkaws-lambdarecognizeversionprops) | Enable this feature flag to opt in to the updated logical id calculation for Lambda Version created using the  `fn.currentVersion`. | (fix) | 1.106.0 | `false` | `true` |
 | [@aws-cdk/aws-cloudfront:defaultSecurityPolicyTLSv1.2\_2021](#aws-cdkaws-cloudfrontdefaultsecuritypolicytlsv12_2021) | Enable this feature flag to have cloudfront distributions use the security policy TLSv1.2_2021 by default. | (fix) | 1.117.0 | `false` | `true` |
 | [@aws-cdk/pipelines:reduceAssetRoleTrustScope](#aws-cdkpipelinesreduceassetroletrustscope) | Remove the root account principal from PipelineAssetsFileRole trust policy | (default) |  | `false` | `true` |
+| [@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask](#aws-cdkaws-stepfunctions-tasksusenews3uriparametersforbedrockinvokemodeltask) | When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model. | (fix) |  | `false` | `true` |
 
 <!-- END diff -->
 
@@ -194,7 +194,8 @@ Here is an example of a `cdk.json` file that restores v1 behavior for these flag
     "@aws-cdk/aws-apigateway:usagePlanKeyOrderInsensitiveId": false,
     "@aws-cdk/aws-lambda:recognizeVersionProps": false,
     "@aws-cdk/aws-cloudfront:defaultSecurityPolicyTLSv1.2_2021": false,
-    "@aws-cdk/pipelines:reduceAssetRoleTrustScope": false
+    "@aws-cdk/pipelines:reduceAssetRoleTrustScope": false,
+    "@aws-cdk/aws-stepfunctions-tasks:useNewS3UriParametersForBedrockInvokeModelTask": false
   }
 }
 ```
@@ -1376,4 +1377,5 @@ When this feature flag is enabled, specify newly introduced props 's3InputUri' a
 
 **Compatibility with old behavior:** Disable the feature flag to use input and output path fields for s3 URI
 
+
 <!-- END details -->
diff --git a/packages/aws-cdk-lib/cx-api/lib/features.ts b/packages/aws-cdk-lib/cx-api/lib/features.ts
index 94f9294a9c85f..c4c5f13a0ee28 100644
--- a/packages/aws-cdk-lib/cx-api/lib/features.ts
+++ b/packages/aws-cdk-lib/cx-api/lib/features.ts
@@ -1111,7 +1111,7 @@ export const FLAGS: Record<string, FlagInfo> = {
 
   //////////////////////////////////////////////////////////////////////
   [USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: {
-    type: FlagType.ApiDefault,
+    type: FlagType.BugFix,
     summary: 'When enabled, use new props for S3 URI field in task definition of state machine for bedrock invoke model.',
     detailsMd: `
     Currently, 'inputPath' and 'outputPath' from the TaskStateBase Props is being used under BedrockInvokeModelProps to define S3URI under 'input' and 'output' fields
diff --git a/packages/aws-cdk-lib/cx-api/test/features.test.ts b/packages/aws-cdk-lib/cx-api/test/features.test.ts
index c1aeb3b80c0d1..470b2644ba2a6 100644
--- a/packages/aws-cdk-lib/cx-api/test/features.test.ts
+++ b/packages/aws-cdk-lib/cx-api/test/features.test.ts
@@ -38,6 +38,7 @@ test('feature flag defaults may not be changed anymore', () => {
     [feats.LAMBDA_RECOGNIZE_VERSION_PROPS]: true,
     [feats.CLOUDFRONT_DEFAULT_SECURITY_POLICY_TLS_V1_2_2021]: true,
     [feats.PIPELINE_REDUCE_ASSET_ROLE_TRUST_SCOPE]: true,
+    [feats.USE_NEW_S3URI_PARAMETERS_FOR_BEDROCK_INVOKE_MODEL_TASK]: true,
     // Add new disabling feature flags below this line
   });
 });