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,117 @@
     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 shouldThrow = false;
+
+    class ReusedErrorWorkflow {
+      public constructor(private _ctx: ExecutionContext) {}
+
+      public async run(_event: Readonly<WorkflowEvent<Params>>, step: WorkflowStep): Promise<void> {
+        await step.do('flaky step', async () => {
+          if (shouldThrow) {
+            shouldThrow = false;
+            throw new Error('second run error');
+          }
+        });
+      }
+    }
+
+    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);
+
+    shouldThrow = true;
+    await workflow.run(event, mockStep);
+    await drainWaitUntilLikeCloudflareVitestPool(waitUntilPromises);
+
+    expect(mockTransport.send).toHaveBeenCalledWith([
+      expect.objectContaining({
+        trace: expect.objectContaining({
+          transaction: 'flaky step',
+          trace_id: TRACE_ID,
+        }),
+      }),
+      [
+        [
+          {
+            type: 'event',
+          },
+          expect.objectContaining({
+            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) {}