From d1ec76cb44fb7826a24cd3070ad34babc620772c Mon Sep 17 00:00:00 2001
From: liurenfeng94-ops <279598012+liurenfeng94-ops@users.noreply.github.com>
Date: Tue, 28 Jul 2026 22:55:30 +0800
Subject: [PATCH] Fix trace tree root duration

---
 src/traceDuration.ts        |  6 ++++++
 src/traceTree.ts            |  8 ++++----
 test/trace-duration.test.ts | 12 ++++++++++++
 3 files changed, 22 insertions(+), 4 deletions(-)
 create mode 100644 src/traceDuration.ts
 create mode 100644 test/trace-duration.test.ts

diff --git a/src/traceDuration.ts b/src/traceDuration.ts
new file mode 100644
index 0000000..158515f
--- /dev/null
+++ b/src/traceDuration.ts
@@ -0,0 +1,6 @@
+export function getTraceRootDuration(currentDuration: number, line: { ts: number, dur?: number }): number {
+  if (line.dur === Number.MAX_SAFE_INTEGER)
+    return currentDuration
+
+  return Math.max(currentDuration, line.ts + (line.dur ?? 0))
+}
diff --git a/src/traceTree.ts b/src/traceTree.ts
index 0618d33..6f4cc6e 100644
--- a/src/traceTree.ts
+++ b/src/traceTree.ts
@@ -4,6 +4,7 @@ import type { TraceData, TraceLine, TypeLine } from '../shared/src/traceData'
 import { getWorkspacePath } from './storage'
 import { postMessage } from './webview'
 import { traceFiles } from './appState'
+import { getTraceRootDuration } from './traceDuration'
 
 export interface Tree { id: number, line: TraceLine, children: Tree[], types: TypeLine[], childCnt: number, childTypeCnt: number, typeCnt: number }
 function getRoot(): Tree {
@@ -31,7 +32,7 @@ export function toTree(traceData: TraceData, workspacePath: string): Tree {
   const tree: Tree = { ...getRoot() }
   let endTs = Number.MAX_SAFE_INTEGER
   let curr = tree
-  let maxDur = 0
+  let maxEndTs = 0
   let id = 0
 
   const stack: Tree[] = []
@@ -44,8 +45,7 @@ export function toTree(traceData: TraceData, workspacePath: string): Tree {
     if ('args' in line && line.args?.path && isAbsolute(line.args?.path))
       line.args.path = relative(workspacePath, line.args.path)
 
-    if (line.dur !== Number.MAX_SAFE_INTEGER && (line.dur ?? 0) > maxDur)
-      maxDur = line.ts
+    maxEndTs = getTraceRootDuration(maxEndTs, line)
 
     while (line.ts > endTs) {
       if (stack.length === 0)
@@ -70,7 +70,7 @@ export function toTree(traceData: TraceData, workspacePath: string): Tree {
     }
   }
 
-  tree.line.dur = maxDur
+  tree.line.dur = maxEndTs
   return tree
 }
 
diff --git a/test/trace-duration.test.ts b/test/trace-duration.test.ts
new file mode 100644
index 0000000..158c46c
--- /dev/null
+++ b/test/trace-duration.test.ts
@@ -0,0 +1,12 @@
+import { describe, expect, it } from 'vitest'
+import { getTraceRootDuration } from '../src/traceDuration'
+
+describe('getTraceRootDuration', () => {
+  it('uses the end timestamp of a finite trace event', () => {
+    expect(getTraceRootDuration(0, { dur: 5, ts: 10 })).toBe(15)
+  })
+
+  it('keeps the current duration for open-ended root events', () => {
+    expect(getTraceRootDuration(22, { dur: Number.MAX_SAFE_INTEGER, ts: 10 })).toBe(22)
+  })
+})