chore: use JsonSchema proto

src/core/ProtobufCommunicator.ts

@@ -384,8 +384,8 @@ export class ProtobufCommunicator {
       submoduleName: span.submoduleName || "",
       inputValue: toStruct(span.inputValue),
       outputValue: toStruct(span.outputValue),
-      inputSchema: toStruct(span.inputSchema),
-      outputSchema: toStruct(span.outputSchema),
+      inputSchema: span.inputSchema,
+      outputSchema: span.outputSchema,
       inputSchemaHash: span.inputSchemaHash || "",
       outputSchemaHash: span.outputSchemaHash || "",
       inputValueHash: span.inputValueHash || "",

src/core/tracing/JsonSchemaHelper.ts

@@ -1,17 +1,11 @@
 import * as crypto from "crypto";
 import { logger } from "../utils/logger";
-
-export enum JsonSchemaType {
-  NUMBER = "NUMBER",
-  STRING = "STRING",
-  BOOLEAN = "BOOLEAN",
-  NULL = "NULL",
-  UNDEFINED = "UNDEFINED",
-  OBJECT = "OBJECT",
-  ORDERED_LIST = "ORDERED_LIST",
-  UNORDERED_LIST = "UNORDERED_LIST",
-  FUNCTION = "FUNCTION",
-}
+import {
+  JsonSchemaType,
+  EncodingType,
+  DecodedType,
+  type JsonSchema,
+} from "@use-tusk/drift-schemas/core/json_schema";
 
 // Standardized schema type mapping
 const jsToJsonSchemaTypeMapping = {
@@ -47,45 +41,8 @@
   ["Arguments"]: JsonSchemaType.ORDERED_LIST,
 } as const;
 
-export enum EncodingType {
-  BASE64 = "BASE64",
-}
-
-export enum DecodedType {
-  JSON = "JSON",
-  HTML = "HTML",
-  CSS = "CSS",
-  JAVASCRIPT = "JAVASCRIPT",
-  XML = "XML",
-  YAML = "YAML",
-  MARKDOWN = "MARKDOWN",
-  CSV = "CSV",
-  SQL = "SQL",
-  GRAPHQL = "GRAPHQL",
-  PLAIN_TEXT = "PLAIN_TEXT",
-  FORM_DATA = "FORM_DATA",
-  MULTIPART_FORM = "MULTIPART_FORM",
-  PDF = "PDF",
-  AUDIO = "AUDIO",
-  VIDEO = "VIDEO",
-  GZIP = "GZIP",
-  BINARY = "BINARY",
-  JPEG = "JPEG",
-  PNG = "PNG",
-  GIF = "GIF",
-  WEBP = "WEBP",
-  SVG = "SVG",
-  ZIP = "ZIP",
-}
-
-export interface JsonSchema {
-  type: JsonSchemaType;
-  properties?: Record<string, JsonSchema>;
-  items?: JsonSchema | null;
-  encoding?: EncodingType;
-  decodedType?: DecodedType;
-  matchImportance?: number; // Should be between 0 and 1, 0 being the lowest importance and 1 being the highest importance
-}
+// Re-export proto types for convenience
+export { JsonSchemaType, EncodingType, DecodedType, type JsonSchema };
 
 // The following types can be merged with the generated schema to provide additional information
 // This is set in the instrumentation layer and merged with the generated schema
@@ -212,39 +169,49 @@
 
   /**
    * Generate schema from data object using standardized types
+   * 
+   * Note: We properties always exists on JsonSchema because proto3 maps cannot be marked optional.
+   * The JSON data is a bit inefficient because of this, but the easiest way to handle this is to keep it for now.
    */
   static generateSchema(data: any, schemaMerges?: SchemaMerges): JsonSchema {
     if (data === null) {
-      return { type: jsToJsonSchemaTypeMapping["null"] };
+      return { type: jsToJsonSchemaTypeMapping["null"], properties: {} };
     }
 
     if (data === undefined) {
-      return { type: jsToJsonSchemaTypeMapping["undefined"] };
+      return { type: jsToJsonSchemaTypeMapping["undefined"], properties: {} };
     }
 
     const detailedType = JsonSchemaHelper.getDetailedType(data);
 
     if (detailedType === JsonSchemaType.ORDERED_LIST) {
       if (Array.isArray(data) && data.length === 0) {
-        return { type: JsonSchemaType.ORDERED_LIST, items: null };
+        return { type: JsonSchemaType.ORDERED_LIST, properties: {} };
+      }
+      const items = Array.isArray(data) && data.length > 0 ? JsonSchemaHelper.generateSchema(data[0]) : undefined;
+      if (items !== undefined) {
+        return {
+          type: JsonSchemaType.ORDERED_LIST,
+          items,
+          properties: {},
+        };
       }
-      return {
-        type: JsonSchemaType.ORDERED_LIST,
-        items:
-          Array.isArray(data) && data.length > 0 ? JsonSchemaHelper.generateSchema(data[0]) : null,
-      };
+      return { type: JsonSchemaType.ORDERED_LIST, properties: {} };
     }
 
     if (detailedType === JsonSchemaType.UNORDERED_LIST) {
       // Handle Set objects
       if (data instanceof Set) {
         const firstItem = data.size > 0 ? data.values().next().value : null;
-        return {
-          type: JsonSchemaType.UNORDERED_LIST,
-          items: firstItem !== null ? JsonSchemaHelper.generateSchema(firstItem) : null,
-        };
+        if (firstItem !== null) {
+          return {
+            type: JsonSchemaType.UNORDERED_LIST,
+            items: JsonSchemaHelper.generateSchema(firstItem),
+            properties: {},
+          };
+        }
       }
-      return { type: JsonSchemaType.UNORDERED_LIST, items: null };
+      return { type: JsonSchemaType.UNORDERED_LIST, properties: {} };
     }
 
     if (detailedType === JsonSchemaType.OBJECT) {
@@ -253,44 +220,40 @@
       // Handle Map objects
       if (data instanceof Map) {
         data.forEach((value, key) => {
-          if (schema.properties) {
-            const keyString = String(key);
-            const generatedSchema = JsonSchemaHelper.generateSchema(value);
-
-            // Check for schema override for this key
-            if (schemaMerges && schemaMerges[keyString]) {
-              schema.properties[keyString] = JsonSchemaHelper.mergeSchemaWithMerges(
-                generatedSchema,
-                schemaMerges[keyString],
-              );
-            } else {
-              schema.properties[keyString] = generatedSchema;
-            }
+          const keyString = String(key);
+          const generatedSchema = JsonSchemaHelper.generateSchema(value);
+
+          // Check for schema override for this key
+          if (schemaMerges && schemaMerges[keyString]) {
+            schema.properties[keyString] = JsonSchemaHelper.mergeSchemaWithMerges(
+              generatedSchema,
+              schemaMerges[keyString],
+            );
+          } else {
+            schema.properties[keyString] = generatedSchema;
           }
         });
       } else if (typeof data === "object") {
         Object.keys(data).forEach((key) => {
-          if (schema.properties) {
-            const generatedSchema = JsonSchemaHelper.generateSchema(data[key]);
-
-            // Check for schema override for this key
-            if (schemaMerges && schemaMerges[key]) {
-              schema.properties[key] = JsonSchemaHelper.mergeSchemaWithMerges(
-                generatedSchema,
-                schemaMerges[key],
-              );
-            } else {
-              schema.properties[key] = generatedSchema;
-            }
+          const generatedSchema = JsonSchemaHelper.generateSchema(data[key]);
+
+          // Check for schema override for this key
+          if (schemaMerges && schemaMerges[key]) {
+            schema.properties[key] = JsonSchemaHelper.mergeSchemaWithMerges(
+              generatedSchema,
+              schemaMerges[key],
+            );
+          } else {
+            schema.properties[key] = generatedSchema;
           }
         });
       }
 
       return schema;
     }
 
-    // For primitive types, return the standardized type
-    return { type: detailedType };
+    // For primitive types, return just the type
+    return { type: detailedType, properties: {} };
   }
 
   /**

src/core/tracing/SpanTransformer.ts

@@ -1,6 +1,6 @@
 import { ReadableSpan } from "@opentelemetry/sdk-trace-base";
 import { SpanKind as OtSpanKind } from "@opentelemetry/api";
-import { JsonSchemaHelper, JsonSchema, JsonSchemaType } from "./JsonSchemaHelper";
+import { JsonSchemaHelper, JsonSchemaType, JsonSchema } from "./JsonSchemaHelper";
 import { CleanSpanData, MetadataObject, TdSpanAttributes } from "../types";
 import { PackageType, StatusCode } from "@use-tusk/drift-schemas/core/span";
 import { logger, OriginalGlobalUtils } from "../utils";

src/core/tracing/adapters/ApiSpanAdapter.ts

@@ -91,8 +91,8 @@ export class ApiSpanAdapter implements SpanExportAdapter {
       packageType: cleanSpan.packageType || PackageType.UNSPECIFIED,
       inputValue: toStruct(cleanSpan.inputValue),
       outputValue: toStruct(cleanSpan.outputValue),
-      inputSchema: toStruct(cleanSpan.inputSchema),
-      outputSchema: toStruct(cleanSpan.outputSchema),
+      inputSchema: cleanSpan.inputSchema,
+      outputSchema: cleanSpan.outputSchema,
       inputSchemaHash: cleanSpan.inputSchemaHash || "",
       outputSchemaHash: cleanSpan.outputSchemaHash || "",
       inputValueHash: cleanSpan.inputValueHash || "",

src/core/types.ts

@@ -1,6 +1,7 @@
 import { createContextKey } from "@opentelemetry/api";
 import { StatusCode, PackageType } from "@use-tusk/drift-schemas/core/span";
 import { SpanKind } from "@opentelemetry/api";
+import type { JsonSchema } from "@use-tusk/drift-schemas/core/json_schema";
 
 export const REPLAY_TRACE_ID_CONTEXT_KEY = createContextKey("td.replayTraceId");
 export const SPAN_KIND_CONTEXT_KEY = createContextKey("td.spanKind");
@@ -60,11 +61,12 @@
 
   packageType?: PackageType;
 
-  // keep these as plain JSON for readability
+  // Values are kept as plain JSON for readability
   inputValue: unknown;
   outputValue: unknown;
-  inputSchema: unknown;
-  outputSchema: unknown;
+  // Schemas use the proto JsonSchema type
+  inputSchema: JsonSchema;
+  outputSchema?: JsonSchema; // optional bc replay outbound spans don't have output schemas
 
   inputSchemaHash: string;
   outputSchemaHash: string;

src/instrumentation/libraries/http/HttpTransformEngine.test.ts

@@ -70,7 +70,7 @@
   spanId: "login_456",
   kind: SpanKind.SERVER,
   protocol: "http",
   inputValue: {
     method: "POST",
     url: "http://localhost:3000/api/auth/login",
     target: "/api/auth/login",
@@ -131,6 +131,11 @@
     headers: {
       "content-type": "application/json",
     },
+    httpVersion: "1.1",
+    httpVersionMajor: 1,
+    httpVersionMinor: 1,
+    complete: true,
+    readable: false,
     body: Buffer.from(JSON.stringify({
       id: "ch_123456",
       amount: 1000,
@@ -145,7 +150,7 @@
   spanId: "multi_999",
   kind: SpanKind.SERVER,
   protocol: "http",
   inputValue: {
     method: "GET",
     url: "http://localhost:3000/api/user/lookup?ssn=123-45-6789&email=user@example.com",
     target: "/api/user/lookup?ssn=123-45-6789&email=user@example.com",
@@ -201,7 +206,7 @@
   t.truthy(transformed);
   const span = transformed as HttpSpanData;
 
   t.regex((span.inputValue as HttpServerInputValue).body.password, /^PWD_[0-9a-f]{12}\.\.\.$/);
   t.is(span.transformMetadata?.actions.length, 1);
   t.is(span.transformMetadata?.actions[0].type, "redact");
   t.is(span.transformMetadata?.actions[0].field, "jsonPath:$.password");
@@ -244,7 +249,7 @@
   const inputValue = span.inputValue as HttpServerInputValue;
   t.is(inputValue.url, "http://localhost:3000/api/user/lookup?ssn=XXXXXXXXXXX&email=user%40example.com");
   t.is(inputValue.target, "/api/user/lookup?ssn=XXXXXXXXXXX&email=user%40example.com");
   t.is(inputValue.body.user.password, "HIDDEN");
   t.is(span.transformMetadata?.actions.length, 2);
   const actionTypes = span.transformMetadata!.actions.map((a) => ({ type: a.type, field: a.field }));
   t.true(actionTypes.some((a) => a.type === "mask" && a.field === "queryParam:ssn"));
@@ -269,7 +274,7 @@
   const engine = new HttpTransformEngine(inboundDropConfig);
 
   t.true(
-    engine.shouldDropInboundRequest("POST", "http://localhost:3000/api/auth/login", "localhost", {
+    engine.shouldDropInboundRequest("POST", "http://localhost:3000/api/auth/login", {
       "content-type": "application/json",
     }),
   );
@@ -278,13 +283,12 @@
     engine.shouldDropInboundRequest(
       "POST",
       "http://localhost:3000/api/other/endpoint",
-      "localhost",
       { "content-type": "application/json" },
     ),
   );
 
   t.false(
-    engine.shouldDropInboundRequest("GET", "http://localhost:3000/api/auth/login", "localhost", {
+    engine.shouldDropInboundRequest("GET", "http://localhost:3000/api/auth/login", {
       "content-type": "application/json",
     }),
   );
@@ -406,7 +410,7 @@
   const spanWithoutBody = cloneSpan(exampleServerSpanData);
   spanWithoutBody.inputValue = {
     ...(spanWithoutBody.inputValue as HttpServerInputValue),
     body: undefined,
   };
 
   const result = engine.applyTransforms(spanWithoutBody as any);
@@ -430,7 +434,7 @@
   const spanWithoutHeaders = cloneSpan(exampleServerSpanData);
   spanWithoutHeaders.inputValue = {
     ...(spanWithoutHeaders.inputValue as HttpServerInputValue),
-    headers: undefined,
+    headers: {} as any,
   };
 
   const result = engine.applyTransforms(spanWithoutHeaders as any);
@@ -454,7 +458,7 @@
   const engine = new HttpTransformEngine(maskConfig);
   const result = engine.applyTransforms(cloneSpan(exampleServerSpanData));
 
   t.is((result.inputValue as HttpServerInputValue).body.password, "#################");
 });
 
 test("applies redact transform with custom hash prefix", (t) => {
@@ -674,7 +678,7 @@
 
   const inputValue = result.inputValue as HttpServerInputValue;
   t.regex(inputValue.url, /http:\/\/localhost:3000\/api\/test\?token=REDACTED_[0-9a-f]{12}\.\.\.&other=value/);
-  t.regex(inputValue.target, /\/api\/test\?token=REDACTED_[0-9a-f]{12}\.\.\.&other=value/);
+  t.regex(inputValue.target!, /\/api\/test\?token=REDACTED_[0-9a-f]{12}\.\.\.&other=value/);
 });
 
 test("transforms query parameter in client path with existing params", (t) => {
@@ -700,7 +704,7 @@
   const result = engine.applyTransforms(spanWithQuery);
 
   const inputValue = result.inputValue as HttpClientInputValue;
-  t.regex(inputValue.path, /\/v1\/charges\?api_key=REDACTED_[0-9a-f]{12}\.\.\.&other=value/);
+  t.regex(inputValue.path!, /\/v1\/charges\?api_key=REDACTED_[0-9a-f]{12}\.\.\.&other=value/);
 });
 
 test("handles query param that doesn't exist", (t) => {

src/instrumentation/libraries/http/integration-tests/outbound-drop.test.ts

@@ -112,11 +112,11 @@ test("should drop outbound span data but keep span present when calling dropped
 
   // Should have transformMetadata indicating it was dropped
   t.truthy(droppedSpan!.transformMetadata, "Expected transformMetadata to be defined");
-  const hasDropAction = droppedSpan!.transformMetadata.actions?.some(
+  const hasDropAction = droppedSpan!.transformMetadata?.actions?.some(
     (action) => action.type === "drop" && action.field === "entire_span",
   );
   t.truthy(
     hasDropAction,
-    `Expected drop action in transformMetadata, got ${JSON.stringify(droppedSpan!.transformMetadata.actions)}`,
+    `Expected drop action in transformMetadata, got ${JSON.stringify(droppedSpan!.transformMetadata?.actions)}`,
   );
 });

src/instrumentation/libraries/http/integration-tests/small.test.ts

@@ -37,15 +37,15 @@ const express = require("express");
 const serviceAApp = express();
 serviceAApp.use(express.json());
 
-serviceAApp.post("/api/sensitive", (req, res) => {
+serviceAApp.post("/api/sensitive", (req: any, res: any) => {
   res.json({
     status: "success",
     sensitiveData: "TOP_SECRET_123",
     userId: req.body.userId,
   });
 });
 
-serviceAApp.get("/api/data", (req, res) => {
+serviceAApp.get("/api/data", (req: any, res: any) => {
   res.json({ data: "public data from service A" });
 });
 
@@ -71,8 +71,8 @@ test.after.always(() => {
 
 test("should be able to reach the server", async (t) => {
   await new Promise<void>((resolve) => {
-    http.get(`http://127.0.0.1:${port}/api/data`, (res) => {
-      res.on("data", function (chunk) {
+    http.get(`http://127.0.0.1:${port}/api/data`, (res: any) => {
+      res.on("data", function (chunk: any) {
         console.log("BODY: " + chunk);
       });
 

src/instrumentation/libraries/http/types.ts

@@ -7,7 +7,6 @@ import type {
   ServerResponse,
   IncomingHttpHeaders,
 } from "http";
-import { HttpBodyType } from "./utils";
 
 export type HttpProtocol = "http" | "https";
 
@@ -21,7 +20,6 @@ export interface HttpClientInputValue {
   timeout?: number;
   body?: unknown;
   bodySize?: number;
-  bodyType?: HttpBodyType;
   hasBodyParsingError?: boolean;
   [key: string]: unknown;
 }