correctly send inbound span for replay

Author: sohankshirsagar

drift/core/batch_processor.py

@@ -165,8 +165,6 @@ def _export_batch(self) -> None:
                 else:
                     adapter.export_spans(batch)  # type: ignore
 
-                logger.debug(f"Exported {len(batch)} spans via {adapter.name}")
-
             except Exception as e:
                 logger.error(f"Failed to export batch via {adapter.name}: {e}")
 

drift/core/communication/communicator.py

@@ -11,7 +11,7 @@
 from typing import Any
 
 from tusk.drift.core.v1 import GetMockRequest as ProtoGetMockRequest
-
+from ..span_serialization import clean_span_to_proto
 from ...version import MIN_CLI_VERSION, SDK_VERSION
 from ..types import CleanSpanData, calling_library_context
 from .types import (
@@ -438,7 +438,7 @@ async def send_inbound_span_for_replay(self, span: CleanSpanData) -> None:
         if not self._socket:
             return
 
-        proto_span = span_to_proto(span)
+        proto_span = clean_span_to_proto(span)
 
         request = SendInboundSpanForReplayRequest(span=proto_span)
 

drift/core/tracing/td_span_processor.py

@@ -15,7 +15,7 @@
 
 from ..sampling import should_sample
 from ..trace_blocking_manager import TraceBlockingManager, should_block_span
-from ..types import TD_INSTRUMENTATION_LIBRARY_NAME, TuskDriftMode
+from ..types import TD_INSTRUMENTATION_LIBRARY_NAME, TuskDriftMode, SpanKind as TdSpanKind, replay_trace_id_context
 from .otel_converter import otel_span_to_clean_span_data
 
 if TYPE_CHECKING:
@@ -152,9 +152,17 @@ def on_end(self, span: ReadableSpan) -> None:
                 return
 
             # Handle REPLAY mode inbound spans
-            if self._mode == TuskDriftMode.REPLAY and clean_span.kind.value == 2:  # SERVER
+            if self._mode == TuskDriftMode.REPLAY and clean_span.kind == TdSpanKind.SERVER:
                 from ..drift_sdk import TuskDrift
 
+                # Set the trace ID to the replay trace ID, CLI will use this to match the trace test with the result
+                replay_trace_id = replay_trace_id_context.get()
+                if replay_trace_id:
+                    clean_span.trace_id = replay_trace_id
+                else:
+                    logger.error("No replay trace ID found, cannot send inbound span for replay")
+                    return
+
                 sdk = TuskDrift.get_instance()
                 # Check for running loop BEFORE creating the coroutine to avoid
                 # "coroutine was never awaited" warnings. If we call the async
@@ -167,6 +175,13 @@ def on_end(self, span: ReadableSpan) -> None:
 
                 if loop is not None:
                     loop.create_task(sdk.send_inbound_span_for_replay(clean_span))
+                else:
+                     # No running loop - run synchronously
+                    try:
+                        asyncio.run(sdk.send_inbound_span_for_replay(clean_span))
+                    except RuntimeError:
+                        logger.error("No running loop, cannot send inbound span for replay")
+                        pass
 
             # Export span via batch processor (RECORD mode)
             if self._mode == TuskDriftMode.RECORD: