fix(cloudflare): Forward ctx argument to Workflow.do user callback

Author: Lms24

packages/cloudflare/src/workflows.ts

@@ -67,24 +67,27 @@ class WrappedWorkflowStep implements WorkflowStep {
     private _step: WorkflowStep,
   ) {}
 
-  public async do<T extends Rpc.Serializable<T>>(name: string, callback: () => Promise<T>): Promise<T>;
+  public async do<T extends Rpc.Serializable<T>>(
+    name: string,
+    callback: (...args: unknown[]) => Promise<T>,
+  ): Promise<T>;
   public async do<T extends Rpc.Serializable<T>>(
     name: string,
     config: WorkflowStepConfig,
-    callback: () => Promise<T>,
+    callback: (...args: unknown[]) => Promise<T>,
   ): Promise<T>;
   public async do<T extends Rpc.Serializable<T>>(
     name: string,
     configOrCallback: WorkflowStepConfig | (() => Promise<T>),
-    maybeCallback?: () => Promise<T>,
+    maybeCallback?: (...args: unknown[]) => Promise<T>,
   ): Promise<T> {
     // Capture the current scope, so parent span (e.g., a startSpan surrounding step.do) is preserved
     const scopeForStep = getCurrentScope();
 
-    const userCallback = (maybeCallback || configOrCallback) as () => Promise<T>;
+    const userCallback = (maybeCallback || configOrCallback) as (...args: unknown[]) => Promise<T>;
     const config = typeof configOrCallback === 'function' ? undefined : configOrCallback;
 
-    const instrumentedCallback: () => Promise<T> = async () => {
+    const instrumentedCallback = async (...args: unknown[]): Promise<T> => {
       return startSpan(
         {
           op: 'function.step.do',
@@ -101,7 +104,7 @@ class WrappedWorkflowStep implements WorkflowStep {
         },
         async span => {
           try {
-            const result = await userCallback();
+            const result = await userCallback(...args);
             span.setStatus({ code: 1 });
             return result;
           } catch (error) {

packages/cloudflare/test/workflow.test.ts

@@ -6,14 +6,16 @@ import { deterministicTraceIdFromInstanceId, instrumentWorkflowWithSentry } from
 
 const NODE_MAJOR_VERSION = parseInt(process.versions.node.split('.')[0]!);
 
+const MOCK_STEP_CTX = { attempt: 1 };
+
 const mockStep: WorkflowStep = {
   do: vi
     .fn()
     .mockImplementation(
       async (
         _name: string,
-        configOrCallback: WorkflowStepConfig | (() => Promise<any>),
-        maybeCallback?: () => Promise<any>,
+        configOrCallback: WorkflowStepConfig | ((...args: unknown[]) => Promise<any>),
+        maybeCallback?: (...args: unknown[]) => Promise<any>,
       ) => {
         let count = 0;
 
@@ -22,9 +24,9 @@ const mockStep: WorkflowStep = {
 
           try {
             if (typeof configOrCallback === 'function') {
-              return await configOrCallback();
+              return await configOrCallback(MOCK_STEP_CTX);
             } else {
-              return await (maybeCallback ? maybeCallback() : Promise.resolve());
+              return await (maybeCallback ? maybeCallback(MOCK_STEP_CTX) : Promise.resolve());
             }
           } catch {
             await new Promise(resolve => setTimeout(resolve, 1000));
@@ -427,6 +429,26 @@ describe.skipIf(NODE_MAJOR_VERSION < 20)('workflows', () => {
     ]);
   });
 
+  test('Forwards step context (ctx) to user callback', async () => {
+    const callbackSpy = vi.fn().mockResolvedValue({ ok: true });
+
+    class CtxTestWorkflow {
+      constructor(_ctx: ExecutionContext, _env: unknown) {}
+
+      async run(_event: Readonly<WorkflowEvent<Params>>, step: WorkflowStep): Promise<void> {
+        await step.do('ctx step', callbackSpy);
+      }
+    }
+
+    const TestWorkflowInstrumented = instrumentWorkflowWithSentry(getSentryOptions, CtxTestWorkflow as any);
+    const workflow = new TestWorkflowInstrumented(mockContext, {}) as CtxTestWorkflow;
+    const event = { payload: {}, timestamp: new Date(), instanceId: INSTANCE_ID };
+    await workflow.run(event, mockStep);
+
+    expect(callbackSpy).toHaveBeenCalledTimes(1);
+    expect(callbackSpy).toHaveBeenCalledWith(MOCK_STEP_CTX);
+  });
+
   test('Step.do span becomes child of surrounding custom span', async () => {
     class ParentChildWorkflow {
       constructor(_ctx: ExecutionContext, _env: unknown) {}