Add extension in front of workflow binding for manually disposing engine stubs

Author: Caio-Nogueira

.changeset/tricky-squids-wait.md

@@ -0,0 +1,7 @@
+---
+"@cloudflare/vitest-pool-workers": minor
+"@cloudflare/workflows-shared": minor
+"miniflare": minor
+---
+
+migrate workflow to use a wrapped binding

fixtures/workflow/tests/index.test.ts

@@ -48,7 +48,11 @@ describe("Workflows", () => {
 		await expect(fetchJson(`http://${ip}:${port}/create?workflowName=test`))
 			.resolves.toMatchInlineSnapshot(`
 			{
-			  "__LOCAL_DEV_STEP_OUTPUTS": [],
+			  "__LOCAL_DEV_STEP_OUTPUTS": [
+			    {
+			      "output": "First step result",
+			    },
+			  ],
 			  "output": null,
 			  "status": "running",
 			}

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

@@ -1,5 +1,6 @@
 import fs from "fs/promises";
 import SCRIPT_WORKFLOWS_BINDING from "worker:workflows/binding";
+import SCRIPT_WORKFLOWS_WRAPPED_BINDING from "worker:workflows/wrapped-binding";
 import { z } from "zod";
 import { Service } from "../../runtime";
 import { getUserServiceName } from "../core";
@@ -43,13 +44,21 @@ export const WORKFLOWS_PLUGIN: Plugin<
 		return Object.entries(options.workflows ?? {}).map(
 			([bindingName, workflow]) => ({
 				name: bindingName,
-				service: {
-					name: getUserBindingServiceName(
-						WORKFLOWS_PLUGIN_NAME,
-						workflow.name,
-						workflow.remoteProxyConnectionString
-					),
-					entrypoint: "WorkflowBinding",
+				wrapped: {
+					moduleName: `${WORKFLOWS_PLUGIN_NAME}:local-wrapped-binding`,
+					innerBindings: [
+						{
+							name: "binding",
+							service: {
+								name: getUserBindingServiceName(
+									WORKFLOWS_PLUGIN_NAME,
+									workflow.name,
+									workflow.remoteProxyConnectionString
+								),
+								entrypoint: "WorkflowBinding",
+							},
+						},
+					],
 				},
 			})
 		);
@@ -64,6 +73,20 @@ export const WORKFLOWS_PLUGIN: Plugin<
 		);
 	},
 
+	getExtensions({}) {
+		return [
+			{
+				modules: [
+					{
+						name: `${WORKFLOWS_PLUGIN_NAME}:local-wrapped-binding`,
+						esModule: SCRIPT_WORKFLOWS_WRAPPED_BINDING(),
+						internal: true,
+					},
+				],
+			},
+		];
+	},
+
 	async getServices({ options, sharedOptions, tmpPath, defaultPersistRoot }) {
 		const persistPath = getPersistPath(
 			WORKFLOWS_PLUGIN_NAME,

packages/miniflare/src/workers/workflows/wrapped-binding.worker.ts

@@ -0,0 +1,57 @@
+import { WorkflowBinding } from "@cloudflare/workflows-shared/src/binding";
+
+class WorkflowImpl {
+	constructor(private binding: WorkflowBinding) {
+		this.binding = binding;
+	}
+
+	async get(id: string): Promise<WorkflowInstance> {
+		using result = await this.binding.get(id);
+		return result;
+	}
+
+	async create(
+		options?: WorkflowInstanceCreateOptions
+	): Promise<WorkflowInstance> {
+		using result = await this.binding.create(options);
+		return result;
+	}
+
+	async createBatch(
+		options: WorkflowInstanceCreateOptions[]
+	): Promise<WorkflowInstance[]> {
+		// no need to dispose here since the binding already does so when mapping
+		// through created instances
+		return this.binding.createBatch(options);
+	}
+
+	async unsafeGetBindingName(): Promise<string> {
+		return this.binding.unsafeGetBindingName();
+	}
+
+	async unsafeAbort(instanceId: string, reason?: string): Promise<void> {
+		return this.binding.unsafeAbort(instanceId, reason);
+	}
+
+	async unsafeGetInstanceModifier(instanceId: string): Promise<unknown> {
+		return this.binding.unsafeGetInstanceModifier(instanceId);
+	}
+
+	async unsafeWaitForStepResult(
+		instanceId: string,
+		name: string,
+		index?: number
+	): Promise<unknown> {
+		return this.binding.unsafeWaitForStepResult(instanceId, name, index);
+	}
+
+	async unsafeWaitForStatus(instanceId: string, status: string): Promise<void> {
+		await this.binding.unsafeWaitForStatus(instanceId, status);
+	}
+}
+
+export function makeBinding(env: { binding: WorkflowBinding }): Workflow {
+	return new WorkflowImpl(env.binding);
+}
+
+export default makeBinding;

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

@@ -45,18 +45,29 @@ class WorkflowInstanceIntrospectorHandle
 {
 	#workflow: WorkflowBinding;
 	#instanceId: string;
-	#instanceModifier: WorkflowInstanceModifier;
+	#instanceModifier: WorkflowInstanceModifier | undefined;
+	#instanceModifierPromise: Promise<WorkflowInstanceModifier> | undefined;
 
 	constructor(workflow: WorkflowBinding, instanceId: string) {
 		this.#workflow = workflow;
 		this.#instanceId = instanceId;
-
-		this.#instanceModifier = workflow.unsafeGetInstanceModifier(
-			instanceId
-		) as WorkflowInstanceModifier;
+		this.#instanceModifierPromise = workflow
+			.unsafeGetInstanceModifier(instanceId)
+			.then((res) => {
+				this.#instanceModifier = res as WorkflowInstanceModifier;
+				this.#instanceModifierPromise = undefined;
+				return this.#instanceModifier;
+			});
 	}
 
 	async modify(fn: ModifierCallback): Promise<WorkflowInstanceIntrospector> {
+		if (this.#instanceModifierPromise !== undefined) {
+			this.#instanceModifier = await this.#instanceModifierPromise;
+		}
+		if (this.#instanceModifier === undefined) {
+			throw new Error("instance modifier is undefined");
+		}
+
 		await fn(this.#instanceModifier);
 
 		return this;

packages/workflows-shared/src/binding.ts

@@ -17,10 +17,16 @@ type Env = {
 
 // this.env.WORKFLOW is WorkflowBinding
 export class WorkflowBinding extends WorkerEntrypoint<Env> implements Workflow {
+	constructor(ctx: ExecutionContext, env: Env) {
+		super(ctx, env);
+	}
+
 	public async create({
 		id = crypto.randomUUID(),
 		params = {},
-	}: WorkflowInstanceCreateOptions = {}): Promise<WorkflowInstance> {
+	}: WorkflowInstanceCreateOptions = {}): Promise<
+		WorkflowInstance & Disposable
+	> {
 		if (!isValidWorkflowInstanceId(id)) {
 			throw new WorkflowError("Workflow instance has invalid id");
 		}
@@ -51,10 +57,11 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> implements Workflow {
 			restart: handle.restart.bind(handle),
 			status: handle.status.bind(handle),
 			sendEvent: handle.sendEvent.bind(handle),
+			[Symbol.dispose]: handle.dispose.bind(handle),
 		};
 	}
 
-	public async get(id: string): Promise<WorkflowInstance> {
+	public async get(id: string): Promise<WorkflowInstance & Disposable> {
 		const engineStubId = this.env.ENGINE.idFromName(id);
 		const engineStub = this.env.ENGINE.get(engineStubId);
 		const handle = new WorkflowHandle(id, engineStub);
@@ -73,6 +80,7 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> implements Workflow {
 			restart: handle.restart.bind(handle),
 			status: handle.status.bind(handle),
 			sendEvent: handle.sendEvent.bind(handle),
+			[Symbol.dispose]: handle.dispose.bind(handle),
 		};
 	}
 	public async createBatch(
@@ -84,14 +92,21 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> implements Workflow {
 			);
 		}
 
-		return await Promise.all(batch.map((val) => this.create(val)));
+		return await Promise.all(
+			batch.map(async (val) => {
+				using res = await this.create(val);
+				return res;
+			})
+		);
 	}
 
-	public unsafeGetBindingName(): string {
+	public async unsafeGetBindingName(): Promise<string> {
+		// async because of rpc
 		return this.env.BINDING_NAME;
 	}
 
-	public unsafeGetInstanceModifier(instanceId: string): unknown {
+	public async unsafeGetInstanceModifier(instanceId: string): Promise<unknown> {
+		// async because of rpc
 		const stubId = this.env.ENGINE.idFromName(instanceId);
 		const stub = this.env.ENGINE.get(stubId);
 
@@ -108,7 +123,8 @@ export class WorkflowBinding extends WorkerEntrypoint<Env> implements Workflow {
 		const stubId = this.env.ENGINE.idFromName(instanceId);
 		const stub = this.env.ENGINE.get(stubId);
 
-		return await stub.waitForStepResult(name, index);
+		using res = await stub.waitForStepResult(name, index);
+		return res;
 	}
 
 	public async unsafeAbort(instanceId: string, reason?: string): Promise<void> {
@@ -207,4 +223,9 @@ export class WorkflowHandle extends RpcTarget implements WorkflowInstance {
 			timestamp: new Date(),
 		});
 	}
+
+	public dispose() {
+		// @ts-expect-error - DurableObjectStub should be disposable but types don't reflect it
+		this.stub[Symbol.dispose]?.();
+	}
 }

packages/workflows-shared/src/engine.ts

@@ -21,7 +21,7 @@ import type { Event } from "./context";
 import type { InstanceMetadata, RawInstanceLog } from "./instance";
 import type { WorkflowEntrypoint, WorkflowEvent } from "cloudflare:workers";
 
-export interface Env {
+interface Env {
 	USER_WORKFLOW: WorkflowEntrypoint;
 }
 

packages/workflows-shared/src/local-binding-worker.ts

@@ -1,2 +1,2 @@
-export { WorkflowBinding } from "./binding";
 export { Engine } from "./engine";
+export { WorkflowBinding } from "./binding";