1// Alchemy modifications are licensed under Apache-2.0.
2// This file includes third-party code; see /THIRD_PARTY_LICENSES.md.
3import type {
4 ModifierCallback,
5 WorkflowBinding,
6 WorkflowInstanceIntrospector,
7 WorkflowInstanceModifier,
8 WorkflowIntrospectionOperation,
9 WorkflowIntrospectionStreamResult,
10 WorkflowIntrospector,
11 WorkflowStepSelector,
12} from "./types.ts";
13
14function normalizeStreamMockChunk(value: unknown): Uint8Array {
15 if (value instanceof Uint8Array) {
16 return value;
17 }
18
19 if (value instanceof ArrayBuffer) {
20 return new Uint8Array(value);
21 }
22
23 if (ArrayBuffer.isView(value) && !(value instanceof DataView)) {
24 return new Uint8Array(value.buffer, value.byteOffset, value.byteLength);
25 }
26
27 throw new TypeError(
28 "Workflow mockStepResult() ReadableStream chunks must be ArrayBuffer or TypedArray values.",
29 );
30}
31
32async function readStreamMockChunks(
33 stream: ReadableStream<unknown>,
34): Promise<Array<Uint8Array>> {
35 if (stream.locked) {
36 throw new TypeError(
37 "Workflow mockStepResult() received a locked or unreadable ReadableStream.",
38 );
39 }
40
41 const chunks: Array<Uint8Array> = [];
42 const reader = stream.getReader();
43 let fullyRead = false;
44 try {
45 while (true) {
46 const result = await reader.read();
47 if (result.done) {
48 fullyRead = true;
49 return chunks;
50 }
51
52 const chunk = normalizeStreamMockChunk(result.value);
53 if (chunk.byteLength === 0) {
54 continue;
55 }
56
57 chunks.push(chunk.slice());
58 }
59 } finally {
60 if (!fullyRead) {
61 await reader
62 .cancel("stream mock consumption stopped before completion")
63 .catch(() => {});
64 }
65 try {
66 reader.releaseLock();
67 } catch {
68 /** Reader may still be processing cancel(). */
69 }
70 }
71}
72
73export class WorkflowInstanceModificationRecorder implements WorkflowInstanceModifier {
74 constructor(
75 private readonly operations: Array<WorkflowIntrospectionOperation>,
76 ) {}
77
78 async disableSleeps(steps?: Array<WorkflowStepSelector>): Promise<void> {
79 this.operations.push({ type: "disableSleeps", steps });
80 }
81
82 async disableRetryDelays(steps?: Array<WorkflowStepSelector>): Promise<void> {
83 this.operations.push({ type: "disableRetryDelays", steps });
84 }
85
86 async mockStepResult(
87 step: WorkflowStepSelector,
88 stepResult: unknown,
89 ): Promise<void> {
90 if (stepResult instanceof ReadableStream) {
91 const streamResult: WorkflowIntrospectionStreamResult = {
92 __workflowIntrospectionStreamResult: true,
93 chunks: await readStreamMockChunks(stepResult),
94 };
95 this.operations.push({
96 type: "mockStepResult",
97 step,
98 stepResult: streamResult,
99 });
100 return;
101 }
102
103 this.operations.push({ type: "mockStepResult", step, stepResult });
104 }
105
106 async mockStepError(
107 step: WorkflowStepSelector,
108 error: Error,
109 times?: number,
110 ): Promise<void> {
111 this.operations.push({
112 type: "mockStepError",
113 step,
114 error: { name: error.name, message: error.message },
115 times,
116 });
117 }
118
119 async forceStepTimeout(
120 step: WorkflowStepSelector,
121 times?: number,
122 ): Promise<void> {
123 this.operations.push({ type: "forceStepTimeout", step, times });
124 }
125
126 async mockEvent(event: { type: string; payload: unknown }): Promise<void> {
127 this.operations.push({ type: "mockEvent", event });
128 }
129
130 async forceEventTimeout(step: WorkflowStepSelector): Promise<void> {
131 this.operations.push({ type: "forceEventTimeout", step });
132 }
133}
134
135export class WorkflowIntrospectorHandle implements WorkflowIntrospector {
136 #disposed = false;
137 #sessionId: string | undefined;
138 #instanceIntrospectors = new Map<string, WorkflowInstanceIntrospector>();
139 #operations: Array<WorkflowIntrospectionOperation> = [];
140
141 constructor(private readonly workflow: WorkflowBinding) {}
142
143 async start(): Promise<void> {
144 this.#sessionId = await this.workflow.unsafeStartIntrospection();
145 }
146
147 private getSessionId(): string {
148 if (this.#sessionId === undefined) {
149 throw new Error("Workflow introspection has not started.");
150 }
151 return this.#sessionId;
152 }
153
154 async modifyAll(fn: ModifierCallback): Promise<void> {
155 const sessionId = this.getSessionId();
156 await fn(new WorkflowInstanceModificationRecorder(this.#operations));
157 await this.workflow.unsafeSetIntrospectionOperations(
158 sessionId,
159 this.#operations,
160 );
161 }
162
163 async get(): Promise<Array<WorkflowInstanceIntrospector>> {
164 const sessionId = this.getSessionId();
165 await this.syncInstanceIntrospectors(sessionId);
166
167 return Array.from(this.#instanceIntrospectors.values());
168 }
169
170 private async syncInstanceIntrospectors(sessionId: string): Promise<void> {
171 const instanceIds =
172 await this.workflow.unsafeGetIntrospectionInstances(sessionId);
173
174 for (const instanceId of instanceIds) {
175 if (!this.#instanceIntrospectors.has(instanceId)) {
176 this.#instanceIntrospectors.set(
177 instanceId,
178 new WorkflowInstanceIntrospectorHandle(this.workflow, instanceId),
179 );
180 }
181 }
182 }
183
184 private async stopIntrospectionSession(sessionId: string): Promise<void> {
185 try {
186 await this.syncInstanceIntrospectors(sessionId);
187 } finally {
188 await this.workflow.unsafeStopIntrospection(sessionId);
189 }
190 }
191
192 private async disposeInstanceIntrospectors(): Promise<void> {
193 try {
194 await Promise.all(
195 Array.from(this.#instanceIntrospectors.values(), (introspector) =>
196 introspector.dispose(),
197 ),
198 );
199 } finally {
200 this.#instanceIntrospectors.clear();
201 }
202 }
203
204 /** Keep this bound; explicit resource management may call the disposer unbound. */
205 dispose = async (): Promise<void> => {
206 if (this.#disposed) {
207 return;
208 }
209 this.#disposed = true;
210 const sessionId = this.#sessionId;
211
212 try {
213 if (sessionId !== undefined) {
214 await this.stopIntrospectionSession(sessionId);
215 }
216 } finally {
217 await this.disposeInstanceIntrospectors();
218 }
219 };
220
221 async [Symbol.asyncDispose](): Promise<void> {
222 await this.dispose();
223 }
224}
225
226export class WorkflowInstanceIntrospectorHandle implements WorkflowInstanceIntrospector {
227 #instanceModifier: WorkflowInstanceModifier | undefined;
228 #instanceModifierPromise: Promise<WorkflowInstanceModifier> | undefined;
229
230 constructor(
231 private readonly workflow: WorkflowBinding,
232 private readonly instanceId: string,
233 ) {
234 this.#instanceModifierPromise = workflow
235 .unsafeGetInstanceModifier(instanceId)
236 .then((modifier) => {
237 this.#instanceModifier = modifier as WorkflowInstanceModifier;
238 this.#instanceModifierPromise = undefined;
239 return this.#instanceModifier;
240 });
241
242 // To avoid an unhandled rejection when the handle is used without modify()
243 void this.#instanceModifierPromise.catch(() => {});
244 }
245
246 async modify(fn: ModifierCallback): Promise<WorkflowInstanceIntrospector> {
247 if (this.#instanceModifierPromise !== undefined) {
248 this.#instanceModifier = await this.#instanceModifierPromise;
249 }
250 if (this.#instanceModifier === undefined) {
251 throw new Error(
252 "could not apply modifications due to internal error. Retrying the test may resolve the issue.",
253 );
254 }
255
256 await fn(this.#instanceModifier);
257
258 return this;
259 }
260
261 async waitForStepResult(step: WorkflowStepSelector): Promise<unknown> {
262 return await this.workflow.unsafeWaitForStepResult(
263 this.instanceId,
264 step.name,
265 step.index,
266 );
267 }
268
269 async waitForStatus(status: string): Promise<void> {
270 if (status === "queued") {
271 return;
272 }
273
274 await this.workflow.unsafeWaitForStatus(this.instanceId, status);
275 }
276
277 async getOutput(): Promise<unknown> {
278 return await this.workflow.unsafeGetOutputOrError(this.instanceId, true);
279 }
280
281 async getError(): Promise<{ name: string; message: string }> {
282 return (await this.workflow.unsafeGetOutputOrError(
283 this.instanceId,
284 false,
285 )) as {
286 name: string;
287 message: string;
288 };
289 }
290
291 /** Keep this bound; explicit resource management may call the disposer unbound. */
292 dispose = async (): Promise<void> => {
293 await this.workflow.unsafeAbort(this.instanceId, "Instance dispose");
294 };
295
296 async [Symbol.asyncDispose](): Promise<void> {
297 await this.dispose();
298 }
299}