Integrate copyExecutionContext across handlers and workflows

Replaced direct usage of `ExecutionContext` with `copyExecutionContext` for consistency and improved flexibility. Applied changes across handlers, durable objects, and workflows to enhance context management and method overriding.

Author: cod1k

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

@@ -14,10 +14,13 @@ export function copyExecutionContext<T extends ContextType>(ctx: T): T {
 
   const overrides: OverridesStore<T> = new Map();
   const contextPrototype = Object.getPrototypeOf(ctx);
-  const methodNames = Object.getOwnPropertyNames(contextPrototype) as unknown as (keyof T)[];
-  const descriptors = methodNames.reduce((prevDescriptors, methodName) => {
-    if (methodName === 'constructor') return prevDescriptors;
+  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,
@@ -48,9 +51,11 @@ function makeOverridableDescriptor<T extends ContextType>(
     configurable: true,
     enumerable: true,
     set: newValue => {
-      if (typeof newValue !== 'function') throw new Error('Cannot override non-function');
-      store.set(method, newValue);
-      return true;
+      if (typeof newValue == 'function') {
+        store.set(method, newValue);
+        return;
+      }
+      Reflect.set(ctx, method, newValue);
     },
 
     get: () => {

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));
                   }
                 });
               });