Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions package-lock.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
91 changes: 1 addition & 90 deletions src/core/ProtobufCommunicator.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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>(
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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;
}
}

/**
Expand Down
97 changes: 60 additions & 37 deletions src/core/TuskDrift.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@ import {
JsonwebtokenInstrumentation,
DateInstrumentation,
JwksRsaInstrumentation,
EnvInstrumentation,
PostgresInstrumentation,
Mysql2Instrumentation,
IORedisInstrumentation,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(),
});
Expand Down Expand Up @@ -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({
Expand Down Expand Up @@ -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");
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Expand Down
7 changes: 4 additions & 3 deletions src/core/tracing/SpanTransformer.ts
Original file line number Diff line number Diff line change
@@ -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";

Expand All @@ -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
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -99,6 +99,7 @@ export class SpanTransformer {
submoduleName: submoduleName || "",

packageType: (attributes[TdSpanAttributes.PACKAGE_TYPE] as PackageType) || undefined,
environment,

inputValue: inputData,
outputValue: outputData,
Expand Down
10 changes: 7 additions & 3 deletions src/core/tracing/SpanUtils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}

Expand All @@ -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<{
Expand Down Expand Up @@ -181,6 +181,7 @@ export class SpanUtils {
packageType,
submodule,
inputValue,
outputValue,
inputSchemaMerges,
isPreAppStart,
metadata,
Expand All @@ -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),
Expand Down
Loading