fix(cloudflare): Get original waituntil in workflows

packages/cloudflare/src/flush.ts

@@ -35,7 +35,7 @@ const flushLockRegistries = new WeakMap<ExecutionContext['waitUntil'], FlushLock
  *
  * By using the original waitUntil for flush operations, we bypass this issue.
  */
-export function getOriginalWaitUntil(context: ExecutionContextCompat): ExecutionContext['waitUntil'] | undefined {
+export function getOriginalWaitUntil(context: ExecutionContextCompat): ExecutionContext['waitUntil'] {
   // eslint-disable-next-line @typescript-eslint/unbound-method
   const currentWaitUntil = context.waitUntil;
   const original = flushLockRegistries.get(currentWaitUntil)?.originalWaitUntil;

packages/cloudflare/src/request.ts

@@ -80,7 +80,7 @@ export function wrapRequestHandlerWithInit(
     // to track pending tasks. If we use the instrumented version for flushAndDispose,
     // it acquires the lock, then flushAndDispose tries to wait for the same lock,
     // creating a deadlock.
-    const waitUntil = context ? getOriginalWaitUntil(context)?.bind(context) : undefined;
+    const waitUntil = context ? getOriginalWaitUntil(context).bind(context) : undefined;
     const errorMechanismType = getRequestErrorMechanismType(context);
 
     const client = initSdk({ ...options, ctx: context });

packages/cloudflare/src/workflows.ts

@@ -25,7 +25,7 @@ import type {
 } from 'cloudflare:workers';
 import { setAsyncLocalStorageAsyncContextStrategy } from '@sentry/server-utils/no-diagnostic-channels';
 import type { CloudflareOptions } from './client';
-import { flushAndDispose } from './flush';
+import { flushAndDispose, getOriginalWaitUntil } from './flush';
 import { instrumentEnv } from './instrumentations/worker/instrumentEnv';
 import { addCloudResourceContext } from './scope-utils';
 import { init } from './sdk';
@@ -218,7 +218,7 @@ export function instrumentWorkflowWithSentry<
               setAsyncLocalStorageAsyncContextStrategy();
 
               return withIsolationScope(async isolationScope => {
-                const waitUntil = context.waitUntil.bind(context);
+                const waitUntil = getOriginalWaitUntil(context).bind(context);
                 const client = init({ ...options, ctx: context, enableDedupe: false });
                 isolationScope.setClient(client);
 

packages/cloudflare/test/flush.test.ts

@@ -165,7 +165,7 @@ describe('getOriginalWaitUntil', () => {
 
     expect(result).not.toBe(context.waitUntil);
     expect(result).toBeDefined();
-    result!(Promise.resolve());
+    result(Promise.resolve());
     expect(originalWaitUntil).toHaveBeenCalled();
   });
 
@@ -183,7 +183,7 @@ describe('getOriginalWaitUntil', () => {
     const result = getOriginalWaitUntil(context);
 
     expect(result).not.toBe(context.waitUntil);
-    result!(Promise.resolve());
+    result(Promise.resolve());
     expect(originalWaitUntil).toHaveBeenCalled();
   });
 
@@ -207,7 +207,7 @@ describe('getOriginalWaitUntil', () => {
     } as unknown as Client;
 
     const originalWaitUntil = getOriginalWaitUntil(context);
-    originalWaitUntil!.call(context, flushAndDispose(mockClient));
+    originalWaitUntil.call(context, flushAndDispose(mockClient));
 
     await vi.waitFor(() => Promise.all(waitUntilPromises));
     expect(mockClient.flush).toHaveBeenCalled();

packages/cloudflare/test/workflow.test.ts

@@ -209,6 +209,135 @@ describe.skipIf(NODE_MAJOR_VERSION < 20)('workflows', () => {
     await expect(drainWaitUntilLikeCloudflareVitestPool(waitUntilPromises)).resolves.toBeUndefined();
   });
 
+  test('teardown does not deadlock when a workflow instance is reused across runs', async () => {
+    const waitUntilPromises: Promise<unknown>[] = [];
+    const context: ExecutionContext = {
+      waitUntil: vi.fn((promise: Promise<unknown>) => {
+        waitUntilPromises.push(promise);
+      }),
+      passThroughOnException: vi.fn(),
+      props: {},
+    };
+
+    let runCount = 0;
+    let releaseAppWork: () => void = () => undefined;
+
+    class ReusedWorkflow {
+      public constructor(private _ctx: ExecutionContext) {}
+
+      public async run(_event: Readonly<WorkflowEvent<Params>>, step: WorkflowStep): Promise<void> {
+        runCount += 1;
+        await step.do('reused step', async () => {
+          if (runCount === 2) {
+            this._ctx.waitUntil(
+              new Promise<void>(resolve => {
+                releaseAppWork = resolve;
+              }),
+            );
+          }
+        });
+      }
+    }
+
+    const TestWorkflowInstrumented = instrumentWorkflowWithSentry(getSentryOptions, ReusedWorkflow as any);
+    // Cloudflare reuses a Workflow instance across runs, so the context
+    // captured at construction is instrumented by the first run's init()
+    const workflow = new TestWorkflowInstrumented(context, {}) as ReusedWorkflow;
+    const event = { payload: {}, timestamp: new Date(), instanceId: INSTANCE_ID };
+
+    await workflow.run(event, mockStep);
+    await drainWaitUntilLikeCloudflareVitestPool(waitUntilPromises);
+
+    await workflow.run(event, mockStep);
+
+    releaseAppWork();
+
+    // Both the application work and the teardown promise must settle
+    await expect(drainWaitUntilLikeCloudflareVitestPool(waitUntilPromises)).resolves.toBeUndefined();
+  });
+
+  test('step errors are still captured when a workflow instance is reused across runs', async () => {
+    const waitUntilPromises: Promise<unknown>[] = [];
+    const context: ExecutionContext = {
+      waitUntil: vi.fn((promise: Promise<unknown>) => {
+        waitUntilPromises.push(promise);
+      }),
+      passThroughOnException: vi.fn(),
+      props: {},
+    };
+
+    let runCount = 0;
+
+    class ReusedErrorWorkflow {
+      public constructor(private _ctx: ExecutionContext) {}
+
+      public async run(_event: Readonly<WorkflowEvent<Params>>, step: WorkflowStep): Promise<void> {
+        runCount += 1;
+        await step.do('flaky step', async () => {
+          if (runCount === 2) {
+            throw new Error('second run error');
+          }
+        });
+      }
+    }
+
+    // Fails the step through every retry without backoff, so the error is
+    // captured on the final attempt and surfaces from run()
+    const alwaysFailStep: WorkflowStep = {
+      do: vi
+        .fn()
+        .mockImplementation(
+          async (
+            _name: string,
+            configOrCallback: WorkflowStepConfig | ((...args: unknown[]) => Promise<any>),
+            maybeCallback?: (...args: unknown[]) => Promise<any>,
+          ) => {
+            const retryLimit = 2;
+            const callback = (typeof configOrCallback === 'function' ? configOrCallback : maybeCallback)!;
+            let lastError: unknown;
+            for (let attempt = 1; attempt <= retryLimit + 1; attempt++) {
+              try {
+                return await callback({ attempt, config: { retries: { limit: retryLimit }, timeout: 60000 } });
+              } catch (err) {
+                lastError = err;
+              }
+            }
+            throw lastError;
+          },
+        ),
+      sleep: vi.fn(),
+      sleepUntil: vi.fn(),
+      waitForEvent: vi.fn(),
+    };
+
+    const TestWorkflowInstrumented = instrumentWorkflowWithSentry(getSentryOptions, ReusedErrorWorkflow as any);
+    const workflow = new TestWorkflowInstrumented(context, {}) as ReusedErrorWorkflow;
+    const event = { payload: {}, timestamp: new Date(), instanceId: INSTANCE_ID };
+
+    await workflow.run(event, mockStep);
+    await drainWaitUntilLikeCloudflareVitestPool(waitUntilPromises);
+
+    await expect(workflow.run(event, alwaysFailStep)).rejects.toThrow('second run error');
+    await expect(drainWaitUntilLikeCloudflareVitestPool(waitUntilPromises)).resolves.toBeUndefined();
+
+    const errorEnvelopes = mockTransport.send.mock.calls.filter(call => {
+      const items = (call[0] as any)[1] as any[];
+      return items.some(i => i[0].type === 'event');
+    });
+    expect(errorEnvelopes).toHaveLength(1);
+    expect(errorEnvelopes[0]![0][1][0][1]).toMatchObject({
+      exception: {
+        values: [
+          expect.objectContaining({
+            type: 'Error',
+            value: 'second run error',
+            mechanism: { type: 'auto.faas.cloudflare.workflow', handled: true },
+          }),
+        ],
+      },
+    });
+  });
+
   test('Wraps env with instrumentEnv', async () => {
     class EnvTestWorkflow {
       constructor(_ctx: ExecutionContext, _env: unknown) {}