fix(cloudflare): copy execution context in durable objects and handlers

.gitignore

@@ -61,3 +61,6 @@ packages/gatsby/gatsby-node.d.ts
 # intellij
 *.iml
 /**/.wrangler/*
+
+#junit reports
+packages/**/*.junit.xml

dev-packages/cloudflare-integration-tests/runner.ts

@@ -86,7 +86,7 @@ export function createRunner(...paths: string[]) {
       }
       return this;
     },
-    start: function (): StartResult {
+    start: function (signal?: AbortSignal): StartResult {
       const { resolve, reject, promise: isComplete } = deferredPromise(cleanupChildProcesses);
       const expectedEnvelopeCount = expectedEnvelopes.length;
 
@@ -155,7 +155,7 @@ export function createRunner(...paths: string[]) {
               '--var',
               `SENTRY_DSN:http://public@localhost:${mockServerPort}/1337`,
             ],
-            { stdio },
+            { stdio, signal },
           );
 
           CLEANUP_STEPS.add(() => {

dev-packages/cloudflare-integration-tests/suites/basic/test.ts

@@ -2,7 +2,7 @@ import { expect, it } from 'vitest';
 import { eventEnvelope } from '../../expect';
 import { createRunner } from '../../runner';
 
-it('Basic error in fetch handler', async () => {
+it('Basic error in fetch handler', async ({ signal }) => {
   const runner = createRunner(__dirname)
     .expect(
       eventEnvelope({
@@ -26,7 +26,7 @@ it('Basic error in fetch handler', async () => {
         },
       }),
     )
-    .start();
+    .start(signal);
   await runner.makeRequest('get', '/', { expectError: true });
   await runner.completed();
 });

dev-packages/cloudflare-integration-tests/suites/tracing/anthropic-ai/test.ts

@@ -6,7 +6,7 @@ import { createRunner } from '../../../runner';
 // want to test that the instrumentation does not break in our
 // cloudflare SDK.
 
-it('traces a basic message creation request', async () => {
+it('traces a basic message creation request', async ({ signal }) => {
   const runner = createRunner(__dirname)
     .ignore('event')
     .expect(envelope => {
@@ -35,7 +35,7 @@ it('traces a basic message creation request', async () => {
         ]),
       );
     })
-    .start();
+    .start(signal);
   await runner.makeRequest('get', '/');
   await runner.completed();
 });

dev-packages/cloudflare-integration-tests/suites/tracing/durableobject/test.ts

@@ -1,7 +1,7 @@
 import { expect, it } from 'vitest';
 import { createRunner } from '../../../runner';
 
-it('traces a durable object method', async () => {
+it('traces a durable object method', async ({ signal }) => {
   const runner = createRunner(__dirname)
     .expect(envelope => {
       const transactionEvent = envelope[1]?.[0]?.[1];
@@ -21,7 +21,7 @@ it('traces a durable object method', async () => {
         }),
       );
     })
-    .start();
+    .start(signal);
   await runner.makeRequest('get', '/hello');
   await runner.completed();
 });

dev-packages/cloudflare-integration-tests/suites/tracing/openai/test.ts

@@ -6,7 +6,7 @@ import { createRunner } from '../../../runner';
 // want to test that the instrumentation does not break in our
 // cloudflare SDK.
 
-it('traces a basic chat completion request', async () => {
+it('traces a basic chat completion request', async ({ signal }) => {
   const runner = createRunner(__dirname)
     .ignore('event')
     .expect(envelope => {
@@ -37,7 +37,7 @@ it('traces a basic chat completion request', async () => {
         ]),
       );
     })
-    .start();
+    .start(signal);
   await runner.makeRequest('get', '/');
   await runner.completed();
 });

packages/cloudflare/src/durableobject.ts

@@ -18,6 +18,7 @@ import { isInstrumented, markAsInstrumented } from './instrument';
 import { getFinalOptions } from './options';
 import { wrapRequestHandler } from './request';
 import { init } from './sdk';
+import { copyExecutionContext } from './utils/copyExecutionContext';
 
 type MethodWrapperOptions = {
   spanName?: string;
@@ -192,8 +193,9 @@ export function instrumentDurableObjectWithSentry<
   C extends new (state: DurableObjectState, env: E) => T,
 >(optionsCallback: (env: E) => CloudflareOptions, DurableObjectClass: C): C {
   return new Proxy(DurableObjectClass, {
-    construct(target, [context, env]) {
+    construct(target, [ctx, env]) {
       setAsyncLocalStorageAsyncContextStrategy();
+      const context = copyExecutionContext(ctx);
 
       const options = getFinalOptions(optionsCallback(env), env);
 

packages/cloudflare/src/handler.ts

@@ -15,6 +15,7 @@ import { getFinalOptions } from './options';
 import { wrapRequestHandler } from './request';
 import { addCloudResourceContext } from './scope-utils';
 import { init } from './sdk';
+import { copyExecutionContext } from './utils/copyExecutionContext';
 
 /**
  * Wrapper for Cloudflare handlers.
@@ -38,7 +39,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('fetch' in handler && typeof handler.fetch === 'function' && !isInstrumented(handler.fetch)) {
       handler.fetch = new Proxy(handler.fetch, {
         apply(target, thisArg, args: Parameters<ExportedHandlerFetchHandler<Env, CfHostMetadata>>) {
-          const [request, env, context] = args;
+          const [request, env, ctx] = args;
+          const context = copyExecutionContext(ctx);
+          args[2] = context;
 
           const options = getFinalOptions(optionsCallback(env), env);
 
@@ -72,7 +75,10 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('scheduled' in handler && typeof handler.scheduled === 'function' && !isInstrumented(handler.scheduled)) {
       handler.scheduled = new Proxy(handler.scheduled, {
         apply(target, thisArg, args: Parameters<ExportedHandlerScheduledHandler<Env>>) {
-          const [event, env, context] = args;
+          const [event, env, ctx] = args;
+          const context = copyExecutionContext(ctx);
+          args[2] = context;
+
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
             const waitUntil = context.waitUntil.bind(context);
@@ -115,7 +121,10 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('email' in handler && typeof handler.email === 'function' && !isInstrumented(handler.email)) {
       handler.email = new Proxy(handler.email, {
         apply(target, thisArg, args: Parameters<EmailExportedHandler<Env>>) {
-          const [emailMessage, env, context] = args;
+          const [emailMessage, env, ctx] = args;
+          const context = copyExecutionContext(ctx);
+          args[2] = context;
+
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
             const waitUntil = context.waitUntil.bind(context);
@@ -156,7 +165,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('queue' in handler && typeof handler.queue === 'function' && !isInstrumented(handler.queue)) {
       handler.queue = new Proxy(handler.queue, {
         apply(target, thisArg, args: Parameters<ExportedHandlerQueueHandler<Env, QueueHandlerMessage>>) {
-          const [batch, env, context] = args;
+          const [batch, env, ctx] = args;
+          const context = copyExecutionContext(ctx);
+          args[2] = context;
 
           return withIsolationScope(isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);
@@ -206,7 +217,9 @@ export function withSentry<Env = unknown, QueueHandlerMessage = unknown, CfHostM
     if ('tail' in handler && typeof handler.tail === 'function' && !isInstrumented(handler.tail)) {
       handler.tail = new Proxy(handler.tail, {
         apply(target, thisArg, args: Parameters<ExportedHandlerTailHandler<Env>>) {
-          const [, env, context] = args;
+          const [, env, ctx] = args;
+          const context = copyExecutionContext(ctx);
+          args[2] = context;
 
           return withIsolationScope(async isolationScope => {
             const options = getFinalOptions(optionsCallback(env), env);

packages/cloudflare/src/utils/copyExecutionContext.ts

@@ -0,0 +1,69 @@
+import { type DurableObjectState, type ExecutionContext } from '@cloudflare/workers-types';
+
+type ContextType = ExecutionContext | DurableObjectState;
+type OverridesStore<T extends ContextType> = Map<keyof T, (...args: unknown[]) => unknown>;
+
+/**
+ * Creates a new copy of the given execution context, optionally overriding methods.
+ *
+ * @param {ContextType|void} ctx - The execution context to be copied. Can be of type `ContextType` or `void`.
+ * @return {ContextType|void} A new execution context with the same properties and overridden methods if applicable.
+ */
+export function copyExecutionContext<T extends ContextType>(ctx: T): T {
+  if (!ctx) return ctx;
+
+  const overrides: OverridesStore<T> = new Map();
+  const contextPrototype = Object.getPrototypeOf(ctx);
+  const prototypeMethodNames = Object.getOwnPropertyNames(contextPrototype) as unknown as (keyof T)[];
+  const ownPropertyNames = Object.getOwnPropertyNames(ctx) as unknown as (keyof T)[];
+  const instrumented = new Set<unknown>(['constructor']);
+  const descriptors = [...ownPropertyNames, ...prototypeMethodNames].reduce((prevDescriptors, methodName) => {
+    if (instrumented.has(methodName)) return prevDescriptors;
+    if (typeof ctx[methodName] !== 'function') return prevDescriptors;
+    instrumented.add(methodName);
+    const overridableDescriptor = makeOverridableDescriptor(overrides, ctx, methodName);
+    return {
+      ...prevDescriptors,
+      [methodName]: overridableDescriptor,
+    };
+  }, {});
+
+  return Object.create(ctx, descriptors);
+}
+
+/**
+ * Creates a property descriptor that allows overriding of a method on the given context object.
+ *
+ * This descriptor supports property overriding with functions only. It delegates method calls to
+ * the provided store if an override exists or to the original method on the context otherwise.
+ *
+ * @param {OverridesStore<ContextType>} store - The storage for overridden methods specific to the context type.
+ * @param {ContextType} ctx - The context object that contains the method to be overridden.
+ * @param {keyof ContextType} method - The method on the context object to create the overridable descriptor for.
+ * @return {PropertyDescriptor} A property descriptor enabling the overriding of the specified method.
+ */
+function makeOverridableDescriptor<T extends ContextType>(
+  store: OverridesStore<T>,
+  ctx: T,
+  method: keyof T,
+): PropertyDescriptor {
+  return {
+    configurable: true,
+    enumerable: true,
+    set: newValue => {
+      if (typeof newValue == 'function') {
+        store.set(method, newValue);
+        return;
+      }
+      Reflect.set(ctx, method, newValue);
+    },
+
+    get: () => {
+      if (store.has(method)) return store.get(method);
+      const methodFunction = Reflect.get(ctx, method);
+      if (typeof methodFunction !== 'function') return methodFunction;
+      // We should do bind() to make sure that the method is bound to the context object - otherwise it will not work
+      return methodFunction.bind(ctx);
+    },
+  };
+}

packages/cloudflare/src/workflows.ts

@@ -22,6 +22,7 @@ import { setAsyncLocalStorageAsyncContextStrategy } from './async';
 import type { CloudflareOptions } from './client';
 import { addCloudResourceContext } from './scope-utils';
 import { init } from './sdk';
+import { copyExecutionContext } from './utils/copyExecutionContext';
 
 const UUID_REGEX = /^[0-9a-f]{8}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{4}-?[0-9a-f]{12}$/i;
 
@@ -157,6 +158,9 @@ export function instrumentWorkflowWithSentry<
   return new Proxy(WorkFlowClass, {
     construct(target: C, args: [ctx: ExecutionContext, env: E], newTarget) {
       const [ctx, env] = args;
+      const context = copyExecutionContext(ctx);
+      args[0] = context;
+
       const options = optionsCallback(env);
       const instance = Reflect.construct(target, args, newTarget) as T;
       return new Proxy(instance, {
@@ -179,10 +183,10 @@ export function instrumentWorkflowWithSentry<
                     return await obj.run.call(
                       obj,
                       event,
-                      new WrappedWorkflowStep(event.instanceId, ctx, options, step),
+                      new WrappedWorkflowStep(event.instanceId, context, options, step),
                     );
                   } finally {
-                    ctx.waitUntil(flush(2000));
+                    context.waitUntil(flush(2000));
                   }
                 });
               });

packages/cloudflare/test/copy-execution-context.test.ts

@@ -0,0 +1,56 @@
+import { type Mocked, describe, expect, it, vi } from 'vitest';
+import { copyExecutionContext } from '../src/utils/copyExecutionContext';
+
+describe('Copy of the execution context', () => {
+  describe.for([
+    'waitUntil',
+    'passThroughOnException',
+    'acceptWebSocket',
+    'blockConcurrencyWhile',
+    'getWebSockets',
+    'arbitraryMethod',
+    'anythingElse',
+  ])('%s', method => {
+    it('Override without changing original', async () => {
+      const context = {
+        [method]: vi.fn(),
+      } as any;
+      const copy = copyExecutionContext(context);
+      copy[method] = vi.fn();
+      expect(context[method]).not.toBe(copy[method]);
+    });
+
+    it('Overridden method was called', async () => {
+      const context = {
+        [method]: vi.fn(),
+      } as any;
+      const copy = copyExecutionContext(context);
+      const overridden = vi.fn();
+      copy[method] = overridden;
+      copy[method]();
+      expect(overridden).toBeCalled();
+      expect(context[method]).not.toBeCalled();
+    });
+  });
+
+  it('No side effects', async () => {
+    const context = makeExecutionContextMock();
+    expect(() => copyExecutionContext(Object.freeze(context))).not.toThrow(
+      /Cannot define property \w+, object is not extensible/,
+    );
+  });
+  it('Respects symbols', async () => {
+    const s = Symbol('test');
+    const context = makeExecutionContextMock<ExecutionContext & { [s]: unknown }>();
+    context[s] = {};
+    const copy = copyExecutionContext(context);
+    expect(copy[s]).toBe(context[s]);
+  });
+});
+
+function makeExecutionContextMock<T extends ExecutionContext>() {
+  return {
+    waitUntil: vi.fn(),
+    passThroughOnException: vi.fn(),
+  } as unknown as Mocked<T>;
+}