Add introspectWorkflow API with proxies

Author: pombosilva

packages/miniflare/src/plugins/workflows/index.ts

@@ -82,7 +82,7 @@ export const WORKFLOWS_PLUGIN: Plugin<
 
 		// this creates one miniflare service per workflow that the user's script has. we should dedupe engine definition later
 		const services = Object.entries(options.workflows ?? {}).map<Service>(
-			([_bindingName, workflow]) => {
+			([bindingName, workflow]) => {
 				// NOTE(lduarte): the engine unique namespace key must be unique per workflow definition
 				// otherwise workerd will crash because there's two equal DO namespaces
 				const uniqueKey = `miniflare-workflows-${workflow.name}`;
@@ -106,7 +106,9 @@ export const WORKFLOWS_PLUGIN: Plugin<
 								className: "Engine",
 								enableSql: true,
 								uniqueKey,
-								preventEviction: true,
+								// Note(osilva): the engine should not be prevented from eviction if we wish
+								// to abort it during tests (vitest-pool-workers)
+								// preventEviction: true,
 							},
 						],
 						durableObjectStorage: {
@@ -128,6 +130,10 @@ export const WORKFLOWS_PLUGIN: Plugin<
 								name: "WORKFLOW_NAME",
 								json: JSON.stringify(workflow.name),
 							},
+							{
+								name: "BINDING_NAME",
+								json: JSON.stringify(bindingName),
+							},
 						],
 					},
 				};

packages/vitest-pool-workers/src/worker/lib/cloudflare/test.ts

@@ -19,4 +19,5 @@ export {
 	applyD1Migrations,
 	createPagesEventContext,
 	introspectWorkflowInstance,
+	introspectWorkflow,
 } from "cloudflare:test-internal";

packages/vitest-pool-workers/src/worker/workflows.ts

@@ -2,98 +2,86 @@ import {
 	instanceStatusName,
 	InstanceStatus as InstanceStatusNumber,
 } from "@cloudflare/workflows-shared/src/instance";
-import { WorkflowIntrospectorError } from "@cloudflare/workflows-shared/src/modifier";
 import { WORKFLOW_ENGINE_BINDING } from "../shared/workflows";
-import { internalEnv } from "./env";
-import type { InstanceModifier } from "@cloudflare/workflows-shared/src/modifier";
+import { runInRunnerObject } from "./durable-objects";
+import { env, internalEnv } from "./env";
+import type {
+	StepSelector,
+	WorkflowInstanceModifier,
+} from "@cloudflare/workflows-shared/src/modifier";
 
-export type StepSelector = {
-	name: string;
-	index?: number;
-};
+type ModifierCallback = (m: WorkflowInstanceModifier) => Promise<void>;
 
-export type WorkflowInstanceIntrospector = {
-	modify(
-		fn: (m: InstanceModifier) => Promise<void>
-	): Promise<WorkflowInstanceIntrospector>;
+// See public facing `cloudflare:test` types for docs
+export interface WorkflowInstanceIntrospector {
+	modify(fn: ModifierCallback): Promise<WorkflowInstanceIntrospector>;
 
 	waitForStepResult(step: StepSelector): Promise<unknown>;
 
 	waitForStatus(status: string): Promise<void>;
 
 	cleanUp(): Promise<void>;
-};
+}
 
