From 5bb85dee5a6d35d5df26fca82e2ce3b70a16953b Mon Sep 17 00:00:00 2001 From: Lukas Stracke Date: Thu, 19 Mar 2026 18:52:06 +0100 Subject: [PATCH] fix(cloudflare): Forward `ctx` argument to `Workflow.do` user callback --- packages/cloudflare/src/workflows.ts | 15 +++++++----- packages/cloudflare/test/workflow.test.ts | 30 ++++++++++++++++++++--- 2 files changed, 35 insertions(+), 10 deletions(-) diff --git a/packages/cloudflare/src/workflows.ts b/packages/cloudflare/src/workflows.ts index 3c40c86ff867..6515a330ca99 100644 --- a/packages/cloudflare/src/workflows.ts +++ b/packages/cloudflare/src/workflows.ts @@ -67,24 +67,27 @@ class WrappedWorkflowStep implements WorkflowStep { private _step: WorkflowStep, ) {} - public async do>(name: string, callback: () => Promise): Promise; + public async do>( + name: string, + callback: (...args: unknown[]) => Promise, + ): Promise; public async do>( name: string, config: WorkflowStepConfig, - callback: () => Promise, + callback: (...args: unknown[]) => Promise, ): Promise; public async do>( name: string, configOrCallback: WorkflowStepConfig | (() => Promise), - maybeCallback?: () => Promise, + maybeCallback?: (...args: unknown[]) => Promise, ): Promise { // 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; + const userCallback = (maybeCallback || configOrCallback) as (...args: unknown[]) => Promise; const config = typeof configOrCallback === 'function' ? undefined : configOrCallback; - const instrumentedCallback: () => Promise = async () => { + const instrumentedCallback = async (...args: unknown[]): Promise => { 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) { diff --git a/packages/cloudflare/test/workflow.test.ts b/packages/cloudflare/test/workflow.test.ts index b460e6bfee5a..f21bee8612a8 100644 --- a/packages/cloudflare/test/workflow.test.ts +++ b/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), - maybeCallback?: () => Promise, + configOrCallback: WorkflowStepConfig | ((...args: unknown[]) => Promise), + maybeCallback?: (...args: unknown[]) => Promise, ) => { 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>, step: WorkflowStep): Promise { + 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) {}