From 1c13faa46528f966b7165211a74760eb214e2aee Mon Sep 17 00:00:00 2001
From: Rico Huijbers <rix0rrr@gmail.com>
Date: Tue, 4 Jun 2019 09:47:27 +0200
Subject: [PATCH] fix(stepfunctions): improve Task payload encoding (#2706)

Improve referencing data fields for StepFunctions tasks, in preparation
of callback task implementaion.

Get rid of `JsonPath`, and in its place we have 2 new classes:

- `Data`, for fields that come from the user payload
  (`$.My.Field`). Settle on the term "data" since that's the
  term used in most of StepFunctions' docs.
- `Context`, for fields that come from the service-defined
  task "context" (like `$$.Execution.StartTime`, and in particular
  `$$.Task.Token`).

These classes have been moved from the `-tasks` module to the
`aws-stepfunctions` module, where it seems to make more sense for them
to live.

Add support for SQS and SNS tasks to publish an arbitrary JSON
structure that can reference fields from context and execution data.

Remove `NumberValue` since we can now encode Tokens in regular number
values.

BREAKING CHANGES:

- **stepfunctions**: `JsonPath.stringFromPath` (and others) are now called `Data.stringAt()`.
  The `DataField` class now lives in the main stepfunctions module.
- **stepfunctions**: `PublishToTopic` property `messageObject` used to take a JSON string, now pass `sfn.TaskInput.fromObject()` or `sfn.TaskInput.fromText()` into the `message` field.
- **stepfunctions**: `SendToQueue` property `messageBody` used to take a JSON string, now pass `sfn.TaskInput.fromObject()` or `sfn.TaskInput.fromText()` into the `message` field.
- **stepfunctions**: Instead of passing `NumberValue`s to StepFunctions tasks, pass regular
  numbers.
---
 .../aws-stepfunctions-tasks/lib/index.ts      |   4 +-
 .../aws-stepfunctions-tasks/lib/json-path.ts  | 119 -----------
 .../lib/number-value.ts                       |  53 -----
 .../lib/publish-to-topic.ts                   |  30 +--
 .../lib/run-ecs-task-base-types.ts            |   8 +-
 .../lib/run-ecs-task-base.ts                  |  19 +-
 .../lib/send-to-queue.ts                      |  18 +-
 .../test/ecs-tasks.test.ts                    |  11 +-
 .../test/integ.ec2-task.ts                    |   2 +-
 .../test/integ.fargate-task.ts                |   2 +-
 .../test/publish-to-topic.test.ts             |  28 ++-
 .../test/send-to-queue.test.ts                |  85 +++++++-
 packages/@aws-cdk/aws-stepfunctions/README.md |  46 +++-
 .../@aws-cdk/aws-stepfunctions/lib/fields.ts  | 136 ++++++++++++
 .../@aws-cdk/aws-stepfunctions/lib/index.ts   |   2 +
 .../@aws-cdk/aws-stepfunctions/lib/input.ts   |  58 ++++++
 .../aws-stepfunctions/lib/json-path.ts        | 196 ++++++++++++++++++
 .../aws-stepfunctions/test/test.fields.ts     | 135 ++++++++++++
 packages/@aws-cdk/cdk/lib/encoding.ts         |   4 +-
 packages/@aws-cdk/cdk/lib/resolve.ts          |   2 +-
 packages/@aws-cdk/cdk/lib/string-fragments.ts |  18 +-
 packages/@aws-cdk/cdk/lib/token-map.ts        |  14 +-
 22 files changed, 739 insertions(+), 251 deletions(-)
 delete mode 100644 packages/@aws-cdk/aws-stepfunctions-tasks/lib/json-path.ts
 delete mode 100644 packages/@aws-cdk/aws-stepfunctions-tasks/lib/number-value.ts
 create mode 100644 packages/@aws-cdk/aws-stepfunctions/lib/fields.ts
 create mode 100644 packages/@aws-cdk/aws-stepfunctions/lib/input.ts
 create mode 100644 packages/@aws-cdk/aws-stepfunctions/lib/json-path.ts
 create mode 100644 packages/@aws-cdk/aws-stepfunctions/test/test.fields.ts

diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/index.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/index.ts
index 3405a459919d1..0decc8f601c18 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/index.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/index.ts
@@ -5,6 +5,4 @@ export * from './run-ecs-task-base-types';
 export * from './publish-to-topic';
 export * from './send-to-queue';
 export * from './run-ecs-ec2-task';
-export * from './run-ecs-fargate-task';
-export * from './number-value';
-export * from './json-path';
+export * from './run-ecs-fargate-task';
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/json-path.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/json-path.ts
deleted file mode 100644
index 9b80f9e37a9d5..0000000000000
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/json-path.ts
+++ /dev/null
@@ -1,119 +0,0 @@
-import { Token, TokenMap } from '@aws-cdk/cdk';
-import { NumberValue } from './number-value';
-
-/**
- * Class to create special parameters for state machine states
- */
-export class JsonPath {
-  /**
-   * Instead of using a literal string, get the value from a JSON path
-   */
-  public static stringFromPath(path: string): string {
-    if (!path.startsWith('$.')) {
-      throw new Error("JSONPath values must start with '$.'");
-    }
-    return new JsonPathToken(path).toString();
-  }
-
-  /**
-   * Instead of using a literal string list, get the value from a JSON path
-   */
-  public static listFromPath(path: string): string[] {
-    if (!path.startsWith('$.')) {
-      throw new Error("JSONPath values must start with '$.'");
-    }
-    return new JsonPathToken(path).toList();
-  }
-
-  /**
-   * Get a number from a JSON path
-   */
-  public static numberFromPath(path: string): NumberValue {
-    return NumberValue.fromJsonPath(path);
-  }
-
-  private constructor() {
-  }
-}
-
-const JSON_PATH_TOKEN_SYMBOL = Symbol.for('JsonPathToken');
-
-class JsonPathToken extends Token {
-  public static isJsonPathToken(x: object): x is JsonPathToken {
-    return (x as any)[JSON_PATH_TOKEN_SYMBOL] === true;
-  }
-
-  constructor(public readonly path: string) {
-    super(() => path); // Make function to prevent eager evaluation in superclass
-    Object.defineProperty(this, JSON_PATH_TOKEN_SYMBOL, { value: true });
-  }
-}
-
-/**
- * Render a parameter string
- *
- * If the string value starts with '$.', render it as a path string, otherwise as a direct string.
- */
-export function renderString(key: string, value: string | undefined): {[key: string]: string} {
-  if (value === undefined) { return {}; }
-
-  const path = jsonPathString(value);
-  if (path !== undefined) {
-    return { [key + '.$']: path };
-  } else {
-    return { [key]: value };
-  }
-}
-
-/**
- * Render a parameter string
- *
- * If the string value starts with '$.', render it as a path string, otherwise as a direct string.
- */
-export function renderStringList(key: string, value: string[] | undefined): {[key: string]: string[] | string} {
-  if (value === undefined) { return {}; }
-
-  const path = jsonPathStringList(value);
-  if (path !== undefined) {
-    return { [key + '.$']: path };
-  } else {
-    return { [key]: value };
-  }
-}
-
-/**
- * Render a parameter string
- *
- * If the string value starts with '$.', render it as a path string, otherwise as a direct string.
- */
-export function renderNumber(key: string, value: NumberValue | undefined): {[key: string]: number | string} {
-  if (value === undefined) { return {}; }
-
-  if (!value.isLiteralNumber) {
-    return { [key + '.$']: value.jsonPath };
-  } else {
-    return { [key]: value.numberValue };
-  }
-}
-
-/**
- * If the indicated string is an encoded JSON path, return the path
- *
- * Otherwise return undefined.
- */
-function jsonPathString(x: string): string | undefined {
-  return pathFromToken(TokenMap.instance().lookupString(x));
-}
-
-/**
- * If the indicated string list is an encoded JSON path, return the path
- *
- * Otherwise return undefined.
- */
-function jsonPathStringList(x: string[]): string | undefined {
-  return pathFromToken(TokenMap.instance().lookupList(x));
-}
-
-function pathFromToken(token: Token | undefined) {
-  return token && (JsonPathToken.isJsonPathToken(token) ? token.path : undefined);
-}
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/number-value.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/number-value.ts
deleted file mode 100644
index 4d01b302ab504..0000000000000
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/number-value.ts
+++ /dev/null
@@ -1,53 +0,0 @@
-/**
- * A number value argument to a Task
- *
- * Either obtained from the current state, or from a literal number.
- *
- * This class is only necessary until https://github.com/awslabs/aws-cdk/issues/1455 is solved,
- * after which time we'll be able to use actual numbers to encode Tokens.
- */
-export class NumberValue {
-  /**
-   * Use a literal number
-   */
-  public static fromNumber(n: number): NumberValue {
-    return new NumberValue(n);
-  }
-
-  /**
-   * Obtain a number from the current state
-   */
-  public static fromJsonPath(path: string): NumberValue {
-    return new NumberValue(undefined, path);
-  }
-
-  private constructor(private readonly n?: number, private readonly path?: string) {
-  }
-
-  /**
-   * Return whether the NumberValue contains a literal number
-   */
-  public get isLiteralNumber(): boolean {
-    return this.n !== undefined;
-  }
-
-  /**
-   * Get the literal number from the NumberValue
-   */
-  public get numberValue(): number {
-    if (this.n === undefined) {
-      throw new Error('NumberValue does not have a number');
-    }
-    return this.n;
-  }
-
-  /**
-   * Get the JSON Path from the NumberValue
-   */
-  public get jsonPath(): string {
-    if (this.path === undefined) {
-      throw new Error('NumberValue does not have a JSONPath');
-    }
-    return this.path;
-  }
-}
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/publish-to-topic.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/publish-to-topic.ts
index 6bcda1589857f..3c50052d3c8c9 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/publish-to-topic.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/publish-to-topic.ts
@@ -1,26 +1,15 @@
 import iam = require('@aws-cdk/aws-iam');
 import sns = require('@aws-cdk/aws-sns');
 import sfn = require('@aws-cdk/aws-stepfunctions');
-import cdk = require('@aws-cdk/cdk');
-import { renderString } from './json-path';
 
 /**
  * Properties for PublishTask
  */
 export interface PublishToTopicProps {
   /**
-   * The text message to send to the queue.
-   *
-   * Exactly one of `message` and `messageObject` is required.
-   */
-  readonly message?: string;
-
-  /**
-   * Object to be JSON-encoded and used as message
-   *
-   * Exactly one of `message`, `messageObject` and `messagePath` is required.
+   * The text message to send to the topic.
    */
-  readonly messageObject?: string;
+  readonly message: sfn.TaskInput;
 
   /**
    * If true, send a different message to every subscription type
@@ -48,12 +37,9 @@ export interface PublishToTopicProps {
  */
 export class PublishToTopic implements sfn.IStepFunctionsTask {
   constructor(private readonly topic: sns.ITopic, private readonly props: PublishToTopicProps) {
-    if ((props.message === undefined) === (props.messageObject === undefined)) {
-      throw new Error(`Supply exactly one of 'message' or 'messageObject'`);
-    }
   }
 
-  public bind(task: sfn.Task): sfn.StepFunctionsTaskProperties {
+  public bind(_task: sfn.Task): sfn.StepFunctionsTaskProperties {
     return {
       resourceArn: 'arn:aws:states:::sns:publish',
       policyStatements: [new iam.PolicyStatement()
@@ -62,11 +48,11 @@ export class PublishToTopic implements sfn.IStepFunctionsTask {
       ],
       parameters: {
         TopicArn: this.topic.topicArn,
-        ...(this.props.messageObject
-            ? { Message: new cdk.Token(() => task.node.stringifyJson(this.props.messageObject)) }
-            : renderString('Message', this.props.message)),
-        MessageStructure: this.props.messagePerSubscriptionType ? "json" : undefined,
-        ...renderString('Subject', this.props.subject),
+        ...sfn.FieldUtils.renderObject({
+          Message: this.props.message.value,
+          MessageStructure: this.props.messagePerSubscriptionType ? "json" : undefined,
+          Subject: this.props.subject,
+        })
       }
     };
   }
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/run-ecs-task-base-types.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/run-ecs-task-base-types.ts
index 7e1ca0e1e21be..a4238287cd819 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/run-ecs-task-base-types.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/run-ecs-task-base-types.ts
@@ -1,5 +1,3 @@
-import { NumberValue } from "./number-value";
-
 export interface ContainerOverride {
   /**
    * Name of the container inside the task definition
@@ -23,21 +21,21 @@ export interface ContainerOverride {
    *
    * @Default The default value from the task definition.
    */
-  readonly cpu?: NumberValue;
+  readonly cpu?: number;
 
   /**
    * Hard memory limit on the container
    *
    * @Default The default value from the task definition.
    */
-  readonly memoryLimit?: NumberValue;
+  readonly memoryLimit?: number;
 
   /**
    * Soft memory limit on the container
    *
    * @Default The default value from the task definition.
    */
-  readonly memoryReservation?: NumberValue;
+  readonly memoryReservation?: number;
 }
 
 /**
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/run-ecs-task-base.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/run-ecs-task-base.ts
index 0e87b79c83720..064ba419c26ac 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/run-ecs-task-base.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/run-ecs-task-base.ts
@@ -3,7 +3,6 @@ import ecs = require('@aws-cdk/aws-ecs');
 import iam = require('@aws-cdk/aws-iam');
 import sfn = require('@aws-cdk/aws-stepfunctions');
 import cdk = require('@aws-cdk/cdk');
-import { renderNumber, renderString, renderStringList } from './json-path';
 import { ContainerOverride } from './run-ecs-task-base-types';
 
 /**
@@ -162,17 +161,17 @@ function renderOverrides(containerOverrides?: ContainerOverride[]) {
 
   const ret = new Array<any>();
   for (const override of containerOverrides) {
-    ret.push({
-      ...renderString('Name', override.containerName),
-      ...renderStringList('Command', override.command),
-      ...renderNumber('Cpu', override.cpu),
-      ...renderNumber('Memory', override.memoryLimit),
-      ...renderNumber('MemoryReservation', override.memoryReservation),
+    ret.push(sfn.FieldUtils.renderObject({
+      Name: override.containerName,
+      Command: override.command,
+      Cpu: override.cpu,
+      Memory: override.memoryLimit,
+      MemoryReservation: override.memoryReservation,
       Environment: override.environment && override.environment.map(e => ({
-        ...renderString('Name', e.name),
-        ...renderString('Value', e.value),
+        Name: e.name,
+        Value: e.value,
       }))
-    });
+    }));
   }
 
   return { ContainerOverrides: ret };
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/send-to-queue.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/send-to-queue.ts
index 820af3e1e83cf..51882f603774c 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/lib/send-to-queue.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/lib/send-to-queue.ts
@@ -1,17 +1,15 @@
 import iam = require('@aws-cdk/aws-iam');
 import sqs = require('@aws-cdk/aws-sqs');
 import sfn = require('@aws-cdk/aws-stepfunctions');
-import { renderNumber, renderString } from './json-path';
-import { NumberValue } from './number-value';
 
 /**
  * Properties for SendMessageTask
  */
 export interface SendToQueueProps {
   /**
-   * The message body to send to the queue.
+   * The text message to send to the queue.
    */
-  readonly messageBody: string;
+  readonly messageBody: sfn.TaskInput;
 
   /**
    * The length of time, in seconds, for which to delay a specific message.
@@ -20,7 +18,7 @@ export interface SendToQueueProps {
    *
    * @default Default value of the queue is used
    */
-  readonly delaySeconds?: NumberValue;
+  readonly delaySeconds?: number;
 
   /**
    * The token used for deduplication of sent messages.
@@ -59,10 +57,12 @@ export class SendToQueue implements sfn.IStepFunctionsTask {
       ],
       parameters: {
         QueueUrl: this.queue.queueUrl,
-        ...renderString('MessageBody', this.props.messageBody),
-        ...renderNumber('DelaySeconds', this.props.delaySeconds),
-        ...renderString('MessageDeduplicationId', this.props.messageDeduplicationId),
-        ...renderString('MessageGroupId', this.props.messageGroupId),
+        ...sfn.FieldUtils.renderObject({
+          MessageBody: this.props.messageBody.value,
+          DelaySeconds: this.props.delaySeconds,
+          MessageDeduplicationId: this.props.messageDeduplicationId,
+          MessageGroupId: this.props.messageGroupId,
+        })
       }
     };
   }
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs-tasks.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs-tasks.test.ts
index 49e37be3036ef..b1edbbd9a1de7 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs-tasks.test.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/ecs-tasks.test.ts
@@ -4,7 +4,6 @@ import ecs = require('@aws-cdk/aws-ecs');
 import sfn = require('@aws-cdk/aws-stepfunctions');
 import { Stack } from '@aws-cdk/cdk';
 import tasks = require('../lib');
-import { JsonPath, NumberValue } from '../lib';
 
 let stack: Stack;
 let vpc: ec2.Vpc;
@@ -64,7 +63,7 @@ test('Running a Fargate Task', () => {
       {
         containerName: 'TheContainer',
         environment: [
-          {name: 'SOME_KEY', value: JsonPath.stringFromPath('$.SomeKey')}
+          {name: 'SOME_KEY', value: sfn.Data.stringAt('$.SomeKey')}
         ]
       }
     ]
@@ -162,7 +161,7 @@ test('Running an EC2 Task with bridge network', () => {
       {
         containerName: 'TheContainer',
         environment: [
-          {name: 'SOME_KEY', value: JsonPath.stringFromPath('$.SomeKey')}
+          {name: 'SOME_KEY', value: sfn.Data.stringAt('$.SomeKey')}
         ]
       }
     ]
@@ -296,9 +295,9 @@ test('Running an EC2 Task with overridden number values', () => {
     containerOverrides: [
       {
         containerName: 'TheContainer',
-        command: JsonPath.listFromPath('$.TheCommand'),
-        cpu: NumberValue.fromNumber(5),
-        memoryLimit: JsonPath.numberFromPath('$.MemoryLimit'),
+        command: sfn.Data.listAt('$.TheCommand'),
+        cpu: 5,
+        memoryLimit: sfn.Data.numberAt('$.MemoryLimit'),
       }
     ]
   });
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.ec2-task.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.ec2-task.ts
index 49d1a0d4dee9c..9ce9f362534d8 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.ec2-task.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.ec2-task.ts
@@ -37,7 +37,7 @@ const definition = new sfn.Pass(stack, 'Start', {
       environment: [
         {
           name: 'SOME_KEY',
-          value: tasks.JsonPath.stringFromPath('$.SomeKey')
+          value: sfn.Data.stringAt('$.SomeKey')
         }
       ]
     }
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.fargate-task.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.fargate-task.ts
index 320c278437ab6..5589cf0e1ef67 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.fargate-task.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/integ.fargate-task.ts
@@ -37,7 +37,7 @@ const definition = new sfn.Pass(stack, 'Start', {
       environment: [
         {
           name: 'SOME_KEY',
-          value: tasks.JsonPath.stringFromPath('$.SomeKey')
+          value: sfn.Data.stringAt('$.SomeKey')
         }
       ]
     }
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/publish-to-topic.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/publish-to-topic.test.ts
index aa2840864c07a..ad846441a73d6 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/publish-to-topic.test.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/publish-to-topic.test.ts
@@ -10,7 +10,7 @@ test('publish to SNS', () => {
 
   // WHEN
   const pub = new sfn.Task(stack, 'Publish', { task: new tasks.PublishToTopic(topic, {
-    message: 'Send this message'
+    message: sfn.TaskInput.fromText('Send this message')
   }) });
 
   // THEN
@@ -24,3 +24,29 @@ test('publish to SNS', () => {
     },
   });
 });
+
+test('publish JSON to SNS', () => {
+  // GIVEN
+  const stack = new cdk.Stack();
+  const topic = new sns.Topic(stack, 'Topic');
+
+  // WHEN
+  const pub = new sfn.Task(stack, 'Publish', { task: new tasks.PublishToTopic(topic, {
+    message: sfn.TaskInput.fromObject({
+      Input: 'Send this message'
+    })
+  }) });
+
+  // THEN
+  expect(stack.node.resolve(pub.toStateJson())).toEqual({
+    Type: 'Task',
+    Resource: 'arn:aws:states:::sns:publish',
+    End: true,
+    Parameters: {
+      TopicArn: { Ref: 'TopicBFC7AF6E' },
+      Message: {
+        Input: 'Send this message'
+      }
+    },
+  });
+});
diff --git a/packages/@aws-cdk/aws-stepfunctions-tasks/test/send-to-queue.test.ts b/packages/@aws-cdk/aws-stepfunctions-tasks/test/send-to-queue.test.ts
index 6af5300dba9aa..6bd53350058fa 100644
--- a/packages/@aws-cdk/aws-stepfunctions-tasks/test/send-to-queue.test.ts
+++ b/packages/@aws-cdk/aws-stepfunctions-tasks/test/send-to-queue.test.ts
@@ -3,15 +3,20 @@ import sfn = require('@aws-cdk/aws-stepfunctions');
 import cdk = require('@aws-cdk/cdk');
 import tasks = require('../lib');
 
-test('publish to queue', () => {
+let stack: cdk.Stack;
+let queue: sqs.Queue;
+
+beforeEach(() => {
   // GIVEN
-  const stack = new cdk.Stack();
-  const queue = new sqs.Queue(stack, 'Queue');
+  stack = new cdk.Stack();
+  queue = new sqs.Queue(stack, 'Queue');
+});
 
+test('publish to queue', () => {
   // WHEN
   const pub = new sfn.Task(stack, 'Send', { task: new tasks.SendToQueue(queue, {
-    messageBody: 'Send this message',
-    messageDeduplicationId: tasks.JsonPath.stringFromPath('$.deduping'),
+    messageBody: sfn.TaskInput.fromText('Send this message'),
+    messageDeduplicationId: sfn.Data.stringAt('$.deduping'),
   }) });
 
   // THEN
@@ -25,4 +30,74 @@ test('publish to queue', () => {
       'MessageDeduplicationId.$': '$.deduping'
     },
   });
+});
+
+test('message body can come from state', () => {
+  // WHEN
+  const pub = new sfn.Task(stack, 'Send', {
+    task: new tasks.SendToQueue(queue, {
+      messageBody: sfn.TaskInput.fromDataAt('$.theMessage')
+    })
+  });
+
+  // THEN
+  expect(stack.node.resolve(pub.toStateJson())).toEqual({
+    Type: 'Task',
+    Resource: 'arn:aws:states:::sqs:sendMessage',
+    End: true,
+    Parameters: {
+      'QueueUrl': { Ref: 'Queue4A7E3555' },
+      'MessageBody.$': '$.theMessage',
+    },
+  });
+});
+
+test('message body can be an object', () => {
+  // WHEN
+  const pub = new sfn.Task(stack, 'Send', {
+    task: new tasks.SendToQueue(queue, {
+      messageBody: sfn.TaskInput.fromObject({
+        literal: 'literal',
+        SomeInput: sfn.Data.stringAt('$.theMessage')
+      })
+    })
+  });
+
+  // THEN
+  expect(stack.node.resolve(pub.toStateJson())).toEqual({
+    Type: 'Task',
+    Resource: 'arn:aws:states:::sqs:sendMessage',
+    End: true,
+    Parameters: {
+      QueueUrl: { Ref: 'Queue4A7E3555' },
+      MessageBody: {
+        'literal': 'literal',
+        'SomeInput.$': '$.theMessage',
+      }
+    },
+  });
+});
+
+test('message body object can contain references', () => {
+  // WHEN
+  const pub = new sfn.Task(stack, 'Send', {
+    task: new tasks.SendToQueue(queue, {
+      messageBody: sfn.TaskInput.fromObject({
+        queueArn: queue.queueArn
+      })
+    })
+  });
+
+  // THEN
+  expect(stack.node.resolve(pub.toStateJson())).toEqual({
+    Type: 'Task',
+    Resource: 'arn:aws:states:::sqs:sendMessage',
+    End: true,
+    Parameters: {
+      QueueUrl: { Ref: 'Queue4A7E3555' },
+      MessageBody: {
+        queueArn: { 'Fn::GetAtt': ['Queue4A7E3555', 'Arn'] }
+      }
+    },
+  });
 });
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions/README.md b/packages/@aws-cdk/aws-stepfunctions/README.md
index 3640e5d6d09ee..4d51b322514a5 100644
--- a/packages/@aws-cdk/aws-stepfunctions/README.md
+++ b/packages/@aws-cdk/aws-stepfunctions/README.md
@@ -119,8 +119,8 @@ couple of the tasks available are:
 
 Many tasks take parameters. The values for those can either be supplied
 directly in the workflow definition (by specifying their values), or at
-runtime by passing a value obtained from the static functions on `JsonPath`,
-such as `JsonPath.stringFromPath()`.
+runtime by passing a value obtained from the static functions on `Data`,
+such as `Data.stringAt()`.
 
 If so, the value is taken from the indicated location in the state JSON,
 similar to (for example) `inputPath`.
@@ -155,9 +155,22 @@ import sns = require('@aws-cdk/aws-sns');
 // ...
 
 const topic = new sns.Topic(this, 'Topic');
-const task = new sfn.Task(this, 'Publish', {
+
+// Use a field from the execution data as message.
+const task1 = new sfn.Task(this, 'Publish1', {
+    task: new tasks.PublishToTopic(topic, {
+        message: TaskInput.fromDataAt('$.state.message'),
+    })
+});
+
+// Combine a field from the execution data with
+// a literal object.
+const task2 = new sfn.Task(this, 'Publish2', {
     task: new tasks.PublishToTopic(topic, {
-        message: JsonPath.stringFromPath('$.state.message'),
+        message: TaskInput.fromObject({
+            field1: 'somedata',
+            field2: Data.stringAt('$.field2'),
+        })
     })
 });
 ```
@@ -170,11 +183,26 @@ import sqs = require('@aws-cdk/aws-sqs');
 // ...
 
 const queue = new sns.Queue(this, 'Queue');
-const task = new sfn.Task(this, 'Send', {
+
+// Use a field from the execution data as message.
+const task1 = new sfn.Task(this, 'Send1', {
+    task: new tasks.SendToQueue(queue, {
+        messageBody: TaskInput.fromDataAt('$.message'),
+        // Only for FIFO queues
+        messageGroupId: '1234'
+    })
+});
+
+// Combine a field from the execution data with
+// a literal object.
+const task2 = new sfn.Task(this, 'Send2', {
     task: new tasks.SendToQueue(queue, {
-        messageBody: JsonPath.stringFromPath('$.message'),
+        messageBody: TaskInput.fromObject({
+            field1: 'somedata',
+            field2: Data.stringAt('$.field2'),
+        }),
         // Only for FIFO queues
-        messageGroupId: JsonPath.stringFromPath('$.messageGroupId'),
+        messageGroupId: '1234'
     })
 });
 ```
@@ -195,7 +223,7 @@ const fargateTask = new ecs.RunEcsFargateTask({
       environment: [
         {
           name: 'CONTAINER_INPUT',
-          value: JsonPath.stringFromPath('$.valueFromStateData')
+          value: Data.stringAt('$.valueFromStateData')
         }
       ]
     }
@@ -464,4 +492,4 @@ Contributions welcome:
 - [ ] A single `LambdaTask` class that is both a `Lambda` and a `Task` in one
   might make for a nice API.
 - [ ] Expression parser for Conditions.
-- [ ] Simulate state machines in unit tests.
\ No newline at end of file
+- [ ] Simulate state machines in unit tests.
diff --git a/packages/@aws-cdk/aws-stepfunctions/lib/fields.ts b/packages/@aws-cdk/aws-stepfunctions/lib/fields.ts
new file mode 100644
index 0000000000000..4fd4e2a693e05
--- /dev/null
+++ b/packages/@aws-cdk/aws-stepfunctions/lib/fields.ts
@@ -0,0 +1,136 @@
+import { findReferencedPaths, JsonPathToken, renderObject } from "./json-path";
+
+/**
+ * Extract a field from the State Machine data that gets passed around between states
+ */
+export class Data {
+  /**
+   * Instead of using a literal string, get the value from a JSON path
+   */
+  public static stringAt(path: string): string {
+    validateDataPath(path);
+    return new JsonPathToken(path).toString();
+  }
+
+  /**
+   * Instead of using a literal string list, get the value from a JSON path
+   */
+  public static listAt(path: string): string[] {
+    validateDataPath(path);
+    return new JsonPathToken(path).toList();
+  }
+
+  /**
+   * Instead of using a literal number, get the value from a JSON path
+   */
+  public static numberAt(path: string): number {
+    validateDataPath(path);
+    return new JsonPathToken(path).toNumber();
+  }
+
+  /**
+   * Use the entire data structure
+   *
+   * Will be an object at invocation time, but is represented in the CDK
+   * application as a string.
+   */
+  public static get entirePayload(): string {
+    return new JsonPathToken('$').toString();
+  }
+
+  private constructor() {
+  }
+}
+
+/**
+ * Extract a field from the State Machine Context data
+ *
+ * @see https://docs.aws.amazon.com/step-functions/latest/dg/connect-to-resource.html#wait-token-contextobject
+ */
+export class Context {
+  /**
+   * Instead of using a literal string, get the value from a JSON path
+   */
+  public static stringAt(path: string): string {
+    validateContextPath(path);
+    return new JsonPathToken(path).toString();
+  }
+
+  /**
+   * Instead of using a literal number, get the value from a JSON path
+   */
+  public static numberAt(path: string): number {
+    validateContextPath(path);
+    return new JsonPathToken(path).toNumber();
+  }
+
+  /**
+   * Return the Task Token field
+   *
+   * External actions will need this token to report step completion
+   * back to StepFunctions using the `SendTaskSuccess` or `SendTaskFailure`
+   * calls.
+   */
+  public static get taskToken(): string {
+    return new JsonPathToken('$$.Task.Token').toString();
+  }
+
+  /**
+   * Use the entire context data structure
+   *
+   * Will be an object at invocation time, but is represented in the CDK
+   * application as a string.
+   */
+  public static get entireContext(): string {
+    return new JsonPathToken('$$').toString();
+  }
+
+  private constructor() {
+  }
+}
+
+/**
+ * Helper functions to work with structures containing fields
+ */
+export class FieldUtils {
+
+  /**
+   * Render a JSON structure containing fields to the right StepFunctions structure
+   */
+  public static renderObject(obj?: {[key: string]: any}): {[key: string]: any}  | undefined {
+    return renderObject(obj);
+  }
+
+  /**
+   * Return all JSON paths used in the given structure
+   */
+  public static findReferencedPaths(obj?: {[key: string]: any}): string[] {
+    return Array.from(findReferencedPaths(obj)).sort();
+  }
+
+  /**
+   * Returns whether the given task structure contains the TaskToken field anywhere
+   *
+   * The field is considered included if the field itself or one of its containing
+   * fields occurs anywhere in the payload.
+   */
+  public static containsTaskToken(obj?: {[key: string]: any}): boolean {
+    const paths = findReferencedPaths(obj);
+    return paths.has('$$.Task.Token') || paths.has('$$.Task') || paths.has('$$');
+  }
+
+  private constructor() {
+  }
+}
+
+function validateDataPath(path: string) {
+  if (!path.startsWith('$.')) {
+    throw new Error("Data JSON path values must start with '$.'");
+  }
+}
+
+function validateContextPath(path: string) {
+  if (!path.startsWith('$$.')) {
+    throw new Error("Context JSON path values must start with '$$.'");
+  }
+}
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions/lib/index.ts b/packages/@aws-cdk/aws-stepfunctions/lib/index.ts
index a550a54412b66..c86d357b3ac02 100644
--- a/packages/@aws-cdk/aws-stepfunctions/lib/index.ts
+++ b/packages/@aws-cdk/aws-stepfunctions/lib/index.ts
@@ -1,4 +1,6 @@
+export * from './fields';
 export * from './activity';
+export * from './input';
 export * from './types';
 export * from './condition';
 export * from './state-machine';
diff --git a/packages/@aws-cdk/aws-stepfunctions/lib/input.ts b/packages/@aws-cdk/aws-stepfunctions/lib/input.ts
new file mode 100644
index 0000000000000..dedb89ad9af87
--- /dev/null
+++ b/packages/@aws-cdk/aws-stepfunctions/lib/input.ts
@@ -0,0 +1,58 @@
+import { Context, Data } from "./fields";
+
+/**
+ * Type union for task classes that accept multiple types of payload
+ */
+export class TaskInput {
+  /**
+   * Use a literal string as task input
+   *
+   * This might be a JSON-encoded object, or just a text.
+   */
+  public static fromText(text: string) {
+    return new TaskInput(InputType.Text, text);
+  }
+
+  /**
+   * Use an object as task input
+   *
+   * This object may contain Data and Context fields
+   * as object values, if desired.
+   */
+  public static fromObject(obj: {[key: string]: any}) {
+    return new TaskInput(InputType.Object, obj);
+  }
+
+  /**
+   * Use a part of the execution data as task input
+   *
+   * Use this when you want to use a subobject or string from
+   * the current state machine execution as complete payload
+   * to a task.
+   */
+  public static fromDataAt(path: string) {
+    return new TaskInput(InputType.Text, Data.stringAt(path));
+  }
+
+  /**
+   * Use a part of the task context as task input
+   *
+   * Use this when you want to use a subobject or string from
+   * the current task context as complete payload
+   * to a task.
+   */
+  public static fromContextAt(path: string) {
+    return new TaskInput(InputType.Text, Context.stringAt(path));
+  }
+
+  private constructor(public readonly type: InputType, public readonly value: any) {
+  }
+}
+
+/**
+ * The type of task input
+ */
+export enum InputType {
+  Text,
+  Object
+}
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions/lib/json-path.ts b/packages/@aws-cdk/aws-stepfunctions/lib/json-path.ts
new file mode 100644
index 0000000000000..eb80f01684e3c
--- /dev/null
+++ b/packages/@aws-cdk/aws-stepfunctions/lib/json-path.ts
@@ -0,0 +1,196 @@
+import { Token, TokenMap } from '@aws-cdk/cdk';
+
+const JSON_PATH_TOKEN_SYMBOL = Symbol.for('@aws-cdk/aws-stepfunctions.JsonPathToken');
+
+export class JsonPathToken extends Token {
+  public static isJsonPathToken(x: any): x is JsonPathToken {
+    return (x as any)[JSON_PATH_TOKEN_SYMBOL] === true;
+  }
+
+  constructor(public readonly path: string) {
+    super(() => path, `field${path}`); // Make function to prevent eager evaluation in superclass
+    Object.defineProperty(this, JSON_PATH_TOKEN_SYMBOL, { value: true });
+  }
+}
+
+/**
+ * Deep render a JSON object to expand JSON path fields, updating the key to end in '.$'
+ */
+export function renderObject(obj: object | undefined): object | undefined {
+  return recurseObject(obj, {
+    handleString: renderString,
+    handleList: renderStringList,
+    handleNumber: renderNumber
+  });
+}
+
+/**
+ * Return all JSON paths that are used in the given structure
+ */
+export function findReferencedPaths(obj: object | undefined): Set<string> {
+  const found = new Set<string>();
+
+  recurseObject(obj, {
+    handleString(_key: string, x: string) {
+      const path = jsonPathString(x);
+      if (path !== undefined) { found.add(path); }
+      return {};
+    },
+
+    handleList(_key: string, x: string[]) {
+      const path = jsonPathStringList(x);
+      if (path !== undefined) { found.add(path); }
+      return {};
+    },
+
+    handleNumber(_key: string, x: number) {
+      const path = jsonPathNumber(x);
+      if (path !== undefined) { found.add(path); }
+      return {};
+    }
+  });
+
+  return found;
+}
+
+interface FieldHandlers {
+  handleString(key: string, x: string): {[key: string]: string};
+  handleList(key: string, x: string[]): {[key: string]: string[] | string };
+  handleNumber(key: string, x: number): {[key: string]: number | string};
+}
+
+export function recurseObject(obj: object | undefined, handlers: FieldHandlers): object | undefined {
+  if (obj === undefined) { return undefined; }
+
+  const ret: any = {};
+  for (const [key, value] of Object.entries(obj)) {
+    if (typeof value === 'string') {
+      Object.assign(ret, handlers.handleString(key, value));
+    } else if (typeof value === 'number') {
+      Object.assign(ret, handlers.handleNumber(key, value));
+    } else if (Array.isArray(value)) {
+      Object.assign(ret, recurseArray(key, value, handlers));
+    } else if (value === null || value === undefined) {
+      // Nothing
+    } else if (typeof value === 'object') {
+      ret[key] = recurseObject(value, handlers);
+    }
+  }
+
+  return ret;
+}
+
+/**
+ * Render an array that may or may not contain a string list token
+ */
+function recurseArray(key: string, arr: any[], handlers: FieldHandlers): {[key: string]: any[] | string} {
+  if (isStringArray(arr)) {
+    const path = jsonPathStringList(arr);
+    if (path !== undefined) {
+      return handlers.handleList(key, arr);
+    }
+
+    // Fall through to correctly reject encoded strings inside an array.
+    // They cannot be represented because there is no key to append a '.$' to.
+  }
+
+  return {
+    [key]: arr.map(value => {
+      if ((typeof value === 'string' && jsonPathString(value) !== undefined)
+        || (typeof value === 'number' && jsonPathNumber(value) !== undefined)
+        || (isStringArray(value) && jsonPathStringList(value) !== undefined)) {
+        throw new Error('Cannot use JsonPath fields in an array, they must be used in objects');
+      }
+      if (typeof value === 'object' && value !== null) {
+        return recurseObject(value, handlers);
+      }
+      return value;
+    })
+  };
+}
+
+function isStringArray(x: any): x is string[] {
+  return Array.isArray(x) && x.every(el => typeof el === 'string');
+}
+
+/**
+ * Render a parameter string
+ *
+ * If the string value starts with '$.', render it as a path string, otherwise as a direct string.
+ */
+function renderString(key: string, value: string): {[key: string]: string} {
+  const path = jsonPathString(value);
+  if (path !== undefined) {
+    return { [key + '.$']: path };
+  } else {
+    return { [key]: value };
+  }
+}
+
+/**
+ * Render a parameter string
+ *
+ * If the string value starts with '$.', render it as a path string, otherwise as a direct string.
+ */
+function renderStringList(key: string, value: string[]): {[key: string]: string[] | string} {
+  const path = jsonPathStringList(value);
+  if (path !== undefined) {
+    return { [key + '.$']: path };
+  } else {
+    return { [key]: value };
+  }
+}
+
+/**
+ * Render a parameter string
+ *
+ * If the string value starts with '$.', render it as a path string, otherwise as a direct string.
+ */
+function renderNumber(key: string, value: number): {[key: string]: number | string} {
+  const path = jsonPathNumber(value);
+  if (path !== undefined) {
+    return { [key + '.$']: path };
+  } else {
+    return { [key]: value };
+  }
+}
+
+/**
+ * If the indicated string is an encoded JSON path, return the path
+ *
+ * Otherwise return undefined.
+ */
+function jsonPathString(x: string): string | undefined {
+  const fragments = TokenMap.instance().splitString(x);
+  const jsonPathTokens = fragments.tokens.filter(JsonPathToken.isJsonPathToken);
+
+  if (jsonPathTokens.length > 0 && fragments.length > 1) {
+    throw new Error(`Field references must be the entire string, cannot concatenate them (found '${x}')`);
+  }
+  if (jsonPathTokens.length > 0) {
+    return jsonPathTokens[0].path;
+  }
+  return undefined;
+}
+
+/**
+ * If the indicated string list is an encoded JSON path, return the path
+ *
+ * Otherwise return undefined.
+ */
+function jsonPathStringList(x: string[]): string | undefined {
+  return pathFromToken(TokenMap.instance().lookupList(x));
+}
+
+/**
+ * If the indicated number is an encoded JSON path, return the path
+ *
+ * Otherwise return undefined.
+ */
+function jsonPathNumber(x: number): string | undefined {
+  return pathFromToken(TokenMap.instance().lookupNumberToken(x));
+}
+
+function pathFromToken(token: Token | undefined) {
+  return token && (JsonPathToken.isJsonPathToken(token) ? token.path : undefined);
+}
\ No newline at end of file
diff --git a/packages/@aws-cdk/aws-stepfunctions/test/test.fields.ts b/packages/@aws-cdk/aws-stepfunctions/test/test.fields.ts
new file mode 100644
index 0000000000000..737ad04d43ec3
--- /dev/null
+++ b/packages/@aws-cdk/aws-stepfunctions/test/test.fields.ts
@@ -0,0 +1,135 @@
+import { Test } from 'nodeunit';
+import { Context, Data, FieldUtils } from "../lib";
+
+export = {
+  'deep replace correctly handles fields in arrays'(test: Test) {
+    test.deepEqual(FieldUtils.renderObject({
+      literal: 'literal',
+      field: Data.stringAt('$.stringField'),
+      listField: Data.listAt('$.listField'),
+      deep: [
+        'literal',
+        {
+          deepField: Data.numberAt('$.numField'),
+        }
+      ]
+    }), {
+      'literal': 'literal',
+      'field.$': '$.stringField',
+      'listField.$': '$.listField',
+      'deep': [
+        'literal',
+        {
+          'deepField.$': '$.numField'
+        }
+      ],
+    });
+
+    test.done();
+  },
+
+  'exercise contextpaths'(test: Test) {
+    test.deepEqual(FieldUtils.renderObject({
+      str: Context.stringAt('$$.Execution.StartTime'),
+      count: Context.numberAt('$$.State.RetryCount'),
+      token: Context.taskToken,
+    }), {
+      'str.$': '$$.Execution.StartTime',
+      'count.$': '$$.State.RetryCount',
+      'token.$': '$$.Task.Token'
+    });
+
+    test.done();
+  },
+
+  'find all referenced paths'(test: Test) {
+    test.deepEqual(FieldUtils.findReferencedPaths({
+      literal: 'literal',
+      field: Data.stringAt('$.stringField'),
+      listField: Data.listAt('$.listField'),
+      deep: [
+        'literal',
+        {
+          field: Data.stringAt('$.stringField'),
+          deepField: Data.numberAt('$.numField'),
+        }
+      ]
+    }), [
+      '$.listField',
+      '$.numField',
+      '$.stringField',
+    ]);
+
+    test.done();
+  },
+
+  'cannot have JsonPath fields in arrays'(test: Test) {
+    test.throws(() => {
+      FieldUtils.renderObject({
+        deep: [Data.stringAt('$.hello')]
+      });
+    }, /Cannot use JsonPath fields in an array/);
+
+    test.done();
+  },
+
+  'datafield path must be correct'(test: Test) {
+    test.throws(() => {
+      Data.stringAt('hello');
+    }, /must start with '\$.'/);
+
+    test.done();
+  },
+
+  'context path must be correct'(test: Test) {
+    test.throws(() => {
+      Context.stringAt('hello');
+    }, /must start with '\$\$.'/);
+
+    test.done();
+  },
+
+  'test contains task token'(test: Test) {
+    test.equal(true, FieldUtils.containsTaskToken({
+      field: Context.taskToken
+    }));
+
+    test.equal(true, FieldUtils.containsTaskToken({
+      field: Context.stringAt('$$.Task'),
+    }));
+
+    test.equal(true, FieldUtils.containsTaskToken({
+      field: Context.entireContext
+    }));
+
+    test.equal(false, FieldUtils.containsTaskToken({
+      oops: 'not here'
+    }));
+
+    test.equal(false, FieldUtils.containsTaskToken({
+      oops: Context.stringAt('$$.Execution.StartTime')
+    }));
+
+    test.done();
+  },
+
+  'arbitrary JSONPath fields are not replaced'(test: Test) {
+    test.deepEqual(FieldUtils.renderObject({
+      field: '$.content',
+    }), {
+      field: '$.content'
+    });
+
+    test.done();
+  },
+
+  'fields cannot be used somewhere in a string interpolation'(test: Test) {
+    test.throws(() => {
+      FieldUtils.renderObject({
+        field: `contains ${Data.stringAt('$.hello')}`
+      });
+    }, /Field references must be the entire string/);
+
+    test.done();
+  }
+};
\ No newline at end of file
diff --git a/packages/@aws-cdk/cdk/lib/encoding.ts b/packages/@aws-cdk/cdk/lib/encoding.ts
index c4d072e64df65..90af8112ff41b 100644
--- a/packages/@aws-cdk/cdk/lib/encoding.ts
+++ b/packages/@aws-cdk/cdk/lib/encoding.ts
@@ -24,7 +24,7 @@ export class TokenString {
   /**
    * Returns a `TokenString` for this string.
    */
-  public static forStringToken(s: string) {
+  public static forString(s: string) {
     return new TokenString(s, STRING_TOKEN_REGEX);
   }
 
@@ -106,7 +106,7 @@ export function containsListTokenElement(xs: any[]) {
  */
 export function unresolved(obj: any): boolean {
   if (typeof(obj) === 'string') {
-    return TokenString.forStringToken(obj).test();
+    return TokenString.forString(obj).test();
   } else if (typeof obj === 'number') {
     return extractTokenDouble(obj) !== undefined;
   } else if (Array.isArray(obj) && obj.length === 1) {
diff --git a/packages/@aws-cdk/cdk/lib/resolve.ts b/packages/@aws-cdk/cdk/lib/resolve.ts
index 0f9bf51a10704..01b970935409e 100644
--- a/packages/@aws-cdk/cdk/lib/resolve.ts
+++ b/packages/@aws-cdk/cdk/lib/resolve.ts
@@ -77,7 +77,7 @@ export function resolve(obj: any, options: IResolveOptions): any {
   // string - potentially replace all stringified Tokens
   //
   if (typeof(obj) === 'string') {
-    const str = TokenString.forStringToken(obj);
+    const str = TokenString.forString(obj);
     if (str.test()) {
       const fragments = str.split(tokenMap.lookupToken.bind(tokenMap));
       return options.resolver.resolveString(fragments, makeContext());
diff --git a/packages/@aws-cdk/cdk/lib/string-fragments.ts b/packages/@aws-cdk/cdk/lib/string-fragments.ts
index 9ca8e2057dd96..f33caf2cd94ec 100644
--- a/packages/@aws-cdk/cdk/lib/string-fragments.ts
+++ b/packages/@aws-cdk/cdk/lib/string-fragments.ts
@@ -12,7 +12,7 @@ type IntrinsicFragment = { type: 'intrinsic'; value: any; };
 type Fragment =  LiteralFragment | TokenFragment | IntrinsicFragment;
 
 /**
- * Fragments of a string with markers
+ * Fragments of a concatenated string containing stringified Tokens
  */
 export class TokenizedStringFragments {
   private readonly fragments = new Array<Fragment>();
@@ -43,6 +43,22 @@ export class TokenizedStringFragments {
     this.fragments.push({ type: 'intrinsic', value });
   }
 
+  /**
+   * Return all Tokens from this string
+   */
+  public get tokens(): Token[] {
+    const ret = new Array<Token>();
+    for (const f of this.fragments) {
+      if (f.type === 'token') {
+        ret.push(f.token);
+      }
+    }
+    return ret;
+  }
+
+  /**
+   * Apply a transformation function to all tokens in the string
+   */
   public mapTokens(mapper: ITokenMapper): TokenizedStringFragments {
     const ret = new TokenizedStringFragments();
 
diff --git a/packages/@aws-cdk/cdk/lib/token-map.ts b/packages/@aws-cdk/cdk/lib/token-map.ts
index 9730cf0f063a4..0adb4ef31fb8a 100644
--- a/packages/@aws-cdk/cdk/lib/token-map.ts
+++ b/packages/@aws-cdk/cdk/lib/token-map.ts
@@ -1,5 +1,6 @@
 import { BEGIN_LIST_TOKEN_MARKER, BEGIN_STRING_TOKEN_MARKER, createTokenDouble,
   END_TOKEN_MARKER, extractTokenDouble, TokenString, VALID_KEY_CHARS } from "./encoding";
+import { TokenizedStringFragments } from "./string-fragments";
 import { Token } from "./token";
 
 const glob = global as any;
@@ -73,13 +74,20 @@ export class TokenMap {
     return createTokenDouble(tokenIndex);
   }
 
+  /**
+   * Split a string into literals and Tokens
+   */
+  public splitString(s: string): TokenizedStringFragments {
+    const str = TokenString.forString(s);
+    return str.split(this.lookupToken.bind(this));
+  }
+
   /**
    * Reverse a string representation into a Token object
    */
   public lookupString(s: string): Token | undefined {
-    const str = TokenString.forStringToken(s);
-    const fragments = str.split(this.lookupToken.bind(this));
-    if (fragments.length === 1) {
+    const fragments = this.splitString(s);
+    if (fragments.tokens.length > 0 && fragments.length === 1) {
       return fragments.firstToken;
     }
     return undefined;