@inference-net/otel-cf-workers 2.0.2 → 2.0.3

package/dist/index.js

@@ -96,7 +96,7 @@ function passthroughGet(target, prop, thisArg) {
   }
 }
 
-const PACKAGE_VERSION = "2.0.2";
+const PACKAGE_VERSION = "2.0.3";
 const PACKAGE_NAME = "@inference-net/otel-cf-workers";
 const ATTR_CLOUDFLARE_COLO = "cloudflare.colo";
 const ATTR_CLOUDFLARE_RAY_ID = "cloudflare.ray_id";
@@ -2924,6 +2924,42 @@ function instrumentSQLStorage(sql) {
   return wrap(sql, sqlHandler);
 }
 
+const RPC_CONTEXT_MARKER = "__otel_rpc_ctx__";
+function isRpcContextCarrier(obj) {
+  return typeof obj === "object" && obj !== null && RPC_CONTEXT_MARKER in obj && obj[RPC_CONTEXT_MARKER] === true;
+}
+function injectRpcContext(ctx = context.active()) {
+  const carrier = {
+    [RPC_CONTEXT_MARKER]: true,
+    headers: {}
+  };
+  propagation.inject(ctx, carrier.headers, {
+    set: (headers, key, value) => {
+      headers[key] = typeof value === "string" ? value : String(value);
+    }
+  });
+  return carrier;
+}
+function extractRpcContext(carrier) {
+  return propagation.extract(context.active(), carrier.headers, {
+    get(headers, key) {
+      return headers[key] || void 0;
+    },
+    keys(headers) {
+      return Object.keys(headers);
+    }
+  });
+}
+function extractAndRemoveRpcContext(args) {
+  if (args.length > 0 && isRpcContextCarrier(args[0])) {
+    const carrier = args[0];
+    const extractedContext = extractRpcContext(carrier);
+    const cleanedArgs = args.slice(1);
+    return [extractedContext, cleanedArgs];
+  }
+  return [void 0, args];
+}
+
 function instrumentRpcMethod(method, methodName, nsName, stubId) {
   const tracer = trace.getTracer("do_rpc_client");
   const rpcHandler = {
@@ -2944,7 +2980,9 @@ function instrumentRpcMethod(method, methodName, nsName, stubId) {
       const spanName = `RPC ${nsName}.${methodName}`;
       return tracer.startActiveSpan(spanName, options, async (span) => {
         try {
-          const result = await Reflect.apply(target, thisArg, argArray);
+          const contextCarrier = injectRpcContext(context.active());
+          const argsWithContext = [contextCarrier, ...argArray];
+          const result = await Reflect.apply(target, thisArg, argsWithContext);
           span.setStatus({ code: SpanStatusCode.OK });
           return result;
         } catch (error) {
@@ -3111,7 +3149,9 @@ function instrumentRpcHandlerMethod(fn, methodName, initialiser, env, id) {
     async apply(target, thisArg, argArray) {
       thisArg = unwrap(thisArg);
       const config = initialiser(env, "do-rpc");
-      const context$1 = setConfig(config);
+      const [extractedContext, cleanedArgs] = extractAndRemoveRpcContext(argArray);
+      let context$1 = extractedContext || context.active();
+      context$1 = setConfig(config, context$1);
       const tracer = trace.getTracer("do_rpc_handler");
       const doName = id.name || id.toString();
       const executeRpc = async () => {
@@ -3131,7 +3171,7 @@ function instrumentRpcHandlerMethod(fn, methodName, initialiser, env, id) {
         return tracer.startActiveSpan(spanName, options, async (span) => {
           try {
             const bound = target.bind(thisArg);
-            const result = await bound.apply(thisArg, argArray);
+            const result = await bound.apply(thisArg, cleanedArgs);
             span.setStatus({ code: SpanStatusCode.OK });
             return result;
           } catch (error) {