feat: export single envvar span during pre app start (#77)

* export single pre app start envvar span

* remove old env var code, update schemas

Author: sohankshirsagar

package-lock.json

@@ -67,7 +67,7 @@
     "@types/supertest": "^6.0.2",
     "@typescript-eslint/eslint-plugin": "^6.0.0",
     "@typescript-eslint/parser": "^6.0.0",
-    "@use-tusk/drift-schemas": "^0.1.18",
+    "@use-tusk/drift-schemas": "^0.1.19",
     "ava": "^6.4.1",
     "axios": "^1.6.0",
     "eslint": "^8.57.1",

package.json

@@ -15,7 +15,6 @@ import {
   SendAlertRequest,
   InstrumentationVersionMismatchAlert,
   UnpatchedDependencyAlert,
-  EnvVarRequest,
 } from "@use-tusk/drift-schemas/core/communication";
 import { context, Context, SpanKind as OtSpanKind } from "@opentelemetry/api";
 import { Value } from "@use-tusk/drift-schemas/google/protobuf/struct";
@@ -229,17 +228,6 @@ try {
     // If CLI rejects -> CLI closes the connection and terminates the service.
   }
 
-  private getStackTrace(): string {
-    Error.stackTraceLimit = 100;
-    const s = new Error().stack || "";
-    Error.stackTraceLimit = 10;
-    return s
-      .split("\n")
-      .slice(2)
-      .filter((l) => !l.includes("ProtobufCommunicator"))
-      .join("\n");
-  }
-
   async requestMockAsync(mockRequest: MockRequestInput): Promise<MockResponseOutput> {
     const requestId = this.generateRequestId();
 
@@ -288,7 +276,7 @@ try {
   /**
    * Generic synchronous request handler that spawns a child process.
    * @param sdkMessage The SDK message to send
-   * @param filePrefix Prefix for temporary files (e.g., 'envvar', 'mock')
+   * @param filePrefix Prefix for temporary files (e.g., 'mock')
    * @param responseHandler Function to extract and return the desired response
    */
   private executeSyncRequest<TResponse>(
@@ -389,56 +377,6 @@ try {
     }
   }
 
-  /**
-   * Request environment variables from CLI synchronously using a child process.
-   * This blocks the main thread, so it should be used carefully.
-   * Similar to requestMockSync but for environment variables.
-   */
-  requestEnvVarsSync(traceTestServerSpanId: string): Record<string, string> {
-    const requestId = this.generateRequestId();
-
-    const envVarRequest = EnvVarRequest.create({
-      traceTestServerSpanId,
-    });
-
-    const sdkMessage = SDKMessage.create({
-      type: MessageType.ENV_VAR_REQUEST,
-      requestId: requestId,
-      payload: {
-        oneofKind: "envVarRequest",
-        envVarRequest,
-      },
-    });
-
-    logger.debug(
-      `[ProtobufCommunicator] Requesting env vars (sync) for trace: ${traceTestServerSpanId}`,
-    );
-
-    return this.executeSyncRequest(sdkMessage, "envvar", (cliMessage) => {
-      if (cliMessage.payload.oneofKind !== "envVarResponse") {
-        throw new Error(`Unexpected response type: ${cliMessage.type}`);
-      }
-
-      const envVarResponse = cliMessage.payload.envVarResponse;
-      if (!envVarResponse) {
-        throw new Error("No env var response received");
-      }
-
-      // Convert protobuf map to Record<string, string>
-      const envVars: Record<string, string> = {};
-      if (envVarResponse.envVars) {
-        Object.entries(envVarResponse.envVars).forEach(([key, value]) => {
-          envVars[key] = value;
-        });
-      }
-
-      logger.debug(
-        `[ProtobufCommunicator] Received env vars (sync), count: ${Object.keys(envVars).length}`,
-      );
-      return envVars;
-    });
-  }
-
   /**
    * This function uses a separate Node.js child process to communicate with the CLI over a socket.
    * The child process creates its own connection and event loop, allowing proper async socket handling.
@@ -737,33 +675,6 @@ try {
         });
       }
     }
-
-    if (message.payload.oneofKind === "envVarResponse") {
-      const envVarResponse = message.payload.envVarResponse;
-      logger.debug(`[ProtobufCommunicator] Received env var response for requestId: ${requestId}`);
-      const pendingRequest = this.pendingRequests.get(requestId);
-
-      if (!pendingRequest) {
-        logger.warn(
-          "[ProtobufCommunicator] received env var response for unknown request:",
-          requestId,
-        );
-        return;
-      }
-
-      this.pendingRequests.delete(requestId);
-
-      // Convert protobuf map to Record<string, string>
-      const envVars: Record<string, string> = {};
-      if (envVarResponse?.envVars) {
-        Object.entries(envVarResponse.envVars).forEach(([key, value]) => {
-          envVars[key] = value;
-        });
-      }
-
-      pendingRequest.resolve(envVars);
-      return;
-    }
   }
 
   /**

src/core/ProtobufCommunicator.ts

@@ -11,7 +11,6 @@ import {
   JsonwebtokenInstrumentation,
   DateInstrumentation,
   JwksRsaInstrumentation,
-  EnvInstrumentation,
   PostgresInstrumentation,
   Mysql2Instrumentation,
   IORedisInstrumentation,
@@ -23,13 +22,15 @@ import {
   MysqlInstrumentation,
 } from "../instrumentation/libraries";
 import { TdSpanExporter } from "./tracing/TdSpanExporter";
-import { trace, Tracer } from "@opentelemetry/api";
+import { trace, Tracer, SpanKind, SpanStatusCode } from "@opentelemetry/api";
 import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
 import { ProtobufCommunicator, MockRequestInput } from "./ProtobufCommunicator";
 import { BatchSpanProcessor } from "@opentelemetry/sdk-trace-node";
 import { CleanSpanData, TD_INSTRUMENTATION_LIBRARY_NAME } from "./types";
 import { TuskDriftInstrumentationModuleNames } from "./TuskDriftInstrumentationModuleNames";
 import { SDK_VERSION } from "../version";
+import { SpanUtils } from "./tracing/SpanUtils";
+import { PackageType } from "@use-tusk/drift-schemas/core/span";
 import {
   LogLevel,
   initializeGlobalLogger,
@@ -254,11 +255,6 @@ export class TuskDriftCore {
       mode: this.mode,
     });
 
-    new EnvInstrumentation({
-      enabled: this.config.recording?.enable_env_var_recording || false,
-      mode: this.mode,
-    });
-
     new PostgresInstrumentation({
       enabled: true,
       mode: this.mode,
@@ -317,7 +313,7 @@ export class TuskDriftCore {
       observableServiceId: this.config.service?.id,
       apiKey: this.initParams.apiKey,
       tuskBackendBaseUrl: this.config.tusk_api?.url || "https://api.usetusk.ai",
-      environment: this.initParams.env || "unknown",
+      environment: this.initParams.env,
       sdkVersion: SDK_VERSION,
       sdkInstanceId: this.generateSdkInstanceId(),
     });
@@ -350,6 +346,55 @@ export class TuskDriftCore {
     return `sdk-${originalDate.getTime()}-${Math.random().toString(36).substr(2, 9)}`;
   }
 
+  /**
+   * Creates a pre-app-start span containing a snapshot of all environment variables.
+   * Only runs in RECORD mode when env var recording is enabled.
+   */
+  private createEnvVarsSnapshot(): void {
+    // Only create snapshot in RECORD mode and if env var recording is enabled
+    if (this.mode !== TuskDriftMode.RECORD || !this.config.recording?.enable_env_var_recording) {
+      return;
+    }
+
+    try {
+      // Capture all env vars from process.env
+      const envVarsSnapshot: Record<string, string | undefined> = {};
+      for (const key of Object.keys(process.env)) {
+        envVarsSnapshot[key] = process.env[key];
+      }
+
+      logger.debug(
+        `Creating env vars snapshot with ${Object.keys(envVarsSnapshot).length} variables`,
+      );
+
+      // Create a span to hold the env vars snapshot
+      SpanUtils.createAndExecuteSpan(
+        this.mode,
+        () => {}, // No-op function since this is just a metadata snapshot
+        {
+          name: "ENV_VARS_SNAPSHOT",
+          kind: SpanKind.INTERNAL,
+          packageName: "process.env",
+          packageType: PackageType.UNSPECIFIED,
+          instrumentationName: "TuskDriftCore",
+          submodule: "env",
+          inputValue: {},
+          outputValue: {
+            ENV_VARS: envVarsSnapshot,
+          },
+          isPreAppStart: true,
+        },
+        (spanInfo) => {
+          // Span is created with metadata, just end it immediately
+          SpanUtils.endSpan(spanInfo.span, { code: SpanStatusCode.OK });
+          logger.debug(`Env vars snapshot span created: ${spanInfo.spanId}`);
+        },
+      );
+    } catch (error) {
+      logger.error("Failed to create env vars snapshot:", error);
+    }
+  }
+
   initialize(initParams: InitParams): void {
     // Initialize logging with provided level or default to 'info'
     initializeGlobalLogger({
@@ -487,6 +532,9 @@ export class TuskDriftCore {
     // which imports the gRPC exporter
     this.initializeTracing({ baseDirectory });
 
+    // Create env vars snapshot span (only in RECORD mode with env var recording enabled)
+    this.createEnvVarsSnapshot();
+
     this.initialized = true;
     logger.info("SDK initialized successfully");
   }
@@ -576,35 +624,6 @@ export class TuskDriftCore {
     }
   }
 
-  /**
-   * Request environment variables from CLI for a specific trace (synchronously).
-   * This blocks the main thread, so it should be used carefully.
-   */
-  requestEnvVarsSync(traceTestServerSpanId: string): Record<string, string> {
-    if (!this.isConnectedWithCLI) {
-      logger.error("Requesting sync env vars but CLI is not ready yet");
-      throw new Error("Requesting sync env vars but CLI is not ready yet");
-    }
-
-    if (!this.communicator || this.mode !== TuskDriftMode.REPLAY) {
-      logger.debug("Cannot request env vars: not in replay mode or no CLI connection");
-      return {};
-    }
-
-    try {
-      logger.debug(`Requesting env vars (sync) for trace: ${traceTestServerSpanId}`);
-      const envVars = this.communicator.requestEnvVarsSync(traceTestServerSpanId);
-      logger.debug(`Received env vars from CLI, count: ${Object.keys(envVars).length}`);
-      logger.debug(
-        `First 10 env vars: ${JSON.stringify(Object.keys(envVars).slice(0, 10), null, 2)}`,
-      );
-      return envVars;
-    } catch (error) {
-      logger.error(`[TuskDrift] Error requesting env vars from CLI:`, error);
-      return {};
-    }
-  }
-
   requestMockSync(mockRequest: MockRequestInput): {
     found: boolean;
     response?: unknown;
@@ -682,6 +701,10 @@ export class TuskDriftCore {
     return this.initParams;
   }
 
+  getEnvironment(): string | undefined {
+    return this.initParams.env;
+  }
+
   getTracer(): Tracer {
     return trace.getTracer(TD_INSTRUMENTATION_LIBRARY_NAME);
   }

src/core/TuskDrift.ts

@@ -1,7 +1,7 @@
 import { ReadableSpan } from "@opentelemetry/sdk-trace-base";
 import { SpanKind as OtSpanKind } from "@opentelemetry/api";
 import { JsonSchemaHelper, JsonSchemaType, JsonSchema } from "./JsonSchemaHelper";
-import { CleanSpanData, MetadataObject, TdSpanAttributes } from "../types";
+import { CleanSpanData, TdSpanAttributes } from "../types";
 import { PackageType, StatusCode } from "@use-tusk/drift-schemas/core/span";
 import { logger, OriginalGlobalUtils } from "../utils";
 
@@ -14,7 +14,7 @@ export class SpanTransformer {
    * Return type is derived from protobuf schema but uses clean JSON.
    * We use JSON because serialized protobuf is extremely verbose and not readable.
    */
-  static transformSpanToCleanJSON(span: ReadableSpan): CleanSpanData {
+  static transformSpanToCleanJSON(span: ReadableSpan, environment?: string): CleanSpanData {
     const isRootSpan = !span.parentSpanId || span.kind === OtSpanKind.SERVER;
 
     // Extract data from span attributes
@@ -68,7 +68,7 @@ export class SpanTransformer {
       } = JsonSchemaHelper.generateSchemaAndHash(outputData));
     }
 
-    let metadata: MetadataObject | undefined = undefined;
+    let metadata: Record<string, unknown> | undefined = undefined;
     if (attributes[TdSpanAttributes.METADATA]) {
       metadata = JSON.parse(attributes[TdSpanAttributes.METADATA] as string);
     }
@@ -99,6 +99,7 @@ export class SpanTransformer {
       submoduleName: submoduleName || "",
 
       packageType: (attributes[TdSpanAttributes.PACKAGE_TYPE] as PackageType) || undefined,
+      environment,
 
       inputValue: inputData,
       outputValue: outputData,

src/core/tracing/SpanTransformer.ts

@@ -10,7 +10,6 @@ import {
 } from "@opentelemetry/api";
 import {
   IS_PRE_APP_START_CONTEXT_KEY,
-  MetadataObject,
   REPLAY_TRACE_ID_CONTEXT_KEY,
   SPAN_KIND_CONTEXT_KEY,
   STOP_RECORDING_CHILD_SPANS_CONTEXT_KEY,
@@ -47,9 +46,10 @@ export interface SpanExecutorOptions {
   instrumentationName: string;
   submodule: string;
   inputValue: Record<string, unknown>;
+  outputValue?: Record<string, unknown>;
   isPreAppStart: boolean;
   inputSchemaMerges?: SchemaMerges;
-  metadata?: MetadataObject;
+  metadata?: Record<string, unknown>;
   stopRecordingChildSpans?: boolean;
 }
 
@@ -64,7 +64,7 @@ export interface AddSpanAttributesOptions {
   outputValue?: Record<string, unknown>;
   inputSchemaMerges?: SchemaMerges;
   outputSchemaMerges?: SchemaMerges;
-  metadata?: MetadataObject;
+  metadata?: Record<string, unknown>;
   transformMetadata?: {
     transformed: boolean;
     actions: Array<{
@@ -181,6 +181,7 @@ export class SpanUtils {
       packageType,
       submodule,
       inputValue,
+      outputValue,
       inputSchemaMerges,
       isPreAppStart,
       metadata,
@@ -202,6 +203,9 @@ export class SpanUtils {
           [TdSpanAttributes.INSTRUMENTATION_NAME]: instrumentationName,
           [TdSpanAttributes.PACKAGE_TYPE]: packageType,
           [TdSpanAttributes.INPUT_VALUE]: createSpanInputValue(inputValue),
+          ...(outputValue && {
+            [TdSpanAttributes.OUTPUT_VALUE]: JSON.stringify(outputValue),
+          }),
           [TdSpanAttributes.IS_PRE_APP_START]: isPreAppStart,
           ...(inputSchemaMerges && {
             [TdSpanAttributes.INPUT_SCHEMA_MERGES]: JSON.stringify(inputSchemaMerges),