1/**
2 * Local interpreter for Step Functions programs — the payoff of the AST.
3 *
4 * `simulate` runs the *same* `SfnEffect` a `StateMachine.fromProgram`
5 * deploys, in-process, as a real `Effect`: Task states dispatch to
6 * caller-provided handler stubs, Choice conditions evaluate against the
7 * live variable environment, `catchTag`/`retry` follow the exact ASL error
8 * semantics (`{ Error, Cause }` error output, `States.ALL` matching), and
9 * `Wait` states are skipped (simulations are instantaneous).
10 */
11import * as Data from "effect/Data";
12import * as Effect from "effect/Effect";
13import * as Result from "effect/Result";
14import {
15 inputExpr,
16 isExpr,
17 nodeOf,
18 variableExpr,
19 type ExprNode,
20} from "./Jsonata.ts";
21import type { AslNode } from "./Node.ts";
22import { isSfnEffect, type SfnEffect } from "./Program.ts";
23
24/**
25 * The simulation itself went wrong — a missing handler, an unsupported raw
26 * JSONata expression, or a structurally invalid program. Distinct from the
27 * program's own typed failures, which flow through the `E` channel.
28 */
29export class SimulateError extends Data.TaggedError("SimulateError")<{
30 readonly message: string;
31}> {}
32
33/**
34 * Handler stubs for the program's Task states, keyed by:
35 * - the function's `LogicalId` for `Sfn.invoke` states,
36 * - the integration resource ARN (including any `.waitForTaskToken`
37 * suffix) for `Sfn.integrate` / `Sfn.waitForTaskToken` states.
38 */
39export interface SimulateHandlers {
40 readonly [key: string]: (payload: unknown) => Effect.Effect<any, any>;
41}
42
43export interface SimulateOptions {
44 /** Task handler stubs. @default {} */
45 readonly handlers?: SimulateHandlers;
46 /** The task token exposed as `Sfn.taskToken`. @default "SIMULATED_TASK_TOKEN" */
47 readonly taskToken?: string;
48}
49
50interface Env {
51 readonly input: unknown;
52 /** Variable scope — prototype chain mirrors ASL variable scoping. */
53 vars: Record<string, unknown>;
54 readonly taskToken: string;
55 readonly handlers: SimulateHandlers;
56 counter: { n: number };
57}
58
59/**
60 * Internal failure wrapper preserving the exact ASL `Error`/`Cause` pair a
61 * `Fail` state declares, so `catch` handlers observe the same
62 * `$states.errorOutput` the cloud would produce. Unwrapped back to the
63 * user's typed failure at the simulation boundary.
64 */
65class FailStateFailure {
66 readonly _sim = "FailStateFailure";
67 constructor(
68 readonly error: string,
69 readonly cause: string,
70 readonly failure: unknown,
71 ) {}
72}
73
74const isFailStateFailure = (value: unknown): value is FailStateFailure =>
75 value instanceof FailStateFailure;
76
77/** The ASL error name of a simulated failure (`_tag` convention). */
78const errorNameOf = (failure: unknown): string => {
79 if (isFailStateFailure(failure)) return failure.error;
80 if (
81 typeof failure === "object" &&
82 failure !== null &&
83 "_tag" in failure &&
84 typeof (failure as { _tag: unknown })._tag === "string"
85 ) {
86 return (failure as { _tag: string })._tag;
87 }
88 return "States.TaskFailed";
89};
90
91const causeOf = (failure: unknown): string => {
92 if (isFailStateFailure(failure)) return failure.cause;
93 if (typeof failure === "object" && failure !== null) {
94 const message = (failure as { message?: unknown }).message;
95 if (typeof message === "string" && message.length > 0) return message;
96 try {
97 return JSON.stringify(failure);
98 } catch {
99 return errorNameOf(failure);
100 }
101 }
102 return String(failure);
103};
104
105const evalExprNode = (node: ExprNode, env: Env): unknown => {
106 switch (node._) {
107 case "root":
108 switch (node.root.kind) {
109 case "input":
110 return env.input;
111 case "variable":
112 return env.vars[node.root.name];
113 case "token":
114 return env.taskToken;
115 }
116 // eslint-disable-next-line no-fallthrough
117 case "prop": {
118 const parent = evalExprNode(node.parent, env) as
119 | Record<string, unknown>
120 | null
121 | undefined;
122 return parent == null ? undefined : parent[node.name];
123 }
124 case "literal":
125 return node.value;
126 case "binop": {
127 const left = evalExprNode(node.left, env) as any;
128 const right = evalExprNode(node.right, env) as any;
129 switch (node.op) {
130 case "=":
131 return left === right;
132 case "!=":
133 return left !== right;
134 case ">":
135 return left > right;
136 case ">=":
137 return left >= right;
138 case "<":
139 return left < right;
140 case "<=":
141 return left <= right;
142 case "and":
143 return Boolean(left) && Boolean(right);
144 case "or":
145 return Boolean(left) || Boolean(right);
146 }
147 }
148 // eslint-disable-next-line no-fallthrough
149 case "not":
150 return !evalExprNode(node.inner, env);
151 case "raw":
152 throw new SimulateError({
153 message: `raw JSONata expressions are not supported by simulate: ${node.jsonata}`,
154 });
155 }
156};
157
158/** Deep-evaluate a payload value: `Expr` references resolve against `env`. */
159const evalValue = (value: unknown, env: Env): unknown => {
160 if (value === undefined) return null;
161 if (isExpr(value)) return evalExprNode(nodeOf(value), env);
162 if (Array.isArray(value)) return value.map((entry) => evalValue(entry, env));
163 if (typeof value === "object" && value !== null) {
164 return Object.fromEntries(
165 Object.entries(value).map(([key, entry]) => [key, evalValue(entry, env)]),
166 );
167 }
168 return value;
169};
170
171const evalExpr = (value: unknown, env: Env) =>
172 Effect.try({
173 try: () => evalValue(value, env),
174 catch: (error) =>
175 error instanceof SimulateError
176 ? error
177 : new SimulateError({ message: String(error) }),
178 });
179
180const allocVar = (env: Env): string => `sim${env.counter.n++}`;
181
182/** Child scope with an ASL scope barrier (Parallel branch / Map iteration). */
183const childEnv = (env: Env): Env => ({
184 ...env,
185 vars: Object.create(env.vars) as Record<string, unknown>,
186});
187
188const runNode = (node: AslNode, env: Env): Effect.Effect<unknown, any> => {
189 switch (node.kind) {
190 case "invoke": {
191 const handler = env.handlers[node.fn.LogicalId];
192 if (handler === undefined) {
193 return Effect.fail(
194 new SimulateError({
195 message: `no simulate handler for function "${node.fn.LogicalId}"`,
196 }),
197 );
198 }
199 return evalExpr(node.payload, env).pipe(Effect.flatMap(handler));
200 }
201
202 case "integrate": {
203 const key = node.waitForTaskToken
204 ? `${node.options.resource}.waitForTaskToken`
205 : node.options.resource;
206 const handler = env.handlers[key];
207 if (handler === undefined) {
208 return Effect.fail(
209 new SimulateError({
210 message: `no simulate handler for integration "${key}"`,
211 }),
212 );
213 }
214 return evalExpr(node.options.arguments, env).pipe(
215 Effect.flatMap(handler),
216 );
217 }
218
219 case "all":
220 return Effect.forEach(node.branches, (branch) =>
221 runNode(branch.node, childEnv(env)),
222 );
223
224 case "forEach":
225 return evalExpr(node.items, env).pipe(
226 Effect.flatMap((items) => {
227 if (!Array.isArray(items)) {
228 return Effect.fail(
229 new SimulateError({
230 message: "Sfn.forEach items did not evaluate to an array",
231 }),
232 );
233 }
234 return Effect.forEach(items, (item) => {
235 // mirror compile: the item is bound to a fresh variable scoped
236 // to the iteration
237 const scope = childEnv(env);
238 const itemVar = allocVar(scope);
239 scope.vars[itemVar] = item;
240 const body = node.body(variableExpr(itemVar));
241 if (!isSfnEffect(body)) {
242 return Effect.fail(
243 new SimulateError({
244 message: "Sfn.forEach body must return an Sfn effect",
245 }),
246 );
247 }
248 return runNode(body.node, scope);
249 });
250 }),
251 );
252
253 case "sleep":
254 // simulations are instantaneous — Wait states are skipped
255 return Effect.succeed(null);
256
257 case "when":
258 return evalExpr(node.condition, env).pipe(
259 Effect.flatMap((condition) =>
260 condition
261 ? runNode(node.onTrue.node, env)
262 : node.onFalse !== undefined
263 ? runNode(node.onFalse.node, env)
264 : Effect.succeed(null),
265 ),
266 );
267
268 case "match":
269 return evalExpr(node.value, env).pipe(
270 Effect.flatMap((value) => {
271 const program = node.cases[String(value)] ?? node.otherwise;
272 if (program === undefined) {
273 return Effect.fail(
274 new SimulateError({
275 message: `States.NoChoiceMatched: no case for ${JSON.stringify(value)}`,
276 }),
277 );
278 }
279 return runNode(program.node, env);
280 }),
281 );
282
283 case "succeed":
284 return evalExpr(node.value, env);
285
286 case "fail":
287 return Effect.fail(
288 new FailStateFailure(node.error, node.cause, node.failure),
289 );
290
291 case "retry": {
292 const tags =
293 node.options.while === undefined
294 ? ["States.ALL"]
295 : typeof node.options.while === "string"
296 ? [node.options.while]
297 : [...node.options.while];
298 const maxAttempts = node.options.maxAttempts ?? 3;
299 const attempt = (retriesLeft: number): Effect.Effect<unknown, any> =>
300 Effect.result(runNode(node.inner.node, env)).pipe(
301 Effect.flatMap((result) => {
302 if (Result.isSuccess(result)) return Effect.succeed(result.success);
303 const failure = result.failure;
304 const retryable =
305 !(failure instanceof SimulateError) &&
306 (tags.includes("States.ALL") ||
307 tags.includes(errorNameOf(failure)));
308 return retryable && retriesLeft > 0
309 ? attempt(retriesLeft - 1)
310 : Effect.fail(failure);
311 }),
312 );
313 return attempt(maxAttempts);
314 }
315
316 case "catch":
317 return Effect.result(runNode(node.inner.node, env)).pipe(
318 Effect.flatMap((result) => {
319 if (Result.isSuccess(result)) return Effect.succeed(result.success);
320 const failure = result.failure;
321 const matches =
322 !(failure instanceof SimulateError) &&
323 (node.tags.includes("States.ALL") ||
324 node.tags.includes(errorNameOf(failure)));
325 if (!matches) return Effect.fail(failure);
326 const errorVar = allocVar(env);
327 env.vars[errorVar] = {
328 Error: errorNameOf(failure),
329 Cause: causeOf(failure),
330 };
331 const handler = node.handler(variableExpr(errorVar));
332 if (!isSfnEffect(handler)) {
333 return Effect.fail(
334 new SimulateError({
335 message:
336 "Sfn.catchTag/catchAll handler must return an Sfn effect",
337 }),
338 );
339 }
340 return runNode(handler.node, env);
341 }),
342 );
343
344 case "gen":
345 return Effect.gen(function* () {
346 const iterator = node.body(inputExpr());
347 let step = iterator.next();
348 while (!step.done) {
349 const program: unknown = step.value;
350 if (!isSfnEffect(program)) {
351 return yield* Effect.fail(
352 new SimulateError({
353 message: "Sfn.gen may only yield Sfn effects",
354 }),
355 );
356 }
357 const value = yield* runNode(program.node, env);
358 const name = allocVar(env);
359 env.vars[name] = value;
360 step = iterator.next(variableExpr(name));
361 }
362 return yield* evalExpr(step.value, env);
363 });
364 }
365};
366
367/**
368 * Run a Step Functions program in-process as an `Effect` — same AST the
369 * cloud executes, interpreted locally for unit tests.
370 *
371 * ```typescript
372 * const result = yield* simulate(program, { value: 6, items: [1, 2, 3] }, {
373 * handlers: {
374 * Doubler: (payload) =>
375 * Effect.succeed({ doubled: (payload as { value: number }).value * 2 }),
376 * },
377 * });
378 * ```
379 */
380export const simulate = <A, E>(
381 program: SfnEffect<A, E>,
382 input: unknown,
383 options?: SimulateOptions,
384): Effect.Effect<A, E | SimulateError> => {
385 const env: Env = {
386 input,
387 vars: Object.create(null) as Record<string, unknown>,
388 taskToken: options?.taskToken ?? "SIMULATED_TASK_TOKEN",
389 handlers: options?.handlers ?? {},
390 counter: { n: 0 },
391 };
392 return runNode(program.node, env).pipe(
393 Effect.mapError((failure) =>
394 isFailStateFailure(failure) ? failure.failure : failure,
395 ),
396 ) as Effect.Effect<A, E | SimulateError>;
397};