From 273f6f0b273192a837d012694660176848b4c87d Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 13 Jul 2026 11:02:18 -0400 Subject: [PATCH 1/4] Add chat history gantt timeline viewer --- pyproject.toml | 1 + src/drs/chat_gantt.py | 346 +++++++++++++++++++++++ src/drs/chat_gantt_tui.py | 458 +++++++++++++++++++++++++++++++ src/drs/commands/chat.py | 38 +++ tests/test_cli.py | 69 +++++ tests/test_commands/test_chat.py | 171 +++++++++++- uv.lock | 57 ++++ 7 files changed, 1132 insertions(+), 8 deletions(-) create mode 100644 src/drs/chat_gantt.py create mode 100644 src/drs/chat_gantt_tui.py diff --git a/pyproject.toml b/pyproject.toml index 7294b9b..8f3ae99 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "pyyaml>=6", "pydantic>=2", "rich>=13", + "textual>=0.79", "prompt-toolkit>=3.0", ] diff --git a/src/drs/chat_gantt.py b/src/drs/chat_gantt.py new file mode 100644 index 0000000..d00217b --- /dev/null +++ b/src/drs/chat_gantt.py @@ -0,0 +1,346 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Shared chat history Gantt parsing and rendering helpers.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Any + + +@dataclass +class ToolSpan: + """A completed tool call span.""" + + lane: int + step: int + name: str + call_id: str + start: datetime + end: datetime + duration_ms: int + offset_ms: int + label: str + arguments: dict[str, Any] | None + title: str | None + summarized_title: str | None + + +@dataclass +class HistoryBounds: + """Overall event timing bounds for a chat history dump.""" + + start: datetime + end: datetime + + @property + def total_ms(self) -> int: + return max(int((self.end - self.start).total_seconds() * 1000), 1) + + +def parse_timestamp(value: str) -> datetime: + """Parse an ISO-8601 timestamp from chat history.""" + return datetime.fromisoformat(value.replace("Z", "+00:00")) + + +def load_history_dump(path: Path) -> dict[str, Any]: + """Load a chat history dump, accepting slightly malformed JSON payloads.""" + text = path.read_text(encoding="utf-8") + + for strict in (True, False): + try: + return json.loads(text, strict=strict) + except json.JSONDecodeError: + continue + + rows = [] + for line in text.splitlines(): + line = line.strip() + if not line: + continue + rows.append(json.loads(line, strict=False)) + return {"data": rows} + + +def summarize_tool_arguments(arguments: Any) -> str: + """Return a compact single-line argument summary for chart labels.""" + if not isinstance(arguments, dict) or not arguments: + return "" + + parts = [] + for value in arguments.values(): + text = str(value).replace("\n", " ").strip() + if text: + parts.append(text) + if len(parts) >= 2: + break + return ", ".join(parts) + + +def truncate_label(label: str, limit: int) -> str: + if len(label) <= limit: + return label + if limit <= 3: + return label[:limit] + return label[: limit - 3] + "..." + + +def build_history_bounds(rows: list[dict]) -> HistoryBounds | None: + """Compute overall elapsed time across all timestamped events.""" + timestamps = [ + parse_timestamp(str(created_at)) + for row in rows + if (created_at := row.get("createdAt")) + ] + if not timestamps: + return None + return HistoryBounds(start=min(timestamps), end=max(timestamps)) + + +def _with_think_time(spans: list[ToolSpan], min_gap_ms: int = 5) -> list[ToolSpan]: + """Insert synthetic think-time spans between non-overlapping steps.""" + if not spans: + return spans + + step_bounds: list[tuple[int, datetime, datetime]] = [] + for step in sorted({span.step for span in spans}): + step_spans = [span for span in spans if span.step == step] + step_bounds.append((step, min(span.start for span in step_spans), max(span.end for span in step_spans))) + + extra_spans: list[ToolSpan] = [] + next_lane = max(span.lane for span in spans) + 1 + for (step, _step_start, step_end), (next_step, next_start, _next_end) in zip(step_bounds, step_bounds[1:]): + gap_ms = int((next_start - step_end).total_seconds() * 1000) + if gap_ms <= min_gap_ms: + continue + extra_spans.append( + ToolSpan( + lane=next_lane, + step=next_step, + name="thinkTime", + call_id=f"think-{step}-to-{next_step}", + start=step_end, + end=next_start, + duration_ms=gap_ms, + offset_ms=0, + label=f"think time (Step {step} -> Step {next_step})", + arguments=None, + title=None, + summarized_title=None, + ) + ) + + if not extra_spans: + return spans + + all_spans = spans + extra_spans + first_start = min(span.start for span in all_spans) + for span in all_spans: + span.offset_ms = max(int((span.start - first_start).total_seconds() * 1000), 0) + return sorted(all_spans, key=lambda span: (span.start, span.end, span.lane, span.call_id)) + + +def build_tool_spans( + rows: list[dict], + *, + include_think_time: bool = False, + min_think_time_ms: int = 5, +) -> tuple[list[ToolSpan], datetime | None, datetime | None]: + """Build timed tool spans and assign each to a visual lane.""" + pending: dict[str, dict[str, Any]] = {} + tool_rows: list[tuple[datetime, datetime, str, str, str, dict[str, Any] | None, str | None, str | None]] = [] + current_title: str | None = None + + for row in rows: + chunk_type = row.get("chunkType") + call_id = row.get("callId") + created_at = row.get("createdAt") + if chunk_type == "model": + result = row.get("result") + if isinstance(result, dict): + title = result.get("title") + if isinstance(title, str) and title.strip(): + current_title = title.strip() + if not call_id or not created_at: + continue + + timestamp = parse_timestamp(str(created_at)) + if chunk_type == "toolRequest": + name = str(row.get("name", "tool")) + arguments = row.get("arguments") if isinstance(row.get("arguments"), dict) else None + args_summary = summarize_tool_arguments(arguments) + label = f"{name}({args_summary})" if args_summary else name + summarized_title = row.get("summarizedTitle") + pending[call_id] = { + "start": timestamp, + "name": name, + "label": label, + "arguments": arguments, + "title": current_title, + "summarized_title": summarized_title.strip() if isinstance(summarized_title, str) else None, + } + elif chunk_type == "toolResponse" and call_id in pending: + start_info = pending.pop(call_id) + start = start_info["start"] + end = timestamp if timestamp >= start else start + tool_rows.append( + ( + start, + end, + start_info["name"], + call_id, + start_info["label"], + start_info["arguments"], + start_info["title"], + start_info["summarized_title"], + ) + ) + + if not tool_rows: + return [], None, None + + tool_rows.sort(key=lambda item: (item[0], item[1], item[2], item[3])) + first_start = min(item[0] for item in tool_rows) + last_end = max(item[1] for item in tool_rows) + lane_ends: list[datetime] = [] + current_step = 0 + current_step_end: datetime | None = None + spans: list[ToolSpan] = [] + + for start, end, name, call_id, label, arguments, title, summarized_title in tool_rows: + if current_step_end is None or start >= current_step_end: + current_step += 1 + current_step_end = end + elif end > current_step_end: + current_step_end = end + + lane = 0 + for idx, lane_end in enumerate(lane_ends): + if lane_end <= start: + lane = idx + lane_ends[idx] = end + break + else: + lane = len(lane_ends) + lane_ends.append(end) + + spans.append( + ToolSpan( + lane=lane, + step=current_step, + name=name, + call_id=call_id, + start=start, + end=end, + duration_ms=max(int((end - start).total_seconds() * 1000), 0), + offset_ms=max(int((start - first_start).total_seconds() * 1000), 0), + label=label, + arguments=arguments, + title=title, + summarized_title=summarized_title, + ) + ) + + if include_think_time: + spans = _with_think_time(spans, min_gap_ms=min_think_time_ms) + + return spans, first_start, last_end + + +def format_duration_ms(duration_ms: int) -> str: + seconds = duration_ms / 1000 + if seconds >= 60: + minutes, rem = divmod(seconds, 60) + return f"{int(minutes)}m{rem:05.2f}s" + return f"{seconds:.3f}s" + + +def render_tool_gantt( + data: dict[str, Any], + width: int = 60, + *, + include_think_time: bool = False, + min_think_time_ms: int = 5, +) -> str: + """Render tool calls from a chat history dump as an ASCII Gantt chart.""" + rows = data.get("data", []) + if not isinstance(rows, list): + raise ValueError("Expected top-level 'data' list in chat history dump") + + history_bounds = build_history_bounds(rows) + spans, first_start, last_end = build_tool_spans( + rows, + include_think_time=include_think_time, + min_think_time_ms=min_think_time_ms, + ) + if not spans or first_start is None or last_end is None: + return "No completed tool calls found." + + total_ms = max(int((last_end - first_start).total_seconds() * 1000), 1) + chart_width = max(width, 20) + label_width = min(max(len(truncate_label(span.label, 40)) for span in spans), 40) + axis_ticks = 5 + tick_positions = [round(idx * (chart_width - 1) / axis_ticks) for idx in range(axis_ticks + 1)] + axis = [" "] * chart_width + for pos in tick_positions: + axis[pos] = "|" + + tick_labels = [" "] * chart_width + for idx, pos in enumerate(tick_positions): + label = format_duration_ms(round(total_ms * idx / axis_ticks)) + start_idx = min(max(pos - len(label) // 2, 0), max(chart_width - len(label), 0)) + for off, ch in enumerate(label): + tick_labels[start_idx + off] = ch + + lines = [ + f"Timeline start: {first_start.isoformat()}", + ( + f"Total time taken: {format_duration_ms(history_bounds.total_ms)} " + f"({history_bounds.start.isoformat()} -> {history_bounds.end.isoformat()})" + if history_bounds is not None + else "Total time taken: unknown" + ), + ( + f"Total tool span: {format_duration_ms(total_ms)} across {max(span.step for span in spans)} step(s), " + f"using {max(span.lane for span in spans) + 1} visual lane(s)" + ), + "".join(tick_labels), + "".join(axis), + ] + + for span in spans: + row = [" "] * chart_width + start_col = min((span.offset_ms * chart_width) // total_ms, chart_width - 1) + end_col = max(((span.offset_ms + span.duration_ms) * chart_width) // total_ms, start_col + 1) + end_col = min(end_col, chart_width) + for idx in range(start_col, end_col): + row[idx] = "#" + if end_col - start_col == 1: + row[start_col] = "*" + else: + row[start_col] = "[" + row[end_col - 1] = "]" + + label = truncate_label(span.label, label_width) + lines.append( + f"S{span.step} {label.ljust(label_width)} {''.join(row)} " + f"+{format_duration_ms(span.offset_ms)} / {format_duration_ms(span.duration_ms)}" + ) + + return "\n".join(lines) diff --git a/src/drs/chat_gantt_tui.py b/src/drs/chat_gantt_tui.py new file mode 100644 index 0000000..6872b0d --- /dev/null +++ b/src/drs/chat_gantt_tui.py @@ -0,0 +1,458 @@ +# +# Copyright (C) 2017-2026 Dremio Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Textual UI for chat history Gantt visualization.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +from rich.console import Group, RenderableType +from rich.panel import Panel +from rich.pretty import Pretty +from rich.text import Text +from textual.app import App, ComposeResult +from textual.containers import Horizontal, ScrollableContainer, Vertical +from textual.reactive import reactive +from textual.widgets import DataTable, Footer, Header, Static + +from drs.chat_gantt import ( + HistoryBounds, + ToolSpan, + build_history_bounds, + build_tool_spans, + format_duration_ms, + load_history_dump, + truncate_label, +) + +LABEL_WIDTH = 32 +MIN_BAR_WIDTH = 24 +ROW_LABEL_WIDTH = 8 +ROW_FIXED_WIDTH = 2 + ROW_LABEL_WIDTH + LABEL_WIDTH + 1 + 12 + + +def _header_chart_width(viewport_width: int) -> int: + """Choose a header chart width that follows the viewport size.""" + return max(viewport_width - 4, 40) + + +def _row_bar_width(viewport_width: int) -> int: + """Choose a row bar width that fits within the current viewport.""" + return max(viewport_width - ROW_FIXED_WIDTH, MIN_BAR_WIDTH) + + +def _duration_color(duration_ms: int, total_ms: int) -> str: + """Color bars by share of total tool span.""" + ratio = duration_ms / max(total_ms, 1) + if ratio < 0.05: + return "green" + if ratio < 0.15: + return "yellow" + if ratio < 0.30: + return "magenta" + return "red" + + +def _tool_time_ms(spans: list[ToolSpan]) -> int: + return sum(span.duration_ms for span in spans if span.name != "thinkTime") + + +def _think_time_ms(spans: list[ToolSpan]) -> int: + return sum(span.duration_ms for span in spans if span.name == "thinkTime") + + +def _tool_call_count(spans: list[ToolSpan]) -> int: + return sum(1 for span in spans if span.name != "thinkTime") + + +@dataclass +class ToolTimeline: + """Prepared tool timeline data for the TUI.""" + + spans: list[ToolSpan] + start: datetime + end: datetime + history_bounds: HistoryBounds + + @property + def total_ms(self) -> int: + return max(int((self.end - self.start).total_seconds() * 1000), 1) + + @property + def lane_count(self) -> int: + return max(span.lane for span in self.spans) + 1 + + +def load_tool_timeline(path: Path) -> ToolTimeline: + """Load and validate timeline data from a history dump.""" + data = load_history_dump(path) + return load_tool_timeline_data(data) + + +def load_tool_timeline_data( + data: dict, + *, + include_think_time: bool = False, + min_think_time_ms: int = 5, +) -> ToolTimeline: + """Load and validate timeline data from an in-memory history dump.""" + rows = data.get("data", []) + if not isinstance(rows, list): + raise ValueError("Expected top-level 'data' list in chat history dump") + history_bounds = build_history_bounds(rows) + spans, start, end = build_tool_spans( + rows, + include_think_time=include_think_time, + min_think_time_ms=min_think_time_ms, + ) + if not spans or start is None or end is None or history_bounds is None: + raise ValueError("No completed tool calls found.") + return ToolTimeline(spans=spans, start=start, end=end, history_bounds=history_bounds) + + +class GanttChart(Static): + """Fixed legend and time axis for the Gantt chart.""" + + def __init__(self, timeline: ToolTimeline, **kwargs) -> None: + super().__init__(**kwargs) + self.timeline = timeline + + def render(self) -> RenderableType: + width = _header_chart_width(self.size.width) + chart_lines = [] + axis_ticks = 5 + tick_positions = [round(idx * (width - 1) / axis_ticks) for idx in range(axis_ticks + 1)] + + labels = [" "] * width + markers = [" "] * width + for idx, pos in enumerate(tick_positions): + markers[pos] = "│" + label = format_duration_ms(round(self.timeline.total_ms * idx / axis_ticks)) + start_idx = min(max(pos - len(label) // 2, 0), max(width - len(label), 0)) + for off, ch in enumerate(label): + labels[start_idx + off] = ch + + chart_lines.append(Text("".join(labels), style="dim")) + chart_lines.append(Text("".join(markers), style="dim")) + legend = Text("Legend: ", style="bold") + legend.append("■ think time ", style="bright_black") + legend.append("■ short (<5%) ", style="green") + legend.append("■ medium (5-15%) ", style="yellow") + legend.append("■ long (15-30%) ", style="magenta") + legend.append("■ very long (30%+) ", style="red") + legend.append("▶ selected", style="bold yellow") + chart_lines.append(legend) + + body = Group(*chart_lines) + subtitle = ( + f"{max(span.step for span in self.timeline.spans)} step(s) " + f"{format_duration_ms(self.timeline.total_ms)} tool span " + f"{format_duration_ms(self.timeline.history_bounds.total_ms)} total time " + "↑/↓ select ←/→ pan Enter details" + ) + return Panel(body, title="Tool Timeline", subtitle=subtitle, border_style="blue") + + +class GanttRows(Static): + """Scrollable tool rows for the Gantt chart.""" + + selected_call_id: reactive[str | None] = reactive(None) + + def __init__(self, timeline: ToolTimeline, **kwargs) -> None: + super().__init__(**kwargs) + self.timeline = timeline + + def render(self) -> RenderableType: + label_width = LABEL_WIDTH + width = _row_bar_width(self.size.width) + chart_lines = [] + + for span in self.timeline.spans: + bar = Text() + start_col = min((span.offset_ms * width) // self.timeline.total_ms, width - 1) + end_col = max(((span.offset_ms + span.duration_ms) * width) // self.timeline.total_ms, start_col + 1) + end_col = min(end_col, width) + color = "bright_black" if span.name == "thinkTime" else _duration_color(span.duration_ms, self.timeline.total_ms) + bar_style = f"bold {color}" if span.call_id == self.selected_call_id else color + selected = span.call_id == self.selected_call_id + line_style = "reverse bold" if selected else "" + fill = [" "] * width + for idx in range(start_col, end_col): + fill[idx] = "█" + if end_col - start_col == 1: + fill[start_col] = "◆" + bar.append(("▶ " if selected else " "), style="bold yellow" if selected else "") + bar.append(f"Step {span.step}".ljust(ROW_LABEL_WIDTH), style=f"bold {line_style}".strip()) + bar.append(truncate_label(span.label, label_width).ljust(label_width), style=f"white {line_style}".strip()) + bar.append(" ", style=line_style) + bar.append("".join(fill), style=bar_style) + bar.append(f" {format_duration_ms(span.duration_ms)}", style=f"dim {line_style}".strip()) + chart_lines.append(bar) + + return Group(*chart_lines) + + +class ChartViewport(ScrollableContainer): + """Scrollable, focusable viewport for the Gantt chart.""" + + can_focus = True + BINDINGS = [ + ("up", "cursor_up", "Up"), + ("down", "cursor_down", "Down"), + ("left", "pan_left", "Pan Left"), + ("right", "pan_right", "Pan Right"), + ("enter", "open_details", "Details"), + ] + + def action_cursor_up(self) -> None: + self.app.action_cursor_up() + + def action_cursor_down(self) -> None: + self.app.action_cursor_down() + + def action_pan_left(self) -> None: + self.app.action_pan_left() + + def action_pan_right(self) -> None: + self.app.action_pan_right() + + def action_open_details(self) -> None: + self.app.action_open_details() + + +class ToolDetails(Static): + """Details pane for the selected tool call.""" + + def show_placeholder(self) -> None: + self.update( + Panel( + Text("Select a Gantt row with ↑/↓ and press Enter to view details.", style="dim"), + title="Selection", + border_style="green", + ) + ) + + def show_span(self, span: ToolSpan) -> None: + body = Text() + if span.name == "thinkTime": + body.append("Tool: think time\n", style="dim") + else: + body.append(f"Tool: {span.name}\n", style="bold") + if span.title: + body.append(f"Title: {span.title}\n") + if span.summarized_title: + body.append(f"Summary: {span.summarized_title}\n") + body.append(f"Step: {span.step}\n") + body.append(f"Call ID: {span.call_id}\n") + body.append(f"Start: {span.start.isoformat()}\n") + body.append(f"End: {span.end.isoformat()}\n") + if span.name == "thinkTime": + body.append(f"Think time: {format_duration_ms(span.duration_ms)}\n", style="dim") + else: + body.append(f"Offset: {format_duration_ms(span.offset_ms)}\n") + body.append(f"Duration: {format_duration_ms(span.duration_ms)}\n") + body.append(f"Label: {span.label}\n", style="dim" if span.name == "thinkTime" else "") + body.append( + f"Run total: {format_duration_ms(self.app.timeline.history_bounds.total_ms)} " + f"({self.app.timeline.history_bounds.start.isoformat()} -> " + f"{self.app.timeline.history_bounds.end.isoformat()})\n" + ) + if span.arguments: + self.update( + Panel( + Group(body, Text("\nArguments:", style="bold"), Pretty(span.arguments, expand_all=True)), + title="Selection", + border_style="green", + ) + ) + return + body.append("Arguments: (none)") + self.update(Panel(body, title="Selection", border_style="green")) + + +class ChatGanttApp(App[None]): + """Interactive Textual app for browsing tool timing.""" + + CSS = """ + Screen { + layout: vertical; + } + #main { + height: 1fr; + } + #table-pane { + width: 48; + min-width: 36; + } + #spans { + height: 1fr; + } + #chart-scroll { + height: 1fr; + width: 1fr; + } + #chart-header-scroll { + height: 6; + width: 1fr; + } + #details { + height: 12; + } + #summary { + height: 7; + min-height: 7; + } + """ + + BINDINGS = [ + ("up", "cursor_up", "Up"), + ("down", "cursor_down", "Down"), + ("left", "pan_left", "Pan Left"), + ("right", "pan_right", "Pan Right"), + ("enter", "open_details", "Details"), + ("q", "quit", "Quit"), + ] + + def __init__(self, timeline: ToolTimeline) -> None: + super().__init__() + self.timeline = timeline + self.selected_index = 0 + + def compose(self) -> ComposeResult: + yield Header(show_clock=True) + with Horizontal(id="main"): + with Vertical(id="table-pane"): + yield DataTable(id="spans") + yield ToolDetails(id="details") + yield Static(id="summary") + with Vertical(): + with ScrollableContainer(id="chart-header-scroll"): + yield GanttChart(self.timeline, id="chart-header") + with ChartViewport(id="chart-scroll"): + yield GanttRows(self.timeline, id="chart-rows") + yield Footer() + + def on_mount(self) -> None: + table = self.query_one("#spans", DataTable) + table.cursor_type = "row" + table.zebra_stripes = True + table.add_columns("Step", "Tool", "Offset", "Duration") + for span in self.timeline.spans: + tool_cell: str | Text + offset_cell: str | Text + duration_cell: str | Text + if span.name == "thinkTime": + tool_cell = Text("think time", style="dim") + offset_cell = Text("", style="dim") + duration_cell = Text(format_duration_ms(span.duration_ms), style="dim") + else: + tool_cell = truncate_label(span.name, 24) + offset_cell = format_duration_ms(span.offset_ms) + duration_cell = format_duration_ms(span.duration_ms) + table.add_row( + str(span.step), + tool_cell, + offset_cell, + duration_cell, + key=span.call_id, + ) + details = self.query_one("#details", ToolDetails) + details.show_placeholder() + summary = self.query_one("#summary", Static) + think_time_ms = _think_time_ms(self.timeline.spans) + tool_time_ms = _tool_time_ms(self.timeline.spans) + summary.update( + Panel( + Text( + "\n".join( + [ + f"Think time: {format_duration_ms(think_time_ms) if think_time_ms else '0.000s'}", + f"Tool calls: {_tool_call_count(self.timeline.spans)}", + f"Tool time: {format_duration_ms(tool_time_ms)}", + f"Total time: {format_duration_ms(self.timeline.history_bounds.total_ms)}", + ] + ) + ), + title="Summary", + border_style="cyan", + ) + ) + chart_scroll = self.query_one("#chart-scroll", ChartViewport) + chart_scroll.focus() + self._highlight_selected_span() + + def action_cursor_up(self) -> None: + if self.selected_index > 0: + self.selected_index -= 1 + self._highlight_selected_span() + + def action_cursor_down(self) -> None: + if self.selected_index < len(self.timeline.spans) - 1: + self.selected_index += 1 + self._highlight_selected_span() + + def action_open_details(self) -> None: + details = self.query_one("#details", ToolDetails) + details.show_span(self.timeline.spans[self.selected_index]) + + def action_pan_left(self) -> None: + chart_header_scroll = self.query_one("#chart-header-scroll", ScrollableContainer) + chart_scroll = self.query_one("#chart-scroll", ChartViewport) + chart_header_scroll.scroll_relative(x=-12, animate=False) + chart_scroll.scroll_relative(x=-12, animate=False) + + def action_pan_right(self) -> None: + chart_header_scroll = self.query_one("#chart-header-scroll", ScrollableContainer) + chart_scroll = self.query_one("#chart-scroll", ChartViewport) + chart_header_scroll.scroll_relative(x=12, animate=False) + chart_scroll.scroll_relative(x=12, animate=False) + + def _highlight_selected_span(self) -> None: + span = self.timeline.spans[self.selected_index] + chart_rows = self.query_one("#chart-rows", GanttRows) + chart_rows.selected_call_id = span.call_id + table = self.query_one("#spans", DataTable) + table.move_cursor(row=self.selected_index, column=0) + chart_scroll = self.query_one("#chart-scroll", ChartViewport) + chart_scroll.scroll_to(y=max(self.selected_index, 0), animate=False, force=True) + + +def run_chat_gantt_tui(path: Path, *, include_think_time: bool = False, min_think_time_ms: int = 5) -> None: + """Launch the Textual app for a history dump.""" + data = load_history_dump(path) + timeline = load_tool_timeline_data( + data, + include_think_time=include_think_time, + min_think_time_ms=min_think_time_ms, + ) + ChatGanttApp(timeline).run() + + +def run_chat_gantt_tui_data( + data: dict, + *, + include_think_time: bool = False, + min_think_time_ms: int = 5, +) -> None: + """Launch the Textual app for an in-memory history dump.""" + timeline = load_tool_timeline_data( + data, + include_think_time=include_think_time, + min_think_time_ms=min_think_time_ms, + ) + ChatGanttApp(timeline).run() diff --git a/src/drs/commands/chat.py b/src/drs/commands/chat.py index a605323..153a566 100644 --- a/src/drs/commands/chat.py +++ b/src/drs/commands/chat.py @@ -37,6 +37,7 @@ from rich.table import Table from rich.text import Text +from drs.chat_gantt import load_history_dump, render_tool_gantt from drs.chat_render import ChatRenderer, PlainRenderer from drs.client import DremioClient from drs.output import error as print_error @@ -662,6 +663,9 @@ async def _run(): def chat_history( conversation_id: str = typer.Argument(help="Conversation ID"), limit: int = typer.Option(50, "--limit", "-n", help="Maximum messages to return"), + gantt: bool = typer.Option(False, "--gantt", help="Render the history as a Gantt timeline instead of transcript output"), + think_time: bool = typer.Option(False, "--think-time", help="With --gantt, include synthetic think-time gaps between steps when the gap exceeds 5 ms"), + ascii: bool = typer.Option(False, "--ascii", help="With --gantt, render plain ASCII output instead of launching the Textual TUI"), fmt: ChatFormat = typer.Option(ChatFormat.table, "--format", "-f", help="Output format: json, table"), ) -> None: """Show message history for a conversation.""" @@ -678,9 +682,43 @@ async def _run(): except DremioAPIError as exc: print_error(str(exc)) raise typer.Exit(1) + if gantt: + try: + if ascii: + print(render_tool_gantt(result, include_think_time=think_time)) + else: + from drs.chat_gantt_tui import run_chat_gantt_tui_data + + run_chat_gantt_tui_data(result, include_think_time=think_time) + return + except ValueError as exc: + print_error(f"Unable to render Gantt chart: {exc}") + raise typer.Exit(1) _chat_output(result, fmt) +@app.command("gantt") +def chat_gantt( + dump_file: Path = typer.Argument(exists=True, dir_okay=False, readable=True, help="Path to chat history dump JSON"), + think_time: bool = typer.Option(False, "--think-time", help="Include synthetic think-time gaps between steps when the gap exceeds 5 ms"), + ascii: bool = typer.Option(False, "--ascii", help="Render plain ASCII output instead of launching the Textual TUI"), + width: int = typer.Option(60, "--width", "-w", min=20, help="Chart width in characters"), +) -> None: + """Render a Gantt chart from a chat history dump.""" + try: + if ascii: + data = load_history_dump(dump_file) + print(render_tool_gantt(data, width=width, include_think_time=think_time)) + return + + from drs.chat_gantt_tui import run_chat_gantt_tui + + run_chat_gantt_tui(dump_file, include_think_time=think_time) + except (OSError, ValueError, json.JSONDecodeError) as exc: + print_error(f"Unable to render Gantt chart: {exc}") + raise typer.Exit(1) + + @app.command("delete") def chat_delete( conversation_id: str = typer.Argument(help="Conversation ID to delete"), diff --git a/tests/test_cli.py b/tests/test_cli.py index 0f6cd3a..baf976d 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -17,6 +17,9 @@ from __future__ import annotations +import json +from unittest.mock import AsyncMock + from typer.testing import CliRunner from drs import __version__ @@ -41,3 +44,69 @@ def test_help_short_flag() -> None: result = runner.invoke(app, ["-h"]) assert result.exit_code == 0 assert f"(version {__version__})" in result.output + + +def test_chat_gantt_command(tmp_path) -> None: + dump_path = tmp_path / "history.json" + dump_path.write_text( + json.dumps( + { + "data": [ + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"arg0": "supplier contract risk exposure"}, + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.719Z", + }, + ] + } + ), + encoding="utf-8", + ) + + result = runner.invoke(app, ["chat", "gantt", str(dump_path), "--ascii", "--width", "20"]) + + assert result.exit_code == 0 + assert "Timeline start: 2026-07-13T12:39:29.860000+00:00" in result.output + assert "S1 searchViewsAndTables" in result.output + + +def test_chat_history_gantt_command(monkeypatch) -> None: + from drs.commands import chat + + client = type("DummyClient", (), {"close": AsyncMock()})() + monkeypatch.setattr(chat, "_get_client", lambda: client) + + async def fake_get_messages(client, conversation_id, limit=50): + return { + "data": [ + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"arg0": "supplier contract risk exposure"}, + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.719Z", + }, + ] + } + + monkeypatch.setattr(chat, "get_messages", fake_get_messages) + + result = runner.invoke(app, ["chat", "history", "conv-1", "--gantt", "--ascii", "--think-time"]) + + assert result.exit_code == 0 + assert "Timeline start: 2026-07-13T12:39:29.860000+00:00" in result.output + assert "S1 searchViewsAndTables" in result.output diff --git a/tests/test_commands/test_chat.py b/tests/test_commands/test_chat.py index e53dae2..874ed36 100644 --- a/tests/test_commands/test_chat.py +++ b/tests/test_commands/test_chat.py @@ -21,14 +21,8 @@ import pytest -from drs.commands.chat import ( - cancel_run, - create_conversation, - delete_conversation, - get_messages, - list_conversations, - send_message, -) +from drs.chat_gantt import build_history_bounds, build_tool_spans, load_history_dump, render_tool_gantt +from drs.commands.chat import cancel_run, create_conversation, delete_conversation, get_messages, list_conversations, send_message @pytest.mark.asyncio @@ -113,3 +107,164 @@ async def test_cancel_run(mock_client) -> None: result = await cancel_run(mock_client, "conv-1", "run-1") mock_client.cancel_conversation_run.assert_called_once_with("conv-1", "run-1") assert result["status"] == "ok" + + +def test_build_tool_spans_assigns_parallel_lanes() -> None: + rows = [ + { + "chunkType": "model", + "createdAt": "2026-07-13T12:39:29.847Z", + "result": { + "title": "Supplier Contract Risk Exposure", + }, + }, + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"arg0": "supplier contract risk exposure"}, + "summarizedTitle": "Search supplier contract risk exposure", + }, + { + "chunkType": "toolRequest", + "callId": "c2", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.900Z", + "arguments": {"arg0": "purchase order supplier spend"}, + "summarizedTitle": "Search purchase order supplier spend", + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.719Z", + }, + { + "chunkType": "toolResponse", + "callId": "c2", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.729Z", + }, + ] + + spans, first_start, last_end = build_tool_spans(rows) + + assert first_start is not None + assert last_end is not None + assert len(spans) == 2 + assert {span.lane for span in spans} == {0, 1} + assert {span.step for span in spans} == {1} + assert spans[0].duration_ms == 2859 + assert spans[1].offset_ms == 40 + assert spans[0].arguments == {"arg0": "supplier contract risk exposure"} + assert spans[0].title == "Supplier Contract Risk Exposure" + assert spans[0].summarized_title == "Search supplier contract risk exposure" + + +def test_load_history_dump_accepts_unescaped_control_chars(tmp_path) -> None: + dump_path = tmp_path / "history.json" + dump_path.write_text('{"data":[{"chunkType":"model","result":{"text":"line 1\nline 2"}}]}', encoding="utf-8") + + loaded = load_history_dump(dump_path) + + assert loaded["data"][0]["result"]["text"] == "line 1\nline 2" + + +def test_build_history_bounds_uses_all_events() -> None: + rows = [ + {"chunkType": "userMessage", "createdAt": "2026-07-13T12:39:17.794Z"}, + {"chunkType": "toolRequest", "callId": "c1", "createdAt": "2026-07-13T12:39:29.860Z"}, + {"chunkType": "toolResponse", "callId": "c1", "createdAt": "2026-07-13T12:39:32.719Z"}, + {"chunkType": "model", "createdAt": "2026-07-13T12:39:35.000Z"}, + ] + + bounds = build_history_bounds(rows) + + assert bounds is not None + assert bounds.total_ms == 17206 + + +def test_build_tool_spans_can_insert_think_time() -> None: + rows = [ + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"arg0": "supplier contract risk exposure"}, + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.719Z", + }, + { + "chunkType": "toolRequest", + "callId": "c2", + "name": "runSql", + "createdAt": "2026-07-13T12:39:32.900Z", + "arguments": {"arg0": "select 1"}, + }, + { + "chunkType": "toolResponse", + "callId": "c2", + "name": "runSql", + "createdAt": "2026-07-13T12:39:33.200Z", + }, + ] + + spans, _, _ = build_tool_spans(rows, include_think_time=True) + + think_spans = [span for span in spans if span.name == "thinkTime"] + assert len(think_spans) == 1 + assert think_spans[0].duration_ms == 181 + assert think_spans[0].step == 2 + + +def test_render_tool_gantt_outputs_chart() -> None: + data = { + "data": [ + { + "chunkType": "userMessage", + "createdAt": "2026-07-13T12:39:17.794Z", + }, + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"arg0": "supplier contract risk exposure"}, + }, + { + "chunkType": "toolRequest", + "callId": "c2", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.900Z", + "arguments": {"arg0": "purchase order supplier spend"}, + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.719Z", + }, + { + "chunkType": "toolResponse", + "callId": "c2", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.729Z", + }, + { + "chunkType": "model", + "createdAt": "2026-07-13T12:39:35.000Z", + }, + ] + } + + rendered = render_tool_gantt(data, width=30) + + assert "Total time taken: 17.206s" in rendered + assert "Total tool span: 2.869s across 1 step(s), using 2 visual lane(s)" in rendered + assert "S1 searchViewsAndTables" in rendered diff --git a/uv.lock b/uv.lock index d8758d2..d89917f 100644 --- a/uv.lock +++ b/uv.lock @@ -90,6 +90,7 @@ dependencies = [ { name = "pydantic" }, { name = "pyyaml" }, { name = "rich" }, + { name = "textual" }, { name = "typer" }, ] @@ -108,6 +109,7 @@ requires-dist = [ { name = "pydantic", specifier = ">=2" }, { name = "pyyaml", specifier = ">=6" }, { name = "rich", specifier = ">=13" }, + { name = "textual", specifier = ">=0.79" }, { name = "typer", specifier = ">=0.9" }, ] @@ -192,6 +194,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "linkify-it-py" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "uc-micro-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2e/c9/06ea13676ef354f0af6169587ae292d3e2406e212876a413bf9eece4eb23/linkify_it_py-2.1.0.tar.gz", hash = "sha256:43360231720999c10e9328dc3691160e27a718e280673d444c38d7d3aaa3b98b", size = 29158, upload-time = "2026-03-01T07:48:47.683Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b4/de/88b3be5c31b22333b3ca2f6ff1de4e863d8fe45aaea7485f591970ec1d3e/linkify_it_py-2.1.0-py3-none-any.whl", hash = "sha256:0d252c1594ecba2ecedc444053db5d3a9b7ec1b0dd929c8f1d74dce89f86c05e", size = 19878, upload-time = "2026-03-01T07:48:46.098Z" }, +] + [[package]] name = "markdown-it-py" version = "4.0.0" @@ -204,6 +218,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, ] +[package.optional-dependencies] +linkify = [ + { name = "linkify-it-py" }, +] + +[[package]] +name = "mdit-py-plugins" +version = "0.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/59/fc/f8d0863f8862f25602c0404d75568e89fb6b4109804645e5cdfb1be5cf56/mdit_py_plugins-0.6.1.tar.gz", hash = "sha256:a2bca0f039f39dbd35fb74ae1b5f998608c437463371f0ff7f49a19a17a114d0", size = 56114, upload-time = "2026-05-13T09:03:38.91Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a5/69/6da5581c6a7fede7dc261bf4e67d6adca4196f176b43288b55b3db395b6e/mdit_py_plugins-0.6.1-py3-none-any.whl", hash = "sha256:214c82fb2ac524472ab6a5bcab1de80f73b50443e187f401bfd77efbc7c6481d", size = 66663, upload-time = "2026-05-13T09:03:37.76Z" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -547,6 +578,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl", hash = "sha256:7ecfff8f2fd72616f7481040475a65b2bf8af90a56c89140852d1120324e8686", size = 9755, upload-time = "2023-10-24T04:13:38.866Z" }, ] +[[package]] +name = "textual" +version = "8.2.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py", extra = ["linkify"] }, + { name = "mdit-py-plugins" }, + { name = "platformdirs" }, + { name = "pygments" }, + { name = "rich" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/00/21/39a76b01bd5eea82a04baaca7580e105d8c59450df03998345bb2cfb307b/textual-8.2.8.tar.gz", hash = "sha256:3f106a9fbc73e39dd266c9712432087de78a6d644084c7c241d6a25c3169115b", size = 1860502, upload-time = "2026-06-30T06:51:24.495Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/be/35261223d9416a0751cdff1c7b4a6f881387218a12d439fe22fefebc8c04/textual-8.2.8-py3-none-any.whl", hash = "sha256:267375fd402dc8d981457212efa71f0e3365fd17bba144ba9bb3ed7563cb374a", size = 731418, upload-time = "2026-06-30T06:51:26.364Z" }, +] + [[package]] name = "typer" version = "0.24.1" @@ -583,6 +631,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] +[[package]] +name = "uc-micro-py" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/78/67/9a363818028526e2d4579334460df777115bdec1bb77c08f9db88f6389f2/uc_micro_py-2.0.0.tar.gz", hash = "sha256:c53691e495c8db60e16ffc4861a35469b0ba0821fe409a8a7a0a71864d33a811", size = 6611, upload-time = "2026-03-01T06:31:27.526Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/61/73/d21edf5b204d1467e06500080a50f79d49ef2b997c79123a536d4a17d97c/uc_micro_py-2.0.0-py3-none-any.whl", hash = "sha256:3603a3859af53e5a39bc7677713c78ea6589ff188d70f4fee165db88e22b242c", size = 6383, upload-time = "2026-03-01T06:31:26.257Z" }, +] + [[package]] name = "virtualenv" version = "21.2.4" From 0e83e9c99e09ea8ff9b6dd51390d59ce24e66b5f Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 13 Jul 2026 12:46:50 -0400 Subject: [PATCH 2/4] Refine chat gantt error details --- src/drs/chat_gantt.py | 72 ++++++++++++- src/drs/chat_gantt_tui.py | 175 +++++++++++++++++++++++-------- tests/test_cli.py | 33 ++++++ tests/test_commands/test_chat.py | 137 +++++++++++++++++++++++- 4 files changed, 365 insertions(+), 52 deletions(-) diff --git a/src/drs/chat_gantt.py b/src/drs/chat_gantt.py index d00217b..31ce171 100644 --- a/src/drs/chat_gantt.py +++ b/src/drs/chat_gantt.py @@ -40,6 +40,8 @@ class ToolSpan: arguments: dict[str, Any] | None title: str | None summarized_title: str | None + failed: bool + error_message: str | None @dataclass @@ -78,6 +80,14 @@ def load_history_dump(path: Path) -> dict[str, Any]: return {"data": rows} +def extract_history_rows(data: dict[str, Any]) -> list[dict[str, Any]]: + """Normalize history rows across dump and API payload shapes.""" + rows = data.get("data", data.get("messages", [])) + if not isinstance(rows, list): + raise ValueError("Expected top-level 'data' or 'messages' list in chat history dump") + return rows + + def summarize_tool_arguments(arguments: Any) -> str: """Return a compact single-line argument summary for chart labels.""" if not isinstance(arguments, dict) or not arguments: @@ -93,6 +103,53 @@ def summarize_tool_arguments(arguments: Any) -> str: return ", ".join(parts) +def _stringify_error(value: Any) -> str | None: + if value is None: + return None + if isinstance(value, str): + text = value.strip() + return text or None + if isinstance(value, dict): + for key in ("message", "error", "errorMessage", "detail", "details"): + nested = _stringify_error(value.get(key)) + if nested: + return nested + for nested_value in value.values(): + nested = _stringify_error(nested_value) + if nested: + return nested + return None + if isinstance(value, list): + parts = [_stringify_error(item) for item in value] + parts = [part for part in parts if part] + return "; ".join(parts) if parts else None + return str(value) + + +def extract_tool_error(row: dict[str, Any]) -> tuple[bool, str | None]: + """Best-effort extraction of tool failure state from a toolResponse row.""" + status = str(row.get("status", "")).strip().lower() + result = row.get("result") + + for key in ("error", "errorMessage", "message"): + direct_error = _stringify_error(row.get(key)) + if direct_error: + return True, direct_error + + nested_error = _stringify_error(result) + if isinstance(result, dict): + if any(key in result for key in ("error", "errorMessage", "message", "detail", "details")) and nested_error: + return True, nested_error + result_status = str(result.get("status", "")).strip().lower() + if result_status in {"error", "failed", "failure", "cancelled", "canceled"}: + return True, nested_error or result_status + + if status in {"error", "failed", "failure", "cancelled", "canceled"}: + return True, nested_error or status + + return False, None + + def truncate_label(label: str, limit: int) -> str: if len(label) <= limit: return label @@ -143,6 +200,8 @@ def _with_think_time(spans: list[ToolSpan], min_gap_ms: int = 5) -> list[ToolSpa arguments=None, title=None, summarized_title=None, + failed=False, + error_message=None, ) ) @@ -164,7 +223,7 @@ def build_tool_spans( ) -> tuple[list[ToolSpan], datetime | None, datetime | None]: """Build timed tool spans and assign each to a visual lane.""" pending: dict[str, dict[str, Any]] = {} - tool_rows: list[tuple[datetime, datetime, str, str, str, dict[str, Any] | None, str | None, str | None]] = [] + tool_rows: list[tuple[datetime, datetime, str, str, str, dict[str, Any] | None, str | None, str | None, bool, str | None]] = [] current_title: str | None = None for row in rows: @@ -199,6 +258,7 @@ def build_tool_spans( start_info = pending.pop(call_id) start = start_info["start"] end = timestamp if timestamp >= start else start + failed, error_message = extract_tool_error(row) tool_rows.append( ( start, @@ -209,6 +269,8 @@ def build_tool_spans( start_info["arguments"], start_info["title"], start_info["summarized_title"], + failed, + error_message, ) ) @@ -223,7 +285,7 @@ def build_tool_spans( current_step_end: datetime | None = None spans: list[ToolSpan] = [] - for start, end, name, call_id, label, arguments, title, summarized_title in tool_rows: + for start, end, name, call_id, label, arguments, title, summarized_title, failed, error_message in tool_rows: if current_step_end is None or start >= current_step_end: current_step += 1 current_step_end = end @@ -254,6 +316,8 @@ def build_tool_spans( arguments=arguments, title=title, summarized_title=summarized_title, + failed=failed, + error_message=error_message, ) ) @@ -279,9 +343,7 @@ def render_tool_gantt( min_think_time_ms: int = 5, ) -> str: """Render tool calls from a chat history dump as an ASCII Gantt chart.""" - rows = data.get("data", []) - if not isinstance(rows, list): - raise ValueError("Expected top-level 'data' list in chat history dump") + rows = extract_history_rows(data) history_bounds = build_history_bounds(rows) spans, first_start, last_end = build_tool_spans( diff --git a/src/drs/chat_gantt_tui.py b/src/drs/chat_gantt_tui.py index 6872b0d..6663268 100644 --- a/src/drs/chat_gantt_tui.py +++ b/src/drs/chat_gantt_tui.py @@ -28,6 +28,7 @@ from textual.app import App, ComposeResult from textual.containers import Horizontal, ScrollableContainer, Vertical from textual.reactive import reactive +from textual.screen import ModalScreen from textual.widgets import DataTable, Footer, Header, Static from drs.chat_gantt import ( @@ -35,6 +36,7 @@ ToolSpan, build_history_bounds, build_tool_spans, + extract_history_rows, format_duration_ms, load_history_dump, truncate_label, @@ -68,6 +70,14 @@ def _duration_color(duration_ms: int, total_ms: int) -> str: return "red" +def _span_marker(span: ToolSpan) -> tuple[str, str]: + if span.failed: + return "●", "red" + if span.name == "thinkTime": + return "•", "bright_black" + return "●", "green" + + def _tool_time_ms(spans: list[ToolSpan]) -> int: return sum(span.duration_ms for span in spans if span.name != "thinkTime") @@ -80,6 +90,45 @@ def _tool_call_count(spans: list[ToolSpan]) -> int: return sum(1 for span in spans if span.name != "thinkTime") +def _build_span_sections(span: ToolSpan, timeline: "ToolTimeline") -> list[RenderableType]: + body = Text() + if span.name == "thinkTime": + body.append("Tool: think time\n", style="dim") + else: + body.append(f"Tool: {span.name}\n", style="bold") + if span.failed: + body.append("Status: failed\n", style="bold red") + elif span.name != "thinkTime": + body.append("Status: success\n", style="green") + if span.title: + body.append(f"Title: {span.title}\n") + if span.summarized_title: + body.append(f"Summary: {span.summarized_title}\n") + body.append(f"Step: {span.step}\n") + body.append(f"Call ID: {span.call_id}\n") + body.append(f"Start: {span.start.isoformat()}\n") + body.append(f"End: {span.end.isoformat()}\n") + if span.name == "thinkTime": + body.append(f"Think time: {format_duration_ms(span.duration_ms)}\n", style="dim") + else: + body.append(f"Offset: {format_duration_ms(span.offset_ms)}\n") + body.append(f"Duration: {format_duration_ms(span.duration_ms)}\n") + body.append(f"Label: {span.label}\n", style="dim" if span.name == "thinkTime" else "") + body.append( + f"Run total: {format_duration_ms(timeline.history_bounds.total_ms)} " + f"({timeline.history_bounds.start.isoformat()} -> " + f"{timeline.history_bounds.end.isoformat()})\n" + ) + sections: list[RenderableType] = [body] + if span.arguments: + sections.extend([Text("\nArguments:", style="bold"), Pretty(span.arguments, expand_all=True)]) + else: + body.append("Arguments: (none)\n") + if span.error_message: + sections.extend([Text("\nError:", style="bold red"), Pretty(span.error_message, expand_all=True)]) + return sections + + @dataclass class ToolTimeline: """Prepared tool timeline data for the TUI.""" @@ -111,9 +160,7 @@ def load_tool_timeline_data( min_think_time_ms: int = 5, ) -> ToolTimeline: """Load and validate timeline data from an in-memory history dump.""" - rows = data.get("data", []) - if not isinstance(rows, list): - raise ValueError("Expected top-level 'data' list in chat history dump") + rows = extract_history_rows(data) history_bounds = build_history_bounds(rows) spans, start, end = build_tool_spans( rows, @@ -150,6 +197,8 @@ def render(self) -> RenderableType: chart_lines.append(Text("".join(labels), style="dim")) chart_lines.append(Text("".join(markers), style="dim")) legend = Text("Legend: ", style="bold") + legend.append("● error ", style="red") + legend.append("● success ", style="green") legend.append("■ think time ", style="bright_black") legend.append("■ short (<5%) ", style="green") legend.append("■ medium (5-15%) ", style="yellow") @@ -187,7 +236,7 @@ def render(self) -> RenderableType: start_col = min((span.offset_ms * width) // self.timeline.total_ms, width - 1) end_col = max(((span.offset_ms + span.duration_ms) * width) // self.timeline.total_ms, start_col + 1) end_col = min(end_col, width) - color = "bright_black" if span.name == "thinkTime" else _duration_color(span.duration_ms, self.timeline.total_ms) + color = "red" if span.failed else "bright_black" if span.name == "thinkTime" else _duration_color(span.duration_ms, self.timeline.total_ms) bar_style = f"bold {color}" if span.call_id == self.selected_call_id else color selected = span.call_id == self.selected_call_id line_style = "reverse bold" if selected else "" @@ -196,9 +245,11 @@ def render(self) -> RenderableType: fill[idx] = "█" if end_col - start_col == 1: fill[start_col] = "◆" + marker, marker_style = _span_marker(span) bar.append(("▶ " if selected else " "), style="bold yellow" if selected else "") bar.append(f"Step {span.step}".ljust(ROW_LABEL_WIDTH), style=f"bold {line_style}".strip()) - bar.append(truncate_label(span.label, label_width).ljust(label_width), style=f"white {line_style}".strip()) + bar.append(f"{marker} ", style=f"{marker_style} {line_style}".strip()) + bar.append(truncate_label(span.label, label_width - 2).ljust(label_width), style=f"white {line_style}".strip()) bar.append(" ", style=line_style) bar.append("".join(fill), style=bar_style) bar.append(f" {format_duration_ms(span.duration_ms)}", style=f"dim {line_style}".strip()) @@ -248,41 +299,58 @@ def show_placeholder(self) -> None: ) def show_span(self, span: ToolSpan) -> None: - body = Text() - if span.name == "thinkTime": - body.append("Tool: think time\n", style="dim") - else: - body.append(f"Tool: {span.name}\n", style="bold") - if span.title: - body.append(f"Title: {span.title}\n") - if span.summarized_title: - body.append(f"Summary: {span.summarized_title}\n") - body.append(f"Step: {span.step}\n") - body.append(f"Call ID: {span.call_id}\n") - body.append(f"Start: {span.start.isoformat()}\n") - body.append(f"End: {span.end.isoformat()}\n") - if span.name == "thinkTime": - body.append(f"Think time: {format_duration_ms(span.duration_ms)}\n", style="dim") - else: - body.append(f"Offset: {format_duration_ms(span.offset_ms)}\n") - body.append(f"Duration: {format_duration_ms(span.duration_ms)}\n") - body.append(f"Label: {span.label}\n", style="dim" if span.name == "thinkTime" else "") - body.append( - f"Run total: {format_duration_ms(self.app.timeline.history_bounds.total_ms)} " - f"({self.app.timeline.history_bounds.start.isoformat()} -> " - f"{self.app.timeline.history_bounds.end.isoformat()})\n" - ) - if span.arguments: - self.update( + sections = _build_span_sections(span, self.app.timeline) + self.update(Panel(Group(*sections), title="Selection", border_style="green")) + + +class ToolDetailModal(ModalScreen[None]): + """Full-screen-ish modal for the selected tool span.""" + + CSS = """ + ToolDetailModal { + align: center middle; + } + #detail-modal { + width: 88%; + height: 88%; + border: round $accent; + background: $surface; + padding: 1 2; + } + #detail-modal-body { + height: 1fr; + width: 1fr; + } + """ + + BINDINGS = [ + ("escape", "dismiss_modal", "Close"), + ("q", "dismiss_modal", "Close"), + ("enter", "dismiss_modal", "Close"), + ] + + def __init__(self, span: ToolSpan, timeline: "ToolTimeline") -> None: + super().__init__() + self.span = span + self.timeline = timeline + + def compose(self) -> ComposeResult: + title = "Tool Call Details" if self.span.name != "thinkTime" else "Think Time Details" + yield ScrollableContainer( + Static( Panel( - Group(body, Text("\nArguments:", style="bold"), Pretty(span.arguments, expand_all=True)), - title="Selection", - border_style="green", - ) - ) - return - body.append("Arguments: (none)") - self.update(Panel(body, title="Selection", border_style="green")) + Group(*_build_span_sections(self.span, self.timeline)), + title=title, + subtitle="Esc/Enter/q closes", + border_style="green" if not self.span.failed else "red", + ), + id="detail-modal-body", + ), + id="detail-modal", + ) + + def action_dismiss_modal(self) -> None: + self.dismiss() class ChatGanttApp(App[None]): @@ -351,19 +419,21 @@ def on_mount(self) -> None: table = self.query_one("#spans", DataTable) table.cursor_type = "row" table.zebra_stripes = True + table.can_focus = True table.add_columns("Step", "Tool", "Offset", "Duration") for span in self.timeline.spans: tool_cell: str | Text offset_cell: str | Text duration_cell: str | Text if span.name == "thinkTime": - tool_cell = Text("think time", style="dim") + tool_cell = Text.assemble(("• ", "bright_black"), ("think time", "dim")) offset_cell = Text("", style="dim") duration_cell = Text(format_duration_ms(span.duration_ms), style="dim") else: - tool_cell = truncate_label(span.name, 24) + marker, marker_style = _span_marker(span) + tool_cell = Text.assemble((f"{marker} ", marker_style), (truncate_label(span.name, 22), "red" if span.failed else "")) offset_cell = format_duration_ms(span.offset_ms) - duration_cell = format_duration_ms(span.duration_ms) + duration_cell = Text(format_duration_ms(span.duration_ms), style="red" if span.failed else "") table.add_row( str(span.step), tool_cell, @@ -392,8 +462,7 @@ def on_mount(self) -> None: border_style="cyan", ) ) - chart_scroll = self.query_one("#chart-scroll", ChartViewport) - chart_scroll.focus() + table.focus() self._highlight_selected_span() def action_cursor_up(self) -> None: @@ -407,8 +476,7 @@ def action_cursor_down(self) -> None: self._highlight_selected_span() def action_open_details(self) -> None: - details = self.query_one("#details", ToolDetails) - details.show_span(self.timeline.spans[self.selected_index]) + self.push_screen(ToolDetailModal(self.timeline.spans[self.selected_index], self.timeline)) def action_pan_left(self) -> None: chart_header_scroll = self.query_one("#chart-header-scroll", ScrollableContainer) @@ -428,9 +496,24 @@ def _highlight_selected_span(self) -> None: chart_rows.selected_call_id = span.call_id table = self.query_one("#spans", DataTable) table.move_cursor(row=self.selected_index, column=0) + details = self.query_one("#details", ToolDetails) + details.show_span(span) chart_scroll = self.query_one("#chart-scroll", ChartViewport) chart_scroll.scroll_to(y=max(self.selected_index, 0), animate=False, force=True) + def on_data_table_row_highlighted(self, event: DataTable.RowHighlighted) -> None: + if event.control.id != "spans" or event.cursor_row is None: + return + self.selected_index = event.cursor_row + self._highlight_selected_span() + + def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None: + if event.control.id != "spans" or event.cursor_row is None: + return + self.selected_index = event.cursor_row + self._highlight_selected_span() + self.action_open_details() + def run_chat_gantt_tui(path: Path, *, include_think_time: bool = False, min_think_time_ms: int = 5) -> None: """Launch the Textual app for a history dump.""" diff --git a/tests/test_cli.py b/tests/test_cli.py index baf976d..4af5829 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -110,3 +110,36 @@ async def fake_get_messages(client, conversation_id, limit=50): assert result.exit_code == 0 assert "Timeline start: 2026-07-13T12:39:29.860000+00:00" in result.output assert "S1 searchViewsAndTables" in result.output + + +def test_chat_history_gantt_command_accepts_messages_envelope(monkeypatch) -> None: + from drs.commands import chat + + client = type("DummyClient", (), {"close": AsyncMock()})() + monkeypatch.setattr(chat, "_get_client", lambda: client) + + async def fake_get_messages(client, conversation_id, limit=50): + return { + "messages": [ + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"arg0": "supplier contract risk exposure"}, + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.719Z", + }, + ] + } + + monkeypatch.setattr(chat, "get_messages", fake_get_messages) + + result = runner.invoke(app, ["chat", "history", "conv-1", "--gantt", "--ascii"]) + + assert result.exit_code == 0 + assert "S1 searchViewsAndTables" in result.output diff --git a/tests/test_commands/test_chat.py b/tests/test_commands/test_chat.py index 874ed36..a85b25a 100644 --- a/tests/test_commands/test_chat.py +++ b/tests/test_commands/test_chat.py @@ -20,8 +20,10 @@ from unittest.mock import AsyncMock import pytest +from rich.console import Console -from drs.chat_gantt import build_history_bounds, build_tool_spans, load_history_dump, render_tool_gantt +from drs.chat_gantt import ToolSpan, build_history_bounds, build_tool_spans, extract_history_rows, load_history_dump, render_tool_gantt +from drs.chat_gantt_tui import ToolTimeline, _build_span_sections from drs.commands.chat import cancel_run, create_conversation, delete_conversation, get_messages, list_conversations, send_message @@ -171,6 +173,14 @@ def test_load_history_dump_accepts_unescaped_control_chars(tmp_path) -> None: assert loaded["data"][0]["result"]["text"] == "line 1\nline 2" +def test_extract_history_rows_accepts_messages_envelope() -> None: + rows = [{"chunkType": "toolRequest", "callId": "c1", "createdAt": "2026-07-13T12:39:29.860Z"}] + + extracted = extract_history_rows({"messages": rows}) + + assert extracted == rows + + def test_build_history_bounds_uses_all_events() -> None: rows = [ {"chunkType": "userMessage", "createdAt": "2026-07-13T12:39:17.794Z"}, @@ -223,6 +233,58 @@ def test_build_tool_spans_can_insert_think_time() -> None: assert think_spans[0].step == 2 +def test_build_tool_spans_marks_failed_tool_calls() -> None: + rows = [ + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "runSql", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"sql": "select * from missing_table"}, + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "runSql", + "createdAt": "2026-07-13T12:39:32.719Z", + "status": "failed", + "result": {"errorMessage": "Table missing_table not found"}, + }, + ] + + spans, _, _ = build_tool_spans(rows) + + assert len(spans) == 1 + assert spans[0].failed is True + assert spans[0].error_message == "Table missing_table not found" + + +def test_build_tool_spans_extracts_nested_error_message() -> None: + rows = [ + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "runSql", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"sql": "select * from missing_table"}, + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "runSql", + "createdAt": "2026-07-13T12:39:32.719Z", + "status": "failed", + "result": {"payload": {"details": {"message": "Nested query failure"}}}, + }, + ] + + spans, _, _ = build_tool_spans(rows) + + assert len(spans) == 1 + assert spans[0].failed is True + assert spans[0].error_message == "Nested query failure" + + def test_render_tool_gantt_outputs_chart() -> None: data = { "data": [ @@ -268,3 +330,76 @@ def test_render_tool_gantt_outputs_chart() -> None: assert "Total time taken: 17.206s" in rendered assert "Total tool span: 2.869s across 1 step(s), using 2 visual lane(s)" in rendered assert "S1 searchViewsAndTables" in rendered + + +def test_render_tool_gantt_accepts_messages_envelope() -> None: + data = { + "messages": [ + { + "chunkType": "toolRequest", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:29.860Z", + "arguments": {"arg0": "supplier contract risk exposure"}, + }, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "searchViewsAndTables", + "createdAt": "2026-07-13T12:39:32.719Z", + }, + ] + } + + rendered = render_tool_gantt(data, width=20) + + assert "S1 searchViewsAndTables" in rendered + + +def test_build_span_sections_includes_error_details() -> None: + span = ToolSpan( + lane=0, + step=1, + name="runSql", + call_id="c1", + start=build_tool_spans( + [ + {"chunkType": "toolRequest", "callId": "c1", "name": "runSql", "createdAt": "2026-07-13T12:39:29.860Z"}, + {"chunkType": "toolResponse", "callId": "c1", "name": "runSql", "createdAt": "2026-07-13T12:39:32.719Z"}, + ] + )[1], + end=build_tool_spans( + [ + {"chunkType": "toolRequest", "callId": "c1", "name": "runSql", "createdAt": "2026-07-13T12:39:29.860Z"}, + {"chunkType": "toolResponse", "callId": "c1", "name": "runSql", "createdAt": "2026-07-13T12:39:32.719Z"}, + ] + )[2], + duration_ms=2859, + offset_ms=0, + label="runSql", + arguments={"sql": "select * from missing_table"}, + title="Broken query", + summarized_title="Run broken query", + failed=True, + error_message="Table missing_table not found", + ) + timeline = ToolTimeline( + spans=[span], + start=span.start, + end=span.end, + history_bounds=build_history_bounds( + [ + {"chunkType": "toolRequest", "createdAt": "2026-07-13T12:39:29.860Z"}, + {"chunkType": "toolResponse", "createdAt": "2026-07-13T12:39:32.719Z"}, + ] + ), + ) + + sections = _build_span_sections(span, timeline) + + console = Console(width=120, record=True) + for section in sections: + console.print(section) + rendered = console.export_text() + assert "Table missing_table not found" in rendered + assert "sql" in rendered From c4492ccf92f511d82365a477f1ceaf192db5fedf Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 13 Jul 2026 13:46:43 -0400 Subject: [PATCH 3/4] Fix chat gantt lint issues --- src/drs/chat_gantt.py | 3 ++- src/drs/chat_gantt_tui.py | 11 ++++++----- tests/test_commands/test_chat.py | 18 ++++++++++++++++-- 3 files changed, 24 insertions(+), 8 deletions(-) diff --git a/src/drs/chat_gantt.py b/src/drs/chat_gantt.py index 31ce171..cc21ab2 100644 --- a/src/drs/chat_gantt.py +++ b/src/drs/chat_gantt.py @@ -20,6 +20,7 @@ import json from dataclasses import dataclass from datetime import datetime +from itertools import pairwise from pathlib import Path from typing import Any @@ -182,7 +183,7 @@ def _with_think_time(spans: list[ToolSpan], min_gap_ms: int = 5) -> list[ToolSpa extra_spans: list[ToolSpan] = [] next_lane = max(span.lane for span in spans) + 1 - for (step, _step_start, step_end), (next_step, next_start, _next_end) in zip(step_bounds, step_bounds[1:]): + for (step, _step_start, step_end), (next_step, next_start, _next_end) in pairwise(step_bounds): gap_ms = int((next_start - step_end).total_seconds() * 1000) if gap_ms <= min_gap_ms: continue diff --git a/src/drs/chat_gantt_tui.py b/src/drs/chat_gantt_tui.py index 6663268..8d498f7 100644 --- a/src/drs/chat_gantt_tui.py +++ b/src/drs/chat_gantt_tui.py @@ -20,6 +20,7 @@ from dataclasses import dataclass from datetime import datetime from pathlib import Path +from typing import ClassVar from rich.console import Group, RenderableType from rich.panel import Panel @@ -90,7 +91,7 @@ def _tool_call_count(spans: list[ToolSpan]) -> int: return sum(1 for span in spans if span.name != "thinkTime") -def _build_span_sections(span: ToolSpan, timeline: "ToolTimeline") -> list[RenderableType]: +def _build_span_sections(span: ToolSpan, timeline: ToolTimeline) -> list[RenderableType]: body = Text() if span.name == "thinkTime": body.append("Tool: think time\n", style="dim") @@ -262,7 +263,7 @@ class ChartViewport(ScrollableContainer): """Scrollable, focusable viewport for the Gantt chart.""" can_focus = True - BINDINGS = [ + BINDINGS: ClassVar = [ ("up", "cursor_up", "Up"), ("down", "cursor_down", "Down"), ("left", "pan_left", "Pan Left"), @@ -323,13 +324,13 @@ class ToolDetailModal(ModalScreen[None]): } """ - BINDINGS = [ + BINDINGS: ClassVar = [ ("escape", "dismiss_modal", "Close"), ("q", "dismiss_modal", "Close"), ("enter", "dismiss_modal", "Close"), ] - def __init__(self, span: ToolSpan, timeline: "ToolTimeline") -> None: + def __init__(self, span: ToolSpan, timeline: ToolTimeline) -> None: super().__init__() self.span = span self.timeline = timeline @@ -387,7 +388,7 @@ class ChatGanttApp(App[None]): } """ - BINDINGS = [ + BINDINGS: ClassVar = [ ("up", "cursor_up", "Up"), ("down", "cursor_down", "Down"), ("left", "pan_left", "Pan Left"), diff --git a/tests/test_commands/test_chat.py b/tests/test_commands/test_chat.py index a85b25a..6579629 100644 --- a/tests/test_commands/test_chat.py +++ b/tests/test_commands/test_chat.py @@ -22,9 +22,23 @@ import pytest from rich.console import Console -from drs.chat_gantt import ToolSpan, build_history_bounds, build_tool_spans, extract_history_rows, load_history_dump, render_tool_gantt +from drs.chat_gantt import ( + ToolSpan, + build_history_bounds, + build_tool_spans, + extract_history_rows, + load_history_dump, + render_tool_gantt, +) from drs.chat_gantt_tui import ToolTimeline, _build_span_sections -from drs.commands.chat import cancel_run, create_conversation, delete_conversation, get_messages, list_conversations, send_message +from drs.commands.chat import ( + cancel_run, + create_conversation, + delete_conversation, + get_messages, + list_conversations, + send_message, +) @pytest.mark.asyncio From 12fe8e42e42252d06946034dd671ba00b9e3d51c Mon Sep 17 00:00:00 2001 From: Aniket Kulkarni Date: Mon, 13 Jul 2026 14:05:54 -0400 Subject: [PATCH 4/4] Apply chat gantt formatting fixes --- src/drs/chat_gantt.py | 10 ++++------ src/drs/chat_gantt_tui.py | 16 +++++++++++++--- src/drs/commands/chat.py | 18 ++++++++++++++---- tests/test_commands/test_chat.py | 14 ++++++++++++-- 4 files changed, 43 insertions(+), 15 deletions(-) diff --git a/src/drs/chat_gantt.py b/src/drs/chat_gantt.py index cc21ab2..a50977e 100644 --- a/src/drs/chat_gantt.py +++ b/src/drs/chat_gantt.py @@ -161,11 +161,7 @@ def truncate_label(label: str, limit: int) -> str: def build_history_bounds(rows: list[dict]) -> HistoryBounds | None: """Compute overall elapsed time across all timestamped events.""" - timestamps = [ - parse_timestamp(str(created_at)) - for row in rows - if (created_at := row.get("createdAt")) - ] + timestamps = [parse_timestamp(str(created_at)) for row in rows if (created_at := row.get("createdAt"))] if not timestamps: return None return HistoryBounds(start=min(timestamps), end=max(timestamps)) @@ -224,7 +220,9 @@ def build_tool_spans( ) -> tuple[list[ToolSpan], datetime | None, datetime | None]: """Build timed tool spans and assign each to a visual lane.""" pending: dict[str, dict[str, Any]] = {} - tool_rows: list[tuple[datetime, datetime, str, str, str, dict[str, Any] | None, str | None, str | None, bool, str | None]] = [] + tool_rows: list[ + tuple[datetime, datetime, str, str, str, dict[str, Any] | None, str | None, str | None, bool, str | None] + ] = [] current_title: str | None = None for row in rows: diff --git a/src/drs/chat_gantt_tui.py b/src/drs/chat_gantt_tui.py index 8d498f7..9d5f16c 100644 --- a/src/drs/chat_gantt_tui.py +++ b/src/drs/chat_gantt_tui.py @@ -237,7 +237,13 @@ def render(self) -> RenderableType: start_col = min((span.offset_ms * width) // self.timeline.total_ms, width - 1) end_col = max(((span.offset_ms + span.duration_ms) * width) // self.timeline.total_ms, start_col + 1) end_col = min(end_col, width) - color = "red" if span.failed else "bright_black" if span.name == "thinkTime" else _duration_color(span.duration_ms, self.timeline.total_ms) + color = ( + "red" + if span.failed + else "bright_black" + if span.name == "thinkTime" + else _duration_color(span.duration_ms, self.timeline.total_ms) + ) bar_style = f"bold {color}" if span.call_id == self.selected_call_id else color selected = span.call_id == self.selected_call_id line_style = "reverse bold" if selected else "" @@ -250,7 +256,9 @@ def render(self) -> RenderableType: bar.append(("▶ " if selected else " "), style="bold yellow" if selected else "") bar.append(f"Step {span.step}".ljust(ROW_LABEL_WIDTH), style=f"bold {line_style}".strip()) bar.append(f"{marker} ", style=f"{marker_style} {line_style}".strip()) - bar.append(truncate_label(span.label, label_width - 2).ljust(label_width), style=f"white {line_style}".strip()) + bar.append( + truncate_label(span.label, label_width - 2).ljust(label_width), style=f"white {line_style}".strip() + ) bar.append(" ", style=line_style) bar.append("".join(fill), style=bar_style) bar.append(f" {format_duration_ms(span.duration_ms)}", style=f"dim {line_style}".strip()) @@ -432,7 +440,9 @@ def on_mount(self) -> None: duration_cell = Text(format_duration_ms(span.duration_ms), style="dim") else: marker, marker_style = _span_marker(span) - tool_cell = Text.assemble((f"{marker} ", marker_style), (truncate_label(span.name, 22), "red" if span.failed else "")) + tool_cell = Text.assemble( + (f"{marker} ", marker_style), (truncate_label(span.name, 22), "red" if span.failed else "") + ) offset_cell = format_duration_ms(span.offset_ms) duration_cell = Text(format_duration_ms(span.duration_ms), style="red" if span.failed else "") table.add_row( diff --git a/src/drs/commands/chat.py b/src/drs/commands/chat.py index 153a566..d80bf21 100644 --- a/src/drs/commands/chat.py +++ b/src/drs/commands/chat.py @@ -663,9 +663,17 @@ async def _run(): def chat_history( conversation_id: str = typer.Argument(help="Conversation ID"), limit: int = typer.Option(50, "--limit", "-n", help="Maximum messages to return"), - gantt: bool = typer.Option(False, "--gantt", help="Render the history as a Gantt timeline instead of transcript output"), - think_time: bool = typer.Option(False, "--think-time", help="With --gantt, include synthetic think-time gaps between steps when the gap exceeds 5 ms"), - ascii: bool = typer.Option(False, "--ascii", help="With --gantt, render plain ASCII output instead of launching the Textual TUI"), + gantt: bool = typer.Option( + False, "--gantt", help="Render the history as a Gantt timeline instead of transcript output" + ), + think_time: bool = typer.Option( + False, + "--think-time", + help="With --gantt, include synthetic think-time gaps between steps when the gap exceeds 5 ms", + ), + ascii: bool = typer.Option( + False, "--ascii", help="With --gantt, render plain ASCII output instead of launching the Textual TUI" + ), fmt: ChatFormat = typer.Option(ChatFormat.table, "--format", "-f", help="Output format: json, table"), ) -> None: """Show message history for a conversation.""" @@ -700,7 +708,9 @@ async def _run(): @app.command("gantt") def chat_gantt( dump_file: Path = typer.Argument(exists=True, dir_okay=False, readable=True, help="Path to chat history dump JSON"), - think_time: bool = typer.Option(False, "--think-time", help="Include synthetic think-time gaps between steps when the gap exceeds 5 ms"), + think_time: bool = typer.Option( + False, "--think-time", help="Include synthetic think-time gaps between steps when the gap exceeds 5 ms" + ), ascii: bool = typer.Option(False, "--ascii", help="Render plain ASCII output instead of launching the Textual TUI"), width: int = typer.Option(60, "--width", "-w", min=20, help="Chart width in characters"), ) -> None: diff --git a/tests/test_commands/test_chat.py b/tests/test_commands/test_chat.py index 6579629..2490418 100644 --- a/tests/test_commands/test_chat.py +++ b/tests/test_commands/test_chat.py @@ -379,13 +379,23 @@ def test_build_span_sections_includes_error_details() -> None: start=build_tool_spans( [ {"chunkType": "toolRequest", "callId": "c1", "name": "runSql", "createdAt": "2026-07-13T12:39:29.860Z"}, - {"chunkType": "toolResponse", "callId": "c1", "name": "runSql", "createdAt": "2026-07-13T12:39:32.719Z"}, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "runSql", + "createdAt": "2026-07-13T12:39:32.719Z", + }, ] )[1], end=build_tool_spans( [ {"chunkType": "toolRequest", "callId": "c1", "name": "runSql", "createdAt": "2026-07-13T12:39:29.860Z"}, - {"chunkType": "toolResponse", "callId": "c1", "name": "runSql", "createdAt": "2026-07-13T12:39:32.719Z"}, + { + "chunkType": "toolResponse", + "callId": "c1", + "name": "runSql", + "createdAt": "2026-07-13T12:39:32.719Z", + }, ] )[2], duration_ms=2859,