Fix trace reuse and packaged model assets

trace_formatter/time_windows.py

@@ -8,6 +8,8 @@ from pathlib import Path
 
 from trace_analyzer.helpers import parse_jsonish, safe_int
 
+from .trace_io import open_trace_text
+
 WINDOW_RE = re.compile(r"(?P<day>\d{4})-(?P<start>\d{4})-(?P<end>\d{4})$")
 UTC_PLUS_8 = timezone(timedelta(hours=8))
 
@@ -32,8 +34,17 @@ def parse_time_to_ms(value: str) -> int:
     raise ValueError(f"Unsupported timestamp format: {value!r}")
 
 
+def _trace_window_name(path: Path) -> str:
+    name = path.name
+    if name.endswith(".jsonl.zst"):
+        return name[: -len(".jsonl.zst")]
+    if name.endswith(".jsonl"):
+        return name[: -len(".jsonl")]
+    return path.stem
+
+
 def _read_first_timestamp(path: Path) -> str:
-    with path.open("r", encoding="utf-8") as handle:
+    with open_trace_text(path) as handle:
         for line in handle:
             stripped = line.strip()
             if not stripped:
@@ -46,7 +57,7 @@ def _read_first_timestamp(path: Path) -> str:
 
 
 def _read_first_timestamp_and_ready_ms(path: Path) -> tuple[str, int]:
-    with path.open("r", encoding="utf-8") as handle:
+    with open_trace_text(path) as handle:
         for line in handle:
             stripped = line.strip()
             if not stripped:
@@ -90,8 +101,8 @@ def infer_time_window(
     if not source_files:
         return None
 
-    first_match = WINDOW_RE.match(source_files[0].stem)
-    last_match = WINDOW_RE.match(source_files[-1].stem)
+    first_match = WINDOW_RE.match(_trace_window_name(source_files[0]))
+    last_match = WINDOW_RE.match(_trace_window_name(source_files[-1]))
     if first_match is None or last_match is None:
         return None