improve trace flag handling and better internal host checking

Author: ericallam

packages/core/src/v3/otel/tracingSDK.ts

@@ -1,4 +1,10 @@
-import { DiagConsoleLogger, DiagLogLevel, TracerProvider, diag } from "@opentelemetry/api";
+import {
+  DiagConsoleLogger,
+  DiagLogLevel,
+  TraceFlags,
+  TracerProvider,
+  diag,
+} from "@opentelemetry/api";
 import { logs } from "@opentelemetry/api-logs";
 import { TraceState } from "@opentelemetry/core";
 import { OTLPLogExporter } from "@opentelemetry/exporter-logs-otlp-http";
@@ -293,13 +299,10 @@ class ExternalSpanExporterWrapper {
     private underlyingExporter: SpanExporter,
     private externalTraceId: string,
     private externalTraceContext:
-      | { traceId: string; spanId: string; tracestate?: string; traceFlags?: string }
+      | { traceId: string; spanId: string; traceFlags: number; tracestate?: string }
       | undefined
   ) {
-    this._isExternallySampled =
-      typeof externalTraceContext?.traceFlags === "string"
-        ? externalTraceContext.traceFlags === "01"
-        : true;
+    this._isExternallySampled = isTraceFlagSampled(externalTraceContext?.traceFlags);
   }
 
   private transformSpan(span: ReadableSpan): ReadableSpan | undefined {
@@ -383,13 +386,10 @@ class ExternalLogRecordExporterWrapper {
     private underlyingExporter: LogRecordExporter,
     private externalTraceId: string,
     private externalTraceContext:
-      | { traceId: string; spanId: string; tracestate?: string; traceFlags?: string }
+      | { traceId: string; spanId: string; tracestate?: string; traceFlags: number }
       | undefined
   ) {
-    this._isExternallySampled =
-      typeof externalTraceContext?.traceFlags === "string"
-        ? externalTraceContext.traceFlags === "01"
-        : true;
+    this._isExternallySampled = isTraceFlagSampled(externalTraceContext?.traceFlags);
   }
 
   export(logs: any[], resultCallback: (result: any) => void): void {
@@ -455,25 +455,42 @@ function isSpanInternalOnly(span: ReadableSpan): boolean {
 
   const httpUrl = span.attributes[SEMATTRS_HTTP_URL] ?? span.attributes["url.full"];
 
-  if (typeof httpUrl === "string" && httpUrl.includes("https://api.trigger.dev")) {
-    return true;
-  }
+  const url = safeParseUrl(httpUrl);
 
-  if (typeof httpUrl === "string" && httpUrl.includes("https://billing.trigger.dev")) {
-    return true;
+  if (!url) {
+    return false;
   }
 
-  if (typeof httpUrl === "string" && httpUrl.includes("https://cloud.trigger.dev")) {
-    return true;
+  const internalHosts = [
+    "api.trigger.dev",
+    "billing.trigger.dev",
+    "cloud.trigger.dev",
+    "engine.trigger.dev",
+    "platform.trigger.dev",
+  ];
+
+  return (
+    internalHosts.some((host) => url.hostname.includes(host)) ||
+    url.pathname.includes("/api/v1/usage/ingest")
+  );
+}
+
+function safeParseUrl(url: unknown): URL | undefined {
+  if (typeof url !== "string") {
+    return undefined;
   }
 
-  if (typeof httpUrl === "string" && httpUrl.includes("https://engine.trigger.dev")) {
-    return true;
+  try {
+    return new URL(url);
+  } catch (e) {
+    return undefined;
   }
+}
 
-  if (typeof httpUrl === "string" && httpUrl.includes("/api/v1/usage/ingest")) {
+function isTraceFlagSampled(traceFlags?: number): boolean {
+  if (typeof traceFlags !== "number") {
     return true;
   }
 
-  return false;
+  return (traceFlags & TraceFlags.SAMPLED) === TraceFlags.SAMPLED;
 }

packages/core/src/v3/traceContext/manager.ts

@@ -1,5 +1,6 @@
 import { Context, context, propagation, trace, TraceFlags } from "@opentelemetry/api";
 import { TraceContextManager } from "./types.js";
+import { parseTraceParent } from "@opentelemetry/core";
 
 export class StandardTraceContextManager implements TraceContextManager {
   public traceContext: Record<string, unknown> = {};
@@ -39,12 +40,7 @@ export class StandardTraceContextManager implements TraceContextManager {
     const spanContext = {
       traceId: externalTraceContext.traceId,
       spanId: currentSpanContext.spanId,
-      traceFlags:
-        typeof externalTraceContext.traceFlags === "string"
-          ? externalTraceContext.traceFlags === "01"
-            ? TraceFlags.SAMPLED
-            : TraceFlags.NONE
-          : TraceFlags.SAMPLED,
+      traceFlags: externalTraceContext.traceFlags,
       isRemote: true,
     };
 
@@ -65,16 +61,16 @@ function extractExternalTraceContext(traceContext: unknown) {
       : undefined;
 
   if ("traceparent" in traceContext && typeof traceContext.traceparent === "string") {
-    const [version, traceId, spanId, traceFlags] = traceContext.traceparent.split("-");
+    const externalSpanContext = parseTraceParent(traceContext.traceparent);
 
-    if (!traceId || !spanId) {
+    if (!externalSpanContext) {
       return undefined;
     }
 
     return {
-      traceId,
-      spanId,
-      traceFlags,
+      traceId: externalSpanContext.traceId,
+      spanId: externalSpanContext.spanId,
+      traceFlags: externalSpanContext.traceFlags,
       tracestate: tracestate,
     };
   }

packages/core/src/v3/traceContext/types.ts

@@ -8,7 +8,7 @@ export interface TraceContextManager {
     | {
         traceId: string;
         spanId: string;
-        traceFlags?: string;
+        traceFlags: number;
         tracestate?: string;
       }
     | undefined;