1// Alchemy modifications are licensed under Apache-2.0.
2// This file includes third-party code; see /THIRD_PARTY_LICENSES.md.
3import { RpcTarget } from "cloudflare:workers";
4import { computeHash } from "./lib/cache.ts";
5import { isReadableStreamLike, writeStreamOutput } from "./lib/streams.ts";
6import type { Event } from "./context.ts";
7import type { Engine } from "./engine.ts";
8
9export type StepSelector = {
10 name: string;
11 index?: number;
12};
13
14type UserEvent = {
15 type: string;
16 payload: unknown;
17};
18
19// KV key prefixes/values used by the modifier/mock system
20export const MODIFIER_KEYS = {
21 REPLACE_RESULT: "replace-result-",
22 MOCK_STEP_ERROR: "mock-step-error-",
23 MOCK_EVENT: "mock-event-",
24 FORCE_STEP_TIMEOUT: "force-step-timeout-",
25 FORCE_EVENT_TIMEOUT: "force-event-timeout-",
26 FAILURE_INDEX: "failure-index-",
27 DISABLE_SLEEP: "disable-sleep-",
28 DISABLE_ALL_SLEEPS: "disableAllSleeps",
29 DISABLE_RETRY_DELAY: "disable-retry-delay-",
30 DISABLE_ALL_RETRY_DELAYS: "disableAllRetryDelays",
31} as const;
32
33export function isModifierKey(key: string): boolean {
34 return Object.values(MODIFIER_KEYS).some((v) => key.startsWith(v));
35}
36
37export class WorkflowInstanceModifier extends RpcTarget {
38 #engine: Engine;
39 #state: DurableObjectState;
40
41 constructor(engine: Engine, state: DurableObjectState) {
42 super();
43 this.#engine = engine;
44 this.#state = state;
45 }
46
47 async #getWaitForEventCacheKey(step: StepSelector): Promise<string> {
48 let count = 1;
49 if (step.index) {
50 count = step.index;
51 }
52 const name = `${step.name}-${count}`;
53 const hash = await computeHash(name);
54 const cacheKey = `${hash}-${count}`;
55 const waitForEventKey = `${cacheKey}-value`;
56
57 return waitForEventKey;
58 }
59
60 async #getBaseCacheKey(step: StepSelector): Promise<string> {
61 const hash = await computeHash(step.name);
62 let count = 1;
63 if (step.index) {
64 count = step.index;
65 }
66 return `${hash}-${count}`;
67 }
68
69 async #getStepCacheKey(step: StepSelector): Promise<string> {
70 const baseCacheKey = await this.#getBaseCacheKey(step);
71 return `${baseCacheKey}-value`;
72 }
73
74 #getAndIncrementCounter = async (valueKey: string, by: number) => {
75 const counterKey = `${MODIFIER_KEYS.FAILURE_INDEX}${valueKey}`;
76 const next = (await this.#state.storage.get<number>(counterKey)) ?? 1;
77 await this.#state.storage.put(counterKey, next + by);
78 return next;
79 };
80
81 async #getSleepStepDisableKey(step: StepSelector): Promise<string> {
82 let count = 1;
83 if (step.index) {
84 count = step.index;
85 }
86 const sleepNameCountHash = await computeHash(step.name + count);
87
88 return `${MODIFIER_KEYS.DISABLE_SLEEP}${sleepNameCountHash}`;
89 }
90
91 async disableSleeps(steps?: Array<StepSelector>): Promise<void> {
92 if (!steps) {
93 await this.#state.storage.put(MODIFIER_KEYS.DISABLE_ALL_SLEEPS, true);
94 } else {
95 for (const step of steps) {
96 const sleepDisableKey = await this.#getSleepStepDisableKey(step);
97 await this.#state.storage.put(sleepDisableKey, true);
98 }
99 }
100 }
101
102 async disableRetryDelays(steps?: Array<StepSelector>): Promise<void> {
103 if (!steps) {
104 await this.#state.storage.put(
105 MODIFIER_KEYS.DISABLE_ALL_RETRY_DELAYS,
106 true,
107 );
108 } else {
109 for (const step of steps) {
110 const valueKey = await this.#getStepCacheKey(step);
111 await this.#state.storage.put(
112 `${MODIFIER_KEYS.DISABLE_RETRY_DELAY}${valueKey}`,
113 true,
114 );
115 }
116 }
117 }
118
119 // step.do() flow: It first checks if a result or error is already in the cache and, if so, returns it immediately.
120 // If nothing is in the cache, it checks for remaining attempts and runs the user's code against the defined timeout.
121 // Since `step.do()` performs this initial cache check, directly changing the `valueKey` would cause it to
122 // assume the value was pre-cached, preventing it from writing any logs about the step's execution state.
123 // Storing the value under a separate key is crucial because it ensures all execution logs for the step are
124 // generated, rather than the step being skipped due to a premature cache hit.
125 async mockStepResult(step: StepSelector, stepResult: unknown): Promise<void> {
126 const valueKey = await this.#getStepCacheKey(step);
127
128 if (
129 await this.#state.storage.get(
130 `${MODIFIER_KEYS.REPLACE_RESULT}${valueKey}`,
131 )
132 ) {
133 throw new Error(
134 `[WorkflowIntrospector] Trying to mock step '${step.name}' multiple times!`,
135 );
136 }
137
138 if (isReadableStreamLike(stepResult)) {
139 // ReadableStream is not structured-cloneable, so we consume it eagerly
140 // and store the chunks via the streaming infrastructure. We use attempt 0
141 // to distinguish mock chunks from real execution attempts (which start at 1).
142 const baseCacheKey = await this.#getBaseCacheKey(step);
143 const streamMeta = await writeStreamOutput({
144 storage: this.#state.storage,
145 cacheKey: baseCacheKey,
146 attempt: 0,
147 stream: stepResult,
148 skipMetaWrite: true,
149 });
150
151 await this.#state.storage.put(
152 `${MODIFIER_KEYS.REPLACE_RESULT}${valueKey}`,
153 {
154 __mockStreamOutput: true,
155 cacheKey: baseCacheKey,
156 meta: streamMeta,
157 },
158 );
159 } else {
160 await this.#state.storage.put(
161 `${MODIFIER_KEYS.REPLACE_RESULT}${valueKey}`,
162 stepResult,
163 );
164 }
165 }
166
167 // Same logic of `mockStepResult` but stores an error instead of a value.
168 async mockStepError(
169 step: StepSelector,
170 error: Error,
171 times?: number,
172 ): Promise<void> {
173 const valueKey = await this.#getStepCacheKey(step);
174 const serializableError = {
175 name: error.name,
176 message: error.message,
177 };
178
179 if (
180 await this.#state.storage.get(
181 `${MODIFIER_KEYS.REPLACE_RESULT}${valueKey}`,
182 )
183 ) {
184 throw new Error(
185 `[WorkflowIntrospector] Trying to mock error on step '${step.name}' after mocking its result!`,
186 );
187 }
188
189 if (times) {
190 const start = await this.#getAndIncrementCounter(valueKey, times);
191 const mockErrorsPuts = Array.from({ length: times }, (_, i) => {
192 const attempt = start + i;
193 const mockErrorKey = `${MODIFIER_KEYS.MOCK_STEP_ERROR}${valueKey}-${attempt}`;
194 return this.#state.storage.put(mockErrorKey, serializableError);
195 });
196
197 await Promise.all(mockErrorsPuts);
198 } else {
199 const mockErrorKey = `${MODIFIER_KEYS.MOCK_STEP_ERROR}${valueKey}`;
200 await this.#state.storage.put(mockErrorKey, serializableError);
201 }
202 }
203
204 async forceStepTimeout(step: StepSelector, times?: number) {
205 const valueKey = await this.#getStepCacheKey(step);
206
207 if (
208 await this.#state.storage.get(
209 `${MODIFIER_KEYS.REPLACE_RESULT}${valueKey}`,
210 )
211 ) {
212 throw new Error(
213 `[WorkflowIntrospector] Trying to force timeout on step '${step.name}' after mocking its result!`,
214 );
215 }
216
217 if (times) {
218 const start = await this.#getAndIncrementCounter(valueKey, times);
219 const forceTimeouts = Array.from({ length: times }, (_, i) => {
220 const attempt = start + i;
221 const forceStepTimeoutKey = `${MODIFIER_KEYS.FORCE_STEP_TIMEOUT}${valueKey}-${attempt}`;
222 return this.#state.storage.put(forceStepTimeoutKey, true);
223 });
224
225 await Promise.all(forceTimeouts);
226 } else {
227 const forceStepTimeoutKey = `${MODIFIER_KEYS.FORCE_STEP_TIMEOUT}${valueKey}`;
228 await this.#state.storage.put(forceStepTimeoutKey, true);
229 }
230 }
231
232 async mockEvent(event: UserEvent): Promise<void> {
233 const myEvent: Event = {
234 timestamp: new Date(),
235 payload: event.payload,
236 type: event.type,
237 };
238
239 await this.#state.storage.put(
240 `${MODIFIER_KEYS.MOCK_EVENT}${event.type}`,
241 true,
242 );
243 await this.#engine.receiveEvent(myEvent);
244 }
245
246 async forceEventTimeout(step: StepSelector): Promise<void> {
247 const waitForEventKey = await this.#getWaitForEventCacheKey(step);
248 await this.#state.storage.put(
249 `${MODIFIER_KEYS.FORCE_EVENT_TIMEOUT}${waitForEventKey}`,
250 true,
251 );
252 }
253}