From 6d279c4fed4a8a78e2258b97081a674f48710cd8 Mon Sep 17 00:00:00 2001
From: Sanket2329 <sanketshakya62@gmail.com>
Date: Sun, 7 Jun 2026 19:43:48 +0530
Subject: [PATCH 1/2] [AGE-3766] feat(frontend): Make collapse headers sticky
 in PrettyJsonView

---
 .../components/DrillInView/PrettyJsonView.tsx | 37 +++++++++++++++----
 .../DrillInView/TraceSpanDrillInView.tsx      | 14 +++----
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/web/oss/src/components/DrillInView/PrettyJsonView.tsx b/web/oss/src/components/DrillInView/PrettyJsonView.tsx
index b8916bf06e..f363424c75 100644
--- a/web/oss/src/components/DrillInView/PrettyJsonView.tsx
+++ b/web/oss/src/components/DrillInView/PrettyJsonView.tsx
@@ -538,7 +538,7 @@ const CopyButton = ({value}: {value: unknown}) => {
         <button
             type="button"
             aria-label={copied ? "Copied" : "Copy to clipboard"}
-            className="opacity-0 group-hover/row:opacity-100 focus-visible:opacity-100 transition-opacity inline-flex items-center justify-center h-[22px] min-w-[22px] px-1.5 ml-auto border border-transparent rounded-sm text-[var(--ant-color-text-quaternary)] cursor-pointer shrink-0 hover:text-[var(--ant-color-text)] hover:bg-[var(--ant-color-bg-container)] hover:border-[var(--ant-color-border)] focus-visible:ring-1 focus-visible:ring-[var(--ant-color-primary)] focus-visible:outline-none"
+            className="opacity-0 group-hover/row:opacity-100 focus-visible:opacity-100 transition-opacity inline-flex items-center justify-center h-[22px] min-w-[22px] px-1.5 ml-auto self-center border border-transparent rounded-sm text-[var(--ant-color-text-quaternary)] cursor-pointer shrink-0 hover:text-[var(--ant-color-text)] hover:bg-[var(--ant-color-bg-container)] hover:border-[var(--ant-color-border)] focus-visible:ring-1 focus-visible:ring-[var(--ant-color-primary)] focus-visible:outline-none"
             onClick={handleCopy}
             onKeyDown={handleKeyDown}
         >
@@ -568,6 +568,7 @@ const NodeRow = memo(function NodeRow({
     value,
     isSection,
     isMessage,
+    depth = 0,
 }: {
     keyLabel: React.ReactNode
     meta?: string
@@ -578,6 +579,7 @@ const NodeRow = memo(function NodeRow({
     value?: unknown
     isSection?: boolean
     isMessage?: boolean
+    depth?: number
 }) {
     const [open, setOpen] = useState(defaultOpen)
     const toggle = useCallback(() => setOpen((o) => !o), [])
@@ -592,7 +594,8 @@ const NodeRow = memo(function NodeRow({
     return (
         <div className={isSection ? "pt-1 first:pt-0" : ""}>
             <div
-                className={`group/row flex items-baseline gap-2 py-1 px-1 rounded-sm min-h-[24px] select-none ${collapsible ? "cursor-pointer sticky top-0 z-[1] bg-[var(--ant-color-bg-container)]" : ""} hover:bg-[var(--ant-color-fill-quaternary)] focus-visible:ring-1 focus-visible:ring-[var(--ant-color-primary)] focus-visible:outline-none`}
+                className={`group/row flex items-baseline gap-2 py-0.5 px-1 rounded-sm min-h-[28px] h-[28px] select-none ${collapsible ? "cursor-pointer sticky z-[1] bg-[var(--ant-color-bg-container)]" : ""} hover:bg-[var(--ant-color-fill-quaternary)] focus-visible:ring-1 focus-visible:ring-[var(--ant-color-primary)] focus-visible:outline-none`}
+                style={collapsible ? {top: `${depth * 28}px`} : undefined}
                 onClick={collapsible ? toggle : undefined}
                 role={collapsible ? "button" : undefined}
                 tabIndex={collapsible ? 0 : undefined}
@@ -791,10 +794,12 @@ const MessageNodeRow = memo(function MessageNodeRow({
     msg,
     index,
     keyPrefix,
+    depth = 0,
 }: {
     msg: {role: string; content: unknown; tool_calls?: unknown[]}
     index: number
     keyPrefix: string
+    depth?: number
 }) {
     const role = (msg.role || "").toLowerCase()
     const roleColor = ROLE_COLOR_CLASSES[role] ?? DEFAULT_ROLE_COLOR_CLASS
@@ -831,7 +836,7 @@ const MessageNodeRow = memo(function MessageNodeRow({
                         name="result"
                         value={parsedContent}
                         keyPrefix={`${editorId}-result`}
-                        depth={1}
+                        depth={depth + 1}
                         expandDepth={2}
                     />
                 </div>
@@ -850,7 +855,7 @@ const MessageNodeRow = memo(function MessageNodeRow({
                         name={part.name}
                         value={part.value}
                         keyPrefix={`${editorId}-part-${i}`}
-                        depth={1}
+                        depth={depth + 1}
                         expandDepth={2}
                     />
                 ))}
@@ -861,13 +866,13 @@ const MessageNodeRow = memo(function MessageNodeRow({
                             name={getToolCallName(tc)}
                             value={getToolCallArgs(tc)}
                             keyPrefix={`${editorId}-tc-${i}`}
-                            depth={1}
+                            depth={depth + 1}
                             expandDepth={2}
                         />
                     ))}
             </div>
         )
-    }, [text, toolCalls, editorId, parsedContent, structuredParts])
+    }, [text, toolCalls, editorId, parsedContent, structuredParts, depth])
 
     return (
         <NodeRow
@@ -878,6 +883,7 @@ const MessageNodeRow = memo(function MessageNodeRow({
             defaultOpen
             value={msg.content ?? (toolCalls ? JSON.stringify(toolCalls) : "")}
             body={body}
+            depth={depth}
         />
     )
 })
@@ -929,10 +935,17 @@ const RecursiveNode = memo(function RecursiveNode({
                 defaultOpen={depth < expandDepth}
                 value={value}
                 isSection={isSection}
+                depth={depth}
                 body={
                     <div className="flex flex-col gap-0.5 py-0.5">
                         {normalized.map((msg, i) => (
-                            <MessageNodeRow key={i} msg={msg} index={i} keyPrefix={nodePrefix} />
+                            <MessageNodeRow
+                                key={i}
+                                msg={msg}
+                                index={i}
+                                keyPrefix={nodePrefix}
+                                depth={depth + 1}
+                            />
                         ))}
                     </div>
                 }
@@ -956,6 +969,7 @@ const RecursiveNode = memo(function RecursiveNode({
                 defaultOpen={depth < expandDepth}
                 value={value}
                 isSection={isSection}
+                depth={depth}
                 body={
                     <>
                         <NodeRow
@@ -964,6 +978,7 @@ const RecursiveNode = memo(function RecursiveNode({
                             collapsible
                             defaultOpen={depth + 1 < expandDepth}
                             value={chatResult.messages}
+                            depth={depth + 1}
                             body={
                                 <div className="flex flex-col gap-0.5 py-0.5">
                                     {normalized.map((msg, i) => (
@@ -972,6 +987,7 @@ const RecursiveNode = memo(function RecursiveNode({
                                             msg={msg}
                                             index={i}
                                             keyPrefix={`${nodePrefix}-${chatKey}`}
+                                            depth={depth + 2}
                                         />
                                     ))}
                                 </div>
@@ -1002,6 +1018,7 @@ const RecursiveNode = memo(function RecursiveNode({
                 collapsible={false}
                 value={value}
                 isSection={isSection}
+                depth={depth}
             />
         )
     }
@@ -1015,6 +1032,7 @@ const RecursiveNode = memo(function RecursiveNode({
                 defaultOpen={depth < expandDepth}
                 value={value}
                 isSection={isSection}
+                depth={depth}
                 body={
                     <EditorProvider
                         id={nodePrefix}
@@ -1057,6 +1075,7 @@ const RecursiveNode = memo(function RecursiveNode({
                     collapsible={false}
                     value={value}
                     isSection={isSection}
+                    depth={depth}
                 />
             )
         }
@@ -1069,6 +1088,7 @@ const RecursiveNode = memo(function RecursiveNode({
                 defaultOpen={depth < expandDepth}
                 value={value}
                 isSection={isSection}
+                depth={depth}
                 body={
                     count > 0
                         ? entries.map(([k, v]) => (
@@ -1100,6 +1120,7 @@ const RecursiveNode = memo(function RecursiveNode({
             collapsible={false}
             value={value}
             isSection={isSection}
+            depth={depth}
         />
     )
 })
@@ -1177,6 +1198,7 @@ export const PrettyJsonView = memo(function PrettyJsonView({
                     defaultOpen
                     value={topChatResult.messages}
                     isSection
+                    depth={0}
                     body={
                         <div className="flex flex-col gap-0.5 py-0.5">
                             {normalized.map((msg, i) => (
@@ -1185,6 +1207,7 @@ export const PrettyJsonView = memo(function PrettyJsonView({
                                     msg={msg}
                                     index={i}
                                     keyPrefix={`${keyPrefix}-${chatKey}`}
+                                    depth={1}
                                 />
                             ))}
                         </div>
diff --git a/web/oss/src/components/DrillInView/TraceSpanDrillInView.tsx b/web/oss/src/components/DrillInView/TraceSpanDrillInView.tsx
index 9ddb1d213e..e4264d3167 100644
--- a/web/oss/src/components/DrillInView/TraceSpanDrillInView.tsx
+++ b/web/oss/src/components/DrillInView/TraceSpanDrillInView.tsx
@@ -466,7 +466,7 @@ export const TraceSpanDrillInView = memo(
         if (rootScope === "span") {
             const showTitle = Boolean(title)
             return (
-                <div className="rounded-md overflow-hidden bg-[var(--ag-c-FFFFFF)]">
+                <div className="rounded-md bg-[var(--ag-c-FFFFFF)]">
                     <div
                         className={`drill-in-field-header rounded-md flex items-center justify-between py-2 px-3 bg-[var(--ag-c-FFFFFF)] border border-solid border-[var(--ag-rgba-051729-06)] ${allowSpanCollapse ? "cursor-pointer" : ""}`}
                         onClick={allowSpanCollapse ? toggleCollapsed : undefined}
@@ -509,7 +509,7 @@ export const TraceSpanDrillInView = memo(
                         </div>
                     </div>
                     {(!allowSpanCollapse || !isCollapsed) && (
-                        <div className="relative overflow-hidden">
+                        <div className="relative">
                             {isSearchOpen && isCodeMode && (
                                 <div className="absolute right-4 top-3 z-20 flex items-center gap-2 rounded-xl border border-[var(--ag-rgba-051729-14)] bg-[var(--ag-c-FFFFFF)] px-2 py-2 shadow-[0_8px_24px_rgba(5,23,41,0.12)] max-w-[calc(100%-2rem)]">
                                     <Input
@@ -576,12 +576,10 @@ export const TraceSpanDrillInView = memo(
                                     </EditorProvider>
                                 </DrillInProvider>
                             ) : isPrettyJson ? (
-                                <div className="overflow-y-auto">
-                                    <PrettyJsonView
-                                        data={prettyJsonSource}
-                                        keyPrefix={`trace-span-${textViewerId}`}
-                                    />
-                                </div>
+                                <PrettyJsonView
+                                    data={prettyJsonSource}
+                                    keyPrefix={`trace-span-${textViewerId}`}
+                                />
                             ) : (
                                 <div className="mx-1 my-2 rounded-md bg-[var(--ag-c-FFFFFF)]">
                                     <TextModeViewer

From a469ce9317579ab5264a5c427bbe05077ebbb98e Mon Sep 17 00:00:00 2001
From: Sanket2329 <sanketshakya62@gmail.com>
Date: Sun, 7 Jun 2026 19:56:04 +0530
Subject: [PATCH 2/2] fix(frontend): address review comments on PrettyJsonView
 and TraceSpanDrillInView

---
 web/oss/src/components/DrillInView/PrettyJsonView.tsx      | 7 +++++--
 .../src/components/DrillInView/TraceSpanDrillInView.tsx    | 4 ++--
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/web/oss/src/components/DrillInView/PrettyJsonView.tsx b/web/oss/src/components/DrillInView/PrettyJsonView.tsx
index f363424c75..4be5270b0f 100644
--- a/web/oss/src/components/DrillInView/PrettyJsonView.tsx
+++ b/web/oss/src/components/DrillInView/PrettyJsonView.tsx
@@ -558,6 +558,9 @@ const CopyButton = ({value}: {value: unknown}) => {
 // so keys align in a single column at every depth. Containers get an
 // interactive chevron; leaves get an invisible 14px spacer.
 
+const ROW_HEIGHT_PX = 28
+const ROW_HEIGHT_CLASS = "min-h-[28px] h-[28px]"
+
 const NodeRow = memo(function NodeRow({
     keyLabel,
     meta,
@@ -594,8 +597,8 @@ const NodeRow = memo(function NodeRow({
     return (
         <div className={isSection ? "pt-1 first:pt-0" : ""}>
             <div
-                className={`group/row flex items-baseline gap-2 py-0.5 px-1 rounded-sm min-h-[28px] h-[28px] select-none ${collapsible ? "cursor-pointer sticky z-[1] bg-[var(--ant-color-bg-container)]" : ""} hover:bg-[var(--ant-color-fill-quaternary)] focus-visible:ring-1 focus-visible:ring-[var(--ant-color-primary)] focus-visible:outline-none`}
-                style={collapsible ? {top: `${depth * 28}px`} : undefined}
+                className={`group/row flex items-baseline gap-2 py-0.5 px-1 rounded-sm ${ROW_HEIGHT_CLASS} select-none ${collapsible ? "cursor-pointer sticky z-[1] bg-[var(--ant-color-bg-container)]" : ""} hover:bg-[var(--ant-color-fill-quaternary)] focus-visible:ring-1 focus-visible:ring-[var(--ant-color-primary)] focus-visible:outline-none`}
+                style={collapsible ? {top: `${depth * ROW_HEIGHT_PX}px`} : undefined}
                 onClick={collapsible ? toggle : undefined}
                 role={collapsible ? "button" : undefined}
                 tabIndex={collapsible ? 0 : undefined}
diff --git a/web/oss/src/components/DrillInView/TraceSpanDrillInView.tsx b/web/oss/src/components/DrillInView/TraceSpanDrillInView.tsx
index e4264d3167..a4c1b45319 100644
--- a/web/oss/src/components/DrillInView/TraceSpanDrillInView.tsx
+++ b/web/oss/src/components/DrillInView/TraceSpanDrillInView.tsx
@@ -653,8 +653,8 @@ export const TraceSpanDrillInView = memo(
 
         // Type assertion needed because traceSpanMolecule.drillIn is optional in the general type
         // but we know it's configured for the trace entity
-        const entityWithDrillIn = traceSpan as typeof traceSpanMolecule & {
-            drillIn: NonNullable<typeof traceSpanMoleculeMolecule.drillIn>
+        const entityWithDrillIn = traceSpanMolecule as typeof traceSpanMolecule & {
+            drillIn: NonNullable<typeof traceSpanMolecule.drillIn>
         }
 
         return (