fix(v9/cloudflare): Avoid turning DurableObject sync methods into async (#17187)

Backport of #17184

Author: andreiborza

packages/cloudflare/src/durableobject.ts

@@ -1,8 +1,10 @@
 /* eslint-disable @typescript-eslint/unbound-method */
 import {
+  type Scope,
   captureException,
   flush,
   getClient,
+  isThenable,
   SEMANTIC_ATTRIBUTE_SENTRY_OP,
   SEMANTIC_ATTRIBUTE_SENTRY_ORIGIN,
   startSpan,
@@ -46,7 +48,8 @@ function wrapMethodWithSentry<T extends OriginalMethod>(
       const currentClient = getClient();
       // if a client is already set, use withScope, otherwise use withIsolationScope
       const sentryWithScope = currentClient ? withScope : withIsolationScope;
-      return sentryWithScope(async scope => {
+
+      const wrappedFunction = (scope: Scope): unknown => {
         // In certain situations, the passed context can become undefined.
         // For example, for Astro while prerendering pages at build time.
         // see: https://github.com/getsentry/sentry-javascript/issues/13217
@@ -65,17 +68,38 @@ function wrapMethodWithSentry<T extends OriginalMethod>(
             if (callback) {
               callback(...args);
             }
-            return await Reflect.apply(target, thisArg, args);
+            const result = Reflect.apply(target, thisArg, args);
+
+            if (isThenable(result)) {
+              return result.then(
+                (res: unknown) => {
+                  waitUntil?.(flush(2000));
+                  return res;
+                },
+                (e: unknown) => {
+                  captureException(e, {
+                    mechanism: {
+                      type: 'cloudflare_durableobject',
+                      handled: false,
+                    },
+                  });
+                  waitUntil?.(flush(2000));
+                  throw e;
+                },
+              );
+            } else {
+              waitUntil?.(flush(2000));
+              return result;
+            }
           } catch (e) {
             captureException(e, {
               mechanism: {
                 type: 'cloudflare_durableobject',
                 handled: false,
               },
             });
-            throw e;
-          } finally {
             waitUntil?.(flush(2000));
+            throw e;
           }
         }
 
@@ -87,22 +111,45 @@ function wrapMethodWithSentry<T extends OriginalMethod>(
           : {};
 
         // Only create these spans if they have a parent span.
-        return startSpan({ name: wrapperOptions.spanName, attributes, onlyIfParent: true }, async () => {
+        return startSpan({ name: wrapperOptions.spanName, attributes, onlyIfParent: true }, () => {
           try {
-            return await Reflect.apply(target, thisArg, args);
+            const result = Reflect.apply(target, thisArg, args);
+
+            if (isThenable(result)) {
+              return result.then(
+                (res: unknown) => {
+                  waitUntil?.(flush(2000));
+                  return res;
+                },
+                (e: unknown) => {
+                  captureException(e, {
+                    mechanism: {
+                      type: 'cloudflare_durableobject',
+                      handled: false,
+                    },
+                  });
+                  waitUntil?.(flush(2000));
+                  throw e;
+                },
+              );
+            } else {
+              waitUntil?.(flush(2000));
+              return result;
+            }
           } catch (e) {
             captureException(e, {
               mechanism: {
                 type: 'cloudflare_durableobject',
                 handled: false,
               },
             });
-            throw e;
-          } finally {
             waitUntil?.(flush(2000));
+            throw e;
           }
         });
-      });
+      };
+
+      return sentryWithScope(wrappedFunction);
     },
   });
 }

packages/cloudflare/test/durableobject.test.ts

@@ -8,20 +8,43 @@ describe('instrumentDurableObjectWithSentry', () => {
   afterEach(() => {
     vi.restoreAllMocks();
   });
+
   it('Generic functionality', () => {
     const options = vi.fn();
     const instrumented = instrumentDurableObjectWithSentry(options, vi.fn());
     expect(instrumented).toBeTypeOf('function');
     expect(() => Reflect.construct(instrumented, [])).not.toThrow();
     expect(options).toHaveBeenCalledOnce();
   });
-  it('Instruments prototype methods and defines implementation in the object', () => {
+
+  it('Instruments sync prototype methods and defines implementation in the object', () => {
     const testClass = class {
-      method() {}
+      method() {
+        return 'sync-result';
+      }
     };
     const obj = Reflect.construct(instrumentDurableObjectWithSentry(vi.fn(), testClass as any), []) as any;
     expect(obj.method).toBe(obj.method);
+
+    const result = obj.method();
+    expect(result).not.toBeInstanceOf(Promise);
+    expect(result).toEqual('sync-result');
   });
+
+  it('Instruments async prototype methods and returns a promise', async () => {
+    const testClass = class {
+      async asyncMethod() {
+        return 'async-result';
+      }
+    };
+    const obj = Reflect.construct(instrumentDurableObjectWithSentry(vi.fn(), testClass as any), []) as any;
+    expect(obj.asyncMethod).toBe(obj.asyncMethod);
+
+    const result = obj.asyncMethod();
+    expect(result).toBeInstanceOf(Promise);
+    expect(await result).toBe('async-result');
+  });
+
   it('Instruments prototype methods without "sticking" to the options', () => {
     const initCore = vi.spyOn(SentryCore, 'initAndBind');
     vi.spyOn(SentryCore, 'getClient').mockReturnValue(undefined);
@@ -41,6 +64,7 @@ describe('instrumentDurableObjectWithSentry', () => {
     expect(initCore).nthCalledWith(1, expect.any(Function), expect.objectContaining({ orgId: 1 }));
     expect(initCore).nthCalledWith(2, expect.any(Function), expect.objectContaining({ orgId: 2 }));
   });
+
   it('All available durable object methods are instrumented', () => {
     const testClass = class {
       propertyFunction = vi.fn();
@@ -72,6 +96,7 @@ describe('instrumentDurableObjectWithSentry', () => {
       expect(isInstrumented((obj as any)[method_name]), `Method ${method_name} is instrumented`).toBeTruthy();
     }
   });
+
   it('flush performs after all waitUntil promises are finished', async () => {
     vi.useFakeTimers();
     onTestFinished(() => {