From 6b8cf46082ceb2c72c2cc078d2ab65bc8461d35d Mon Sep 17 00:00:00 2001
From: Sohan Kshirsagar <sohan.kshirsagar@gmail.com>
Date: Tue, 14 Apr 2026 12:01:35 -0700
Subject: [PATCH] fix: add output schema merges to Django replay span for
 correct UI display

---
 drift/instrumentation/django/middleware.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/drift/instrumentation/django/middleware.py b/drift/instrumentation/django/middleware.py
index 6514855..b534f8d 100644
--- a/drift/instrumentation/django/middleware.py
+++ b/drift/instrumentation/django/middleware.py
@@ -274,9 +274,10 @@ def process_view(
     def _capture_replay_output(self, request: HttpRequest, response: HttpResponse, span_info: SpanInfo) -> None:
         """Capture response data on the span for REPLAY mode.
 
-        Sets OUTPUT_VALUE so the inbound replay span sent to the CLI includes
-        the actual response for UI comparison. Skips RECORD-mode concerns like
-        transforms, trace blocking, and schema merges.
+        Sets OUTPUT_VALUE and OUTPUT_SCHEMA_MERGES so the inbound replay span
+        sent to the CLI includes the actual response with correct schema hints
+        for UI display. Skips RECORD-mode concerns like transforms and trace
+        blocking.
 
         Args:
             request: Django HttpRequest object
@@ -313,6 +314,9 @@ def _capture_replay_output(self, request: HttpRequest, response: HttpResponse, s
 
         span_info.span.set_attribute(TdSpanAttributes.OUTPUT_VALUE, json.dumps(output_value))
 
+        output_schema_merges = build_output_schema_merges(output_value)
+        span_info.span.set_attribute(TdSpanAttributes.OUTPUT_SCHEMA_MERGES, json.dumps(output_schema_merges))
+
     def _normalize_html_response(self, response: HttpResponse) -> HttpResponse:
         """Normalize HTML response body for REPLAY mode comparison.