-/**
- * Entry point that targets a single Workflow instance
- * This would allow to apply test rules and mocks to the given instance
- */
 export async function introspectWorkflowInstance(
 	workflow: Workflow,
 	instanceId: string
 ): Promise<WorkflowInstanceIntrospector> {
 	if (!workflow || !instanceId) {
-		throw new Error("Workflow binding and instance id are required.");
+		throw new Error(
+			"[WorkflowIntrospector] Workflow binding and instance id are required."
+		);
 	}
 
-	console.log(
-		`[Vitest-Workflows] Introspecting workflow instance: ${instanceId}`
-	);
-
 	// @ts-expect-error getWorkflowName() not exposed
 	const engineBindingName = `${WORKFLOW_ENGINE_BINDING}${(await workflow.getWorkflowName()).toUpperCase()}`;
+
 	// @ts-expect-error DO binding created in runner worker start
 	const engineStubId = internalEnv[engineBindingName].idFromName(instanceId);
 	// @ts-expect-error DO binding created in runner worker start
 	const engineStub = internalEnv[engineBindingName].get(engineStubId);
 
-	const instanceModifier = await engineStub.getInstanceModifier();
+	const instanceModifier = engineStub.getInstanceModifier();
 
 	return new WorkflowInstanceIntrospectorHandle(engineStub, instanceModifier);
 }
 
 class WorkflowInstanceIntrospectorHandle
 	implements WorkflowInstanceIntrospector
 {
-	engineStub: DurableObjectStub;
-	instanceModifier: InstanceModifier;
+	#engineStub: DurableObjectStub;
+	#instanceModifier: WorkflowInstanceModifier;
 	constructor(
 		engineStub: DurableObjectStub,
-		instanceModifier: InstanceModifier
+		instanceModifier: WorkflowInstanceModifier
 	) {
-		this.engineStub = engineStub;
-		this.instanceModifier = instanceModifier;
+		this.#engineStub = engineStub;
+		this.#instanceModifier = instanceModifier;
 	}
 
-	async modify(
-		fn: (m: InstanceModifier) => Promise<void>
-	): Promise<WorkflowInstanceIntrospector> {
-		console.log("[Vitest-Workflows] I should go call a modifier");
-		await fn(this.instanceModifier);
+	async modify(fn: ModifierCallback): Promise<WorkflowInstanceIntrospector> {
+		await fn(this.#instanceModifier);
+
 		return this;
 	}
 
 	async waitForStepResult(step: StepSelector): Promise<unknown> {
-		console.log("waiting for step result of step", step.name);
-		// @ts-expect-error waitForStepResult not exposed
-		const stepResult = await this.engineStub.waitForStepResult(
+		// @ts-expect-error DO binding created in runner worker start
+		const stepResult = await this.#engineStub.waitForStepResult(
 			step.name,
 			step.index
 		);
 
-		console.log("result of step", step.name, "awaited");
 		return stepResult;
 	}
 
 	async waitForStatus(status: InstanceStatus["status"]): Promise<void> {
-		console.log("[Vitest-Workflows] waiting for status");
-
 		if (
 			status === instanceStatusName(InstanceStatusNumber.Terminated) ||
 			status === instanceStatusName(InstanceStatusNumber.Paused)
 		) {
-			throw new WorkflowIntrospectorError(
-				`InstanceStatus '${status}' is not implemented yet and cannot be waited.`
+			throw new Error(
+				`[WorkflowIntrospector] InstanceStatus '${status}' is not implemented yet and cannot be waited.`
 			);
 		}
 
@@ -102,19 +90,165 @@ class WorkflowInstanceIntrospectorHandle
 			// starts running, so waiting for it to be queued should always return
 			return;
 		}
-		// @ts-expect-error waitForStatus not exposed
-		await this.engineStub.waitForStatus(status);
-
-		console.log("[Vitest-Workflows] status awaited");
+		// @ts-expect-error DO binding created in runner worker start
+		await this.#engineStub.waitForStatus(status);
 	}
 
 	async cleanUp(): Promise<void> {
 		// this cleans state with isolatedStorage = false
 		try {
 			// @ts-expect-error DO binding created in runner worker start
-			await this.engineStub.unsafeAbort("user called delete");
+			await this.#engineStub._unsafeAbort("Instance clean up");
 		} catch {
 			// do nothing because we want to clean up this instance
 		}
 	}
 }
+
+// See public facing `cloudflare:test` types for docs
+export interface WorkflowIntrospector {
+	modifyAll(fn: ModifierCallback): void;
+
+	get(): WorkflowInstanceIntrospector[];
+
+	cleanUp(): void;
+}
+
+export async function introspectWorkflow(
+	workflow: Workflow
+): Promise<WorkflowIntrospectorHandle> {
+	if (!workflow) {
+		throw new Error("[WorkflowIntrospector] Workflow binding is required.");
+	}
+
+	const modifierCallbacks: ModifierCallback[] = [];
+	const instanceIntrospectors: WorkflowInstanceIntrospector[] = [];
+	// @ts-expect-error getBindingName not exposed
+	const bindingName = await workflow.getBindingName();
+	const internalOriginalWorkflow = internalEnv[bindingName] as Workflow;
+	const externalOriginalWorkflow = env[bindingName] as Workflow;
+
+	const introspectAndModifyInstance = async (instanceId: string) => {
+		try {
+			await runInRunnerObject(internalEnv, async () => {
+				const introspector = await introspectWorkflowInstance(
+					workflow,
+					instanceId
+				);
+				instanceIntrospectors.push(introspector);
+				// Apply any stored modifier functions
+				for (const callback of modifierCallbacks) {
+					await introspector.modify(callback);
+				}
+			});
+		} catch (error) {
+			console.error(
+				`[WorkflowIntrospector] Error during introspection for instance ${instanceId}:`,
+				error
+			);
+			throw new Error(
+				`[WorkflowIntrospector] Failed to introspect Workflow instance ${instanceId}.`
+			);
+		}
+	};
+
+	const createWorkflowProxyGetHandler = <
+		T extends Workflow,
+	>(): ProxyHandler<T>["get"] => {
+		return (target, property) => {
+			if (property === "create") {
+				return new Proxy(target[property], {
+					async apply(func, thisArg, argArray) {
+						const hasId = Object.hasOwn(argArray[0] ?? {}, "id");
+						if (!hasId) {
+							argArray = [{ id: crypto.randomUUID(), ...(argArray[0] ?? {}) }];
+						}
+						const instanceId = (argArray[0] as { id: string }).id;
+
+						await introspectAndModifyInstance(instanceId);
+
+						return target[property](...argArray);
+					},
+				});
+			}
+
+			if (property === "createBatch") {
+				return new Proxy(target[property], {
+					async apply(func, thisArg, argArray) {
+						for (const [index, arg] of argArray[0]?.entries() ?? []) {
+							const hasId = Object.hasOwn(arg, "id");
+							if (!hasId) {
+								argArray[0][index] = { id: crypto.randomUUID(), ...arg };
+							}
+						}
+
+						await Promise.all(
+							argArray[0].map((options: { id: string }) =>
+								introspectAndModifyInstance(options.id)
+							)
+						);
+						return target[property](argArray[0]);
+					},
+				});
+			}
+			// @ts-expect-error index signature
+			return target[property];
+		};
+	};
+
+	const cleanup = () => {
+		internalEnv[bindingName] = internalOriginalWorkflow;
+		env[bindingName] = externalOriginalWorkflow;
+	};
+
+	// Create a single handler instance to be reused
+	const proxyGetHandler = createWorkflowProxyGetHandler();
+
+	// Apply the proxies using the shared handler logic
+	internalEnv[bindingName] = new Proxy(internalOriginalWorkflow, {
+		get: proxyGetHandler,
+	});
+	env[bindingName] = new Proxy(externalOriginalWorkflow, {
+		get: proxyGetHandler,
+	});
+
+	return new WorkflowIntrospectorHandle(
+		workflow,
+		modifierCallbacks,
+		instanceIntrospectors,
+		cleanup
+	);
+}
+
+class WorkflowIntrospectorHandle implements WorkflowIntrospector {
+	workflow: Workflow;
+	#modifierCallbacks: ModifierCallback[];
+	#instanceIntrospectors: WorkflowInstanceIntrospector[];
+	#cleanupCallback: () => void;
+
+	constructor(
+		workflow: Workflow,
+		modifierCallbacks: ModifierCallback[],
+		instanceIntrospectors: WorkflowInstanceIntrospector[],
+		cleanupCallback: () => void
+	) {
+		this.workflow = workflow;
+		this.#modifierCallbacks = modifierCallbacks;
+		this.#instanceIntrospectors = instanceIntrospectors;
+		this.#cleanupCallback = cleanupCallback;
+	}
+
+	modifyAll(fn: ModifierCallback): void {
+		this.#modifierCallbacks.push(fn);
+	}
+
+	get(): WorkflowInstanceIntrospector[] {
+		return this.#instanceIntrospectors;
+	}
+
+	cleanUp(): void {
+		this.#modifierCallbacks = [];
+		this.#instanceIntrospectors = [];
+		this.#cleanupCallback();
+	}
+}

packages/workflows-shared/src/binding.ts

@@ -11,6 +11,7 @@ import type {
 type Env = {
 	ENGINE: DurableObjectNamespace<Engine>;
 	WORKFLOW_NAME: string;
+	BINDING_NAME: string;
 };
 
 // this.env.WORKFLOW is WorkflowBinding
@@ -79,9 +80,13 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> implements Workflow {
 
 		return await Promise.all(batch.map((val) => this.create(val)));
 	}
-	public async getWorkflowName(): Promise<string> {
+	public getWorkflowName(): string {
 		return this.env.WORKFLOW_NAME;
 	}
+
+	public getBindingName(): string {
+		return this.env.BINDING_NAME;
+	}
 }
 
 export class WorkflowHandle extends RpcTarget implements WorkflowInstance {