Merge pull request #3 from phiat/feature/enhanced-metadata

Add agent type labels, tool duration, and token usage display

Author: phiat

internal/parser/parser.go

@@ -22,29 +22,45 @@ const (
 
 // StreamItem represents a single item in the output stream
 type StreamItem struct {
-	Type      StreamItemType
-	SessionID string // which session this belongs to
-	AgentID   string // empty for main session, "abc123" for subagents
-	AgentName string // human-readable name derived from agent type or ID
-	Timestamp time.Time
-	Content   string
-	ToolName  string // for tool_input/tool_output
-	ToolID    string // to correlate input with output
+	Type         StreamItemType
+	SessionID    string // which session this belongs to
+	AgentID      string // empty for main session, "abc123" for subagents
+	AgentName    string // human-readable name derived from agent type or ID
+	Timestamp    time.Time
+	Content      string
+	ToolName     string // for tool_input/tool_output
+	ToolID       string // to correlate input with output
+	DurationMs   int64  // tool execution duration in ms (0 = not available)
+	InputTokens  int64  // usage.input_tokens from assistant messages
+	OutputTokens int64  // usage.output_tokens from assistant messages
 }
 
 // RawMessage represents a line from the JSONL file
 type RawMessage struct {
-	Type      string          `json:"type"`
-	AgentID   string          `json:"agentId,omitempty"`
-	SessionID string          `json:"sessionId"`
-	Timestamp string          `json:"timestamp"`
-	Message   json.RawMessage `json:"message"`
+	Type          string          `json:"type"`
+	AgentID       string          `json:"agentId,omitempty"`
+	SessionID     string          `json:"sessionId"`
+	Timestamp     string          `json:"timestamp"`
+	Message       json.RawMessage `json:"message"`
+	ToolUseResult json.RawMessage `json:"toolUseResult,omitempty"`
+}
+
+// RawToolUseResult represents the toolUseResult field on user messages
+type RawToolUseResult struct {
+	DurationMs int64 `json:"durationMs"`
 }
 
 // AssistantMessage represents the message field for assistant responses
 type AssistantMessage struct {
 	Role    string         `json:"role"`
 	Content []ContentBlock `json:"content"`
+	Usage   *UsageInfo     `json:"usage,omitempty"`
+}
+
+// UsageInfo represents token usage from assistant messages
+type UsageInfo struct {
+	InputTokens  int64 `json:"input_tokens"`
+	OutputTokens int64 `json:"output_tokens"`
 }
 
 // ContentBlock represents a single content item in assistant response
@@ -159,6 +175,12 @@ func parseAssistantMessage(raw RawMessage, timestamp time.Time) []StreamItem {
 		}
 	}
 
+	// Attach token usage to the first item only
+	if len(items) > 0 && msg.Usage != nil {
+		items[0].InputTokens = msg.Usage.InputTokens
+		items[0].OutputTokens = msg.Usage.OutputTokens
+	}
+
 	return items
 }
 
@@ -171,6 +193,15 @@ func parseUserMessage(raw RawMessage, timestamp time.Time) []StreamItem {
 		return nil
 	}
 
+	// Parse toolUseResult for duration
+	var durationMs int64
+	if len(raw.ToolUseResult) > 0 {
+		var tur RawToolUseResult
+		if err := json.Unmarshal(raw.ToolUseResult, &tur); err == nil {
+			durationMs = tur.DurationMs
+		}
+	}
+
 	var items []StreamItem
 	agentName := "Main"
 	if raw.AgentID != "" {
@@ -180,12 +211,13 @@ func parseUserMessage(raw RawMessage, timestamp time.Time) []StreamItem {
 	for _, result := range results {
 		if result.Type == "tool_result" {
 			items = append(items, StreamItem{
-				Type:      TypeToolOutput,
-				AgentID:   raw.AgentID,
-				AgentName: agentName,
-				Timestamp: timestamp,
-				Content:   extractToolResultContent(result.Content),
-				ToolID:    result.ToolUseID,
+				Type:       TypeToolOutput,
+				AgentID:    raw.AgentID,
+				AgentName:  agentName,
+				Timestamp:  timestamp,
+				Content:    extractToolResultContent(result.Content),
+				ToolID:     result.ToolUseID,
+				DurationMs: durationMs,
 			})
 		}
 	}

internal/tui/model.go

@@ -22,19 +22,21 @@ const (
 
 // Model is the main TUI model
 type Model struct {
-	tree         *TreeView
-	stream       *StreamView
-	watcher      *watcher.Watcher
-	focus        Focus
-	showTree     bool
-	width        int
-	height       int
-	treeWidth    int
-	sessionID    string
-	skipHistory  bool
-	pollInterval time.Duration
-	err          error
-	quitting     bool
+	tree              *TreeView
+	stream            *StreamView
+	watcher           *watcher.Watcher
+	focus             Focus
+	showTree          bool
+	width             int
+	height            int
+	treeWidth         int
+	sessionID         string
+	skipHistory       bool
+	pollInterval      time.Duration
+	err               error
+	quitting          bool
+	totalInputTokens  int64
+	totalOutputTokens int64
 }
 
 // NewModel creates a new TUI model
@@ -87,7 +89,8 @@ func (m *Model) initWatcher() tea.Cmd {
 		for _, session := range w.GetSessions() {
 			m.tree.AddSession(session.ID, session.ProjectPath)
 			for agentID := range session.Subagents {
-				m.tree.AddAgent(session.ID, agentID)
+				agentType := session.SubagentTypes[agentID]
+				m.tree.AddAgent(session.ID, agentID, agentType)
 			}
 		}
 
@@ -125,11 +128,19 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.updateActivityStatus()
 
 	case streamItemMsg:
-		m.stream.AddItem(parser.StreamItem(msg))
+		item := parser.StreamItem(msg)
+		// Accumulate token usage (includes history — shows total session cost)
+		if item.InputTokens > 0 {
+			m.totalInputTokens += item.InputTokens
+		}
+		if item.OutputTokens > 0 {
+			m.totalOutputTokens += item.OutputTokens
+		}
+		m.stream.AddItem(item)
 		m.stream.SetEnabledFilters(m.tree.GetEnabledFilters())
 
 	case newAgentMsg:
-		m.tree.AddAgent(msg.SessionID, msg.AgentID)
+		m.tree.AddAgent(msg.SessionID, msg.AgentID, msg.AgentType)
 		m.stream.SetEnabledFilters(m.tree.GetEnabledFilters())
 
 	case newSessionMsg:
@@ -391,14 +402,36 @@ func (m *Model) renderHeader() string {
 		}
 	}
 
+	// Token usage display
+	tokenInfo := ""
+	if m.totalInputTokens > 0 || m.totalOutputTokens > 0 {
+		tokenInfo = fmt.Sprintf("│ %s in / %s out",
+			formatTokenCount(m.totalInputTokens),
+			formatTokenCount(m.totalOutputTokens))
+	}
+
 	// Build header - use plain text and apply headerStyle uniformly (like Rust version)
 	// Don't use Width() as it causes truncation on narrow terminals
 	headerText := fmt.Sprintf("%s  │  %s", toggles, sessionInfo)
+	if tokenInfo != "" {
+		headerText += "  " + tokenInfo
+	}
 	header := headerStyle.Render(headerText)
 
 	return header
 }
 
+// formatTokenCount formats token counts for display
+func formatTokenCount(n int64) string {
+	if n < 1000 {
+		return fmt.Sprintf("%d", n)
+	}
+	if n < 1000000 {
+		return fmt.Sprintf("%.1fk", float64(n)/1000.0)
+	}
+	return fmt.Sprintf("%.1fm", float64(n)/1000000.0)
+}
+
 func (m *Model) renderToggle(name string, enabled bool, key string) string {
 	checkbox := "☐"
 	if enabled {

internal/tui/stream.go

@@ -215,7 +215,11 @@ func (s *StreamView) renderItem(item parser.StreamItem, width int) string {
 		b.WriteString(toolInputContentStyle.Render(content))
 
 	case parser.TypeToolOutput:
-		header := toolOutputStyle.Render(toolOutputIcon + " Output")
+		outputLabel := toolOutputIcon + " Output"
+		if item.DurationMs > 0 {
+			outputLabel += " " + formatDuration(item.DurationMs)
+		}
+		header := toolOutputStyle.Render(outputLabel)
 		b.WriteString(fmt.Sprintf("%s%s%s\n", agentName, sep, header))
 		content := s.truncateContent(item.Content, width)
 		b.WriteString(toolOutputContentStyle.Render(content))
@@ -263,6 +267,19 @@ func (s *StreamView) truncateContent(content string, width int) string {
 	return strings.Join(wrapped, "\n")
 }
 
+// formatDuration formats a duration in milliseconds to a human-readable string
+func formatDuration(ms int64) string {
+	if ms < 1000 {
+		return fmt.Sprintf("(%dms)", ms)
+	}
+	secs := float64(ms) / 1000.0
+	if secs < 60 {
+		return fmt.Sprintf("(%.1fs)", secs)
+	}
+	mins := secs / 60.0
+	return fmt.Sprintf("(%.1fm)", mins)
+}
+
 // View renders the stream
 func (s *StreamView) View() string {
 	return s.viewport.View()

internal/tui/tree.go

@@ -105,8 +105,10 @@ func (t *TreeView) AddSession(sessionID, projectPath string) *TreeNode {
 	return session
 }
 
-// AddAgent adds a subagent under a session
-func (t *TreeView) AddAgent(sessionID, agentID string) {
+// AddAgent adds a subagent under a session.
+// If agentType is non-empty, it is used as the display name.
+// For compound types like "feature-dev:code-reviewer", only the part after ":" is used.
+func (t *TreeView) AddAgent(sessionID, agentID, agentType string) {
 	// Find the session node
 	var session *TreeNode
 	for _, child := range t.Root.Children {
@@ -127,11 +129,21 @@ func (t *TreeView) AddAgent(sessionID, agentID string) {
 		}
 	}
 
+	displayName := fmt.Sprintf("Agent-%s", agentID[:min(AgentIDDisplayLength, len(agentID))])
+	if agentType != "" {
+		// For compound types like "feature-dev:code-reviewer", use part after ":"
+		if idx := strings.LastIndex(agentType, ":"); idx >= 0 && idx < len(agentType)-1 {
+			displayName = agentType[idx+1:]
+		} else {
+			displayName = agentType
+		}
+	}
+
 	node := &TreeNode{
 		Type:      NodeTypeAgent,
 		ID:        agentID,
 		SessionID: sessionID,
-		Name:      fmt.Sprintf("Agent-%s", agentID[:min(AgentIDDisplayLength, len(agentID))]),
+		Name:      displayName,
 		Enabled:   true,
 		IsActive:  true,
 		Parent:    session,

internal/tui/tree_test.go

@@ -44,7 +44,7 @@ func TestTreeView_AddSessionDuplicate(t *testing.T) {
 func TestTreeView_AddAgent(t *testing.T) {
 	tv := NewTreeView()
 	tv.AddSession("sess1", "project")
-	tv.AddAgent("sess1", "agent123456789")
+	tv.AddAgent("sess1", "agent123456789", "")
 
 	session := tv.Root.Children[0]
 	if len(session.Children) != 2 {
@@ -65,7 +65,7 @@ func TestTreeView_AddAgent(t *testing.T) {
 func TestTreeView_AddAgentNoSession(t *testing.T) {
 	tv := NewTreeView()
 	// Should not panic when adding agent to non-existent session
-	tv.AddAgent("nonexistent", "agent1")
+	tv.AddAgent("nonexistent", "agent1", "")
 	if len(tv.Root.Children) != 0 {
 		t.Error("should not add anything for non-existent session")
 	}
@@ -74,8 +74,8 @@ func TestTreeView_AddAgentNoSession(t *testing.T) {
 func TestTreeView_AddAgentDuplicate(t *testing.T) {
 	tv := NewTreeView()
 	tv.AddSession("sess1", "project")
-	tv.AddAgent("sess1", "agent1")
-	tv.AddAgent("sess1", "agent1")
+	tv.AddAgent("sess1", "agent1", "")
+	tv.AddAgent("sess1", "agent1", "")
 
 	session := tv.Root.Children[0]
 	if len(session.Children) != 2 {
@@ -107,7 +107,7 @@ func TestTreeView_AddBackgroundTask(t *testing.T) {
 func TestTreeView_AddBackgroundTaskUnderAgent(t *testing.T) {
 	tv := NewTreeView()
 	tv.AddSession("sess1", "project")
-	tv.AddAgent("sess1", "agent1")
+	tv.AddAgent("sess1", "agent1", "")
 	tv.AddBackgroundTask("sess1", "agent1", "toolu_456", "Task: explore", "/path/out.txt", true)
 
 	agent := tv.Root.Children[0].Children[1]
@@ -180,7 +180,7 @@ func TestTreeView_RemoveNonExistent(t *testing.T) {
 func TestTreeView_GetEnabledFilters(t *testing.T) {
 	tv := NewTreeView()
 	tv.AddSession("sess1", "project")
-	tv.AddAgent("sess1", "agent1")
+	tv.AddAgent("sess1", "agent1", "")
 
 	filters := tv.GetEnabledFilters()
 	if len(filters) != 2 {
@@ -222,7 +222,7 @@ func TestTreeView_GetEnabledFiltersDisabled(t *testing.T) {
 func TestTreeView_Navigation(t *testing.T) {
 	tv := NewTreeView()
 	tv.AddSession("sess1", "project")
-	tv.AddAgent("sess1", "agent1")
+	tv.AddAgent("sess1", "agent1", "")
 	// Nodes: session, main, agent = 3 nodes
 
 	if tv.cursor != 0 {
@@ -300,7 +300,7 @@ func TestTreeView_GetSelectedNode(t *testing.T) {
 func TestTreeView_GetSelectedSession(t *testing.T) {
 	tv := NewTreeView()
 	tv.AddSession("sess1", "project")
-	tv.AddAgent("sess1", "agent1")
+	tv.AddAgent("sess1", "agent1", "")
 
 	// At session
 	tv.cursor = 0