Fix

Author: jy-tan

drift/core/mock_utils.py

@@ -17,7 +17,7 @@
     from .json_schema_helper import SchemaMerges
     from .types import CleanSpanData
 
-from .json_schema_helper import JsonSchema, JsonSchemaHelper
+from .json_schema_helper import JsonSchemaHelper
 from .types import (
     Duration,
     PackageType,
@@ -92,7 +92,7 @@ def convert_mock_request_to_clean_span(
         input_schema=input_result.schema,
         input_schema_hash=input_result.decoded_schema_hash,
         input_value_hash=input_result.decoded_value_hash,
-        output_schema=JsonSchema(),
+        output_schema=None,  # type: ignore[arg-type] - Must be None to avoid betterproto serialization issues
         output_schema_hash="",
         output_value_hash="",
         kind=kind,

drift/instrumentation/psycopg2/instrumentation.py

@@ -22,7 +22,7 @@
 
 from ...core.communication.types import MockRequestInput
 from ...core.drift_sdk import TuskDrift
-from ...core.json_schema_helper import JsonSchema, JsonSchemaHelper
+from ...core.json_schema_helper import JsonSchemaHelper
 from ...core.tracing import TdSpanAttributes
 from ...core.types import (
     CleanSpanData,
@@ -675,8 +675,8 @@ def _try_get_mock(
                 submodule_name="query",
                 input_value=input_value,
                 output_value=None,
-                input_schema=JsonSchema(),
-                output_schema=JsonSchema(),
+                input_schema=None,  # type: ignore[arg-type]
+                output_schema=None,  # type: ignore[arg-type]
                 input_schema_hash=input_result.decoded_schema_hash,
                 output_schema_hash="",
                 input_value_hash=input_result.decoded_value_hash,