Instrument methods at prototype level not instance level

Author: andreiborza

packages/cloudflare/src/durableobject.ts

@@ -110,8 +110,7 @@ function wrapMethodWithSentry<T extends OriginalMethod>(
             }
           : {};
 
-        // Only create these spans if they have a parent span.
-        return startSpan({ name: wrapperOptions.spanName, attributes, onlyIfParent: true }, () => {
+        return startSpan({ name: wrapperOptions.spanName, attributes }, () => {
           try {
             const result = Reflect.apply(target, thisArg, args);
 
@@ -273,46 +272,87 @@ export function instrumentDurableObjectWithSentry<
           );
         }
       }
-      const instrumentedPrototype = instrumentPrototype(target, options, context);
-      Object.setPrototypeOf(obj, instrumentedPrototype);
+
+      // Store context and options on the instance for prototype methods to access
+      Object.defineProperty(obj, '__SENTRY_CONTEXT__', {
+        value: context,
+        enumerable: false,
+        writable: false,
+        configurable: false,
+      });
+
+      Object.defineProperty(obj, '__SENTRY_OPTIONS__', {
+        value: options,
+        enumerable: false,
+        writable: false,
+        configurable: false,
+      });
+
+      instrumentPrototype(target);
 
       return obj;
     },
   });
 }
 
-function instrumentPrototype<T extends NewableFunction>(
-  target: T,
-  options: CloudflareOptions,
-  context: MethodWrapperOptions['context'],
-): T {
-  return new Proxy(target.prototype, {
-    get(target, prop, receiver) {
-      const value = Reflect.get(target, prop, receiver);
-      if (prop === 'constructor' || typeof value !== 'function') {
-        return value;
+function instrumentPrototype<T extends NewableFunction>(target: T): void {
+  const proto = target.prototype;
+
+  // Get all methods from the prototype chain
+  const methodNames = new Set<string>();
+  let current = proto;
+
+  while (current && current !== Object.prototype) {
+    Object.getOwnPropertyNames(current).forEach(name => {
+      if (name !== 'constructor' && typeof current[name] === 'function') {
+        methodNames.add(name);
+      }
+    });
+    current = Object.getPrototypeOf(current);
+  }
+
+  // Instrument each method on the prototype
+  methodNames.forEach(methodName => {
+    const originalMethod = proto[methodName];
+
+    if (!originalMethod || isInstrumented(originalMethod)) {
+      return;
+    }
+
+    // Create a wrapper that gets context/options from the instance at runtime
+    const wrappedMethod = function (this: any, ...args: any[]) {
+      const instanceContext = this.__SENTRY_CONTEXT__;
+      const instanceOptions = this.__SENTRY_OPTIONS__;
+
+      if (!instanceOptions) {
+        // Fallback to original method if no Sentry data found
+        return originalMethod.apply(this, args);
       }
-      const wrapped = wrapMethodWithSentry(
-        { options, context, spanName: prop.toString(), spanOp: 'rpc' },
-        value,
+
+      // Use the existing wrapper but with instance-specific context/options
+      const wrapper = wrapMethodWithSentry(
+        {
+          options: instanceOptions,
+          context: instanceContext,
+          spanName: methodName,
+          spanOp: 'rpc',
+        },
+        originalMethod,
         undefined,
-        true,
+        true, // noMark = true since we'll mark the prototype method
       );
-      const instrumented = new Proxy(wrapped, {
-        get(target, p, receiver) {
-          if ('__SENTRY_INSTRUMENTED__' === p) {
-            return true;
-          }
-          return Reflect.get(target, p, receiver);
-        },
-      });
-      Object.defineProperty(receiver, prop, {
-        value: instrumented,
-        enumerable: true,
-        writable: true,
-        configurable: true,
-      });
-      return instrumented;
-    },
+
+      return wrapper.apply(this, args);
+    };
+
+    markAsInstrumented(wrappedMethod);
+
+    // Replace the prototype method
+    Object.defineProperty(proto, methodName, {
+      value: wrappedMethod,
+      enumerable: false,
+      writable: true,
+      configurable: true,
+    });
   });
 }

packages/cloudflare/test/durableobject.test.ts

@@ -46,8 +46,12 @@ describe('instrumentDurableObjectWithSentry', () => {
   });
 
   it('Instruments prototype methods without "sticking" to the options', () => {
+    const mockContext = {
+      waitUntil: vi.fn(),
+    } as any;
+    const mockEnv = {} as any; // Environment mock
     const initCore = vi.spyOn(SentryCore, 'initAndBind');
-    vi.spyOn(SentryCore, 'getClient').mockReturnValue(undefined);
+    const getClientSpy = vi.spyOn(SentryCore, 'getClient').mockReturnValue(undefined);
     const options = vi
       .fn()
       .mockReturnValueOnce({
@@ -59,8 +63,12 @@ describe('instrumentDurableObjectWithSentry', () => {
     const testClass = class {
       method() {}
     };
-    (Reflect.construct(instrumentDurableObjectWithSentry(options, testClass as any), []) as any).method();
-    (Reflect.construct(instrumentDurableObjectWithSentry(options, testClass as any), []) as any).method();
+    const instance1 = Reflect.construct(instrumentDurableObjectWithSentry(options, testClass as any), [mockContext, mockEnv]) as any;
+    instance1.method();
+    
+    const instance2 = Reflect.construct(instrumentDurableObjectWithSentry(options, testClass as any), [mockContext, mockEnv]) as any;
+    instance2.method();
+    
     expect(initCore).nthCalledWith(1, expect.any(Function), expect.objectContaining({ orgId: 1 }));
     expect(initCore).nthCalledWith(2, expect.any(Function), expect.objectContaining({ orgId: 2 }));
   });
@@ -83,7 +91,6 @@ describe('instrumentDurableObjectWithSentry', () => {
     };
     const instrumented = instrumentDurableObjectWithSentry(vi.fn(), testClass as any);
     const obj = Reflect.construct(instrumented, []);
-    expect(Object.getPrototypeOf(obj), 'Prototype is instrumented').not.toBe(testClass.prototype);
     for (const method_name of [
       'propertyFunction',
       'fetch',