Fix/tui aggregated token cost

pkg/runtime/event.go

@@ -183,8 +183,9 @@ func Warning(message, agentName string) Event {
 }
 
 type TokenUsageEvent struct {
-	Type  string `json:"type"`
-	Usage *Usage `json:"usage"`
+	Type      string `json:"type"`
+	SessionID string `json:"session_id"`
+	Usage     *Usage `json:"usage"`
 	AgentContext
 }
 
@@ -196,9 +197,10 @@ type Usage struct {
 	Cost          float64 `json:"cost"`
 }
 
-func TokenUsage(inputTokens, outputTokens, contextLength, contextLimit int, cost float64, agentName string) Event {
+func TokenUsage(sessionID, agentName string, inputTokens, outputTokens, contextLength, contextLimit int, cost float64) Event {
 	return &TokenUsageEvent{
-		Type: "token_usage",
+		Type:      "token_usage",
+		SessionID: sessionID,
 		Usage: &Usage{
 			ContextLength: contextLength,
 			ContextLimit:  contextLimit,

pkg/runtime/runtime.go

@@ -257,12 +257,13 @@ func (r *LocalRuntime) forwardRAGEvents(ctx context.Context, ragManagers map[str
 					case "usage":
 						// Convert RAG usage to TokenUsageEvent so TUI displays it
 						sendEvent(TokenUsage(
+							"",
+							agentName,
 							ragEvent.TotalTokens, // input tokens (embeddings)
 							0,                    // output tokens (0 for embeddings)
 							ragEvent.TotalTokens, // context length
 							0,                    // context limit (not applicable)
 							ragEvent.Cost,
-							agentName,
 						))
 					case "ready":
 						sendEvent(RAGReady(ragName, agentName))
@@ -618,7 +619,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 			if m != nil {
 				contextLimit = m.Limit.Context
 			}
-			events <- TokenUsage(sess.InputTokens, sess.OutputTokens, sess.InputTokens+sess.OutputTokens, contextLimit, sess.Cost, r.currentAgent)
+			events <- TokenUsage(sess.ID, r.currentAgent, sess.InputTokens, sess.OutputTokens, sess.InputTokens+sess.OutputTokens, contextLimit, sess.Cost)
 
 			if m != nil && r.sessionCompaction {
 				if sess.InputTokens+sess.OutputTokens > int(float64(contextLimit)*0.9) {
@@ -627,7 +628,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 					if len(res.Calls) == 0 {
 						events <- SessionCompaction(sess.ID, "start", r.currentAgent)
 						r.Summarize(ctx, sess, events)
-						events <- TokenUsage(sess.InputTokens, sess.OutputTokens, sess.InputTokens+sess.OutputTokens, contextLimit, sess.Cost, r.currentAgent)
+						events <- TokenUsage(sess.ID, r.currentAgent, sess.InputTokens, sess.OutputTokens, sess.InputTokens+sess.OutputTokens, contextLimit, sess.Cost)
 						events <- SessionCompaction(sess.ID, "completed", r.currentAgent)
 					}
 				}
@@ -641,7 +642,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 				if sess.InputTokens+sess.OutputTokens > int(float64(contextLimit)*0.9) {
 					events <- SessionCompaction(sess.ID, "start", r.currentAgent)
 					r.Summarize(ctx, sess, events)
-					events <- TokenUsage(sess.InputTokens, sess.OutputTokens, sess.InputTokens+sess.OutputTokens, contextLimit, sess.Cost, r.currentAgent)
+					events <- TokenUsage(sess.ID, r.currentAgent, sess.InputTokens, sess.OutputTokens, sess.InputTokens+sess.OutputTokens, contextLimit, sess.Cost)
 					events <- SessionCompaction(sess.ID, "completed", r.currentAgent)
 				}
 			}
@@ -1265,7 +1266,6 @@ func (r *LocalRuntime) handleTaskTransfer(ctx context.Context, sess *session.Ses
 	}
 
 	sess.ToolsApproved = s.ToolsApproved
-	sess.Cost += s.Cost
 
 	sess.AddSubSession(s)
 

pkg/runtime/runtime_test.go

@@ -223,7 +223,7 @@ func TestSimple(t *testing.T) {
 		UserMessage("Hi"),
 		StreamStarted(sess.ID, "root"),
 		AgentChoice("root", "Hello"),
-		TokenUsage(3, 2, 5, 0, 0, "root"),
+		TokenUsage(sess.ID, "root", 3, 2, 5, 0, 0),
 		StreamStopped(sess.ID, "root"),
 	}
 
@@ -255,7 +255,7 @@ func TestMultipleContentChunks(t *testing.T) {
 		AgentChoice("root", "how "),
 		AgentChoice("root", "are "),
 		AgentChoice("root", "you?"),
-		TokenUsage(8, 12, 20, 0, 0, "root"),
+		TokenUsage(sess.ID, "root", 8, 12, 20, 0, 0),
 		StreamStopped(sess.ID, "root"),
 	}
 
@@ -283,7 +283,7 @@ func TestWithReasoning(t *testing.T) {
 		AgentChoiceReasoning("root", "Let me think about this..."),
 		AgentChoiceReasoning("root", " I should respond politely."),
 		AgentChoice("root", "Hello, how can I help you?"),
-		TokenUsage(10, 15, 25, 0, 0, "root"),
+		TokenUsage(sess.ID, "root", 10, 15, 25, 0, 0),
 		StreamStopped(sess.ID, "root"),
 	}
 
@@ -313,7 +313,7 @@ func TestMixedContentAndReasoning(t *testing.T) {
 		AgentChoice("root", "Hello!"),
 		AgentChoiceReasoning("root", " I should be friendly"),
 		AgentChoice("root", " How can I help you today?"),
-		TokenUsage(15, 20, 35, 0, 0, "root"),
+		TokenUsage(sess.ID, "root", 15, 20, 35, 0, 0),
 		StreamStopped(sess.ID, "root"),
 	}
 

pkg/tui/components/sidebar/sidebar.go

@@ -31,7 +31,7 @@ type Model interface {
 	layout.Model
 	layout.Sizeable
 
-	SetTokenUsage(usage *runtime.Usage)
+	SetTokenUsage(event *runtime.TokenUsageEvent)
 	SetTodos(toolCall tools.ToolCall) error
 	SetWorking(working bool) tea.Cmd
 	SetMode(mode Mode)
@@ -54,7 +54,8 @@ type ragIndexingState struct {
 type model struct {
 	width            int
 	height           int
-	usage            *runtime.Usage
+	sessionUsage     map[string]*runtime.Usage // sessionID -> latest usage snapshot
+	sessionAgent     map[string]string         // sessionID -> agent name
 	todoComp         *todotool.SidebarComponent
 	working          bool
 	mcpInit          bool
@@ -76,7 +77,8 @@ func New(manager *service.TodoManager) Model {
 	return &model{
 		width:          20,
 		height:         24,
-		usage:          &runtime.Usage{},
+		sessionUsage:   make(map[string]*runtime.Usage),
+		sessionAgent:   make(map[string]string),
 		todoComp:       todotool.NewSidebarComponent(manager),
 		spinner:        spinner.New(spinner.ModeSpinnerOnly),
 		sessionTitle:   "New session",
@@ -89,8 +91,15 @@ func (m *model) Init() tea.Cmd {
 	return nil
 }
 
-func (m *model) SetTokenUsage(usage *runtime.Usage) {
-	m.usage = usage
+func (m *model) SetTokenUsage(event *runtime.TokenUsageEvent) {
+	if event == nil || event.Usage == nil || event.SessionID == "" || event.AgentContext.AgentName == "" {
+		return
+	}
+
+	// Store/replace by session ID (each event has cumulative totals for that session)
+	usage := *event.Usage
+	m.sessionUsage[event.SessionID] = &usage
+	m.sessionAgent[event.SessionID] = event.AgentContext.AgentName
 }
 
 func (m *model) SetTodos(toolCall tools.ToolCall) error {
@@ -156,6 +165,24 @@ func formatTokenCount(count int) string {
 	return fmt.Sprintf("%d", count)
 }
 
+func formatCost(cost float64) string {
+	return fmt.Sprintf("%.2f", cost)
+}
+
+// contextPercent returns a context usage percentage string when a single session has a limit.
+func (m *model) contextPercent() (string, bool) {
+	if len(m.sessionUsage) != 1 {
+		return "", false
+	}
+	for _, usage := range m.sessionUsage {
+		if usage.ContextLimit > 0 {
+			percent := (float64(usage.ContextLength) / float64(usage.ContextLimit)) * 100
+			return fmt.Sprintf("Context: %.0f%%", percent), true
+		}
+	}
+	return "", false
+}
+
 // getCurrentWorkingDirectory returns the current working directory with home directory replaced by ~/
 func getCurrentWorkingDirectory() string {
 	pwd, err := os.Getwd()
@@ -171,12 +198,15 @@ func getCurrentWorkingDirectory() string {
 	return pwd
 }
 
-// Update handles messages and updates the component state
+// Update handles messages and updates the component state.
 func (m *model) Update(msg tea.Msg) (layout.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case tea.WindowSizeMsg:
 		cmd := m.SetSize(msg.Width, msg.Height)
 		return m, cmd
+	case *runtime.TokenUsageEvent:
+		m.SetTokenUsage(msg)
+		return m, nil
 	case *runtime.MCPInitStartedEvent:
 		m.mcpInit = true
 		return m, m.spinner.Init()
@@ -272,13 +302,14 @@ func (m *model) View() string {
 
 func (m *model) horizontalView() string {
 	pwd := getCurrentWorkingDirectory()
+	usageSummary := m.tokenUsageSummary()
 
 	wi := m.workingIndicatorHorizontal()
 	titleGapWidth := m.width - lipgloss.Width(m.sessionTitle) - lipgloss.Width(wi) - 2
 	title := fmt.Sprintf("%s%*s%s", m.sessionTitle, titleGapWidth, "", wi)
 
-	gapWidth := m.width - lipgloss.Width(pwd) - lipgloss.Width(m.tokenUsage()) - 2
-	return lipgloss.JoinVertical(lipgloss.Top, title, fmt.Sprintf("%s%*s%s", styles.MutedStyle.Render(pwd), gapWidth, "", m.tokenUsage()))
+	gapWidth := m.width - lipgloss.Width(pwd) - lipgloss.Width(usageSummary) - 2
+	return lipgloss.JoinVertical(lipgloss.Top, title, fmt.Sprintf("%s%*s%s", styles.MutedStyle.Render(pwd), gapWidth, "", usageSummary))
 }
 
 func (m *model) verticalView() string {
@@ -473,17 +504,85 @@ func (m *model) workingIndicatorHorizontal() string {
 }
 
 func (m *model) tokenUsage() string {
-	totalTokens := m.usage.InputTokens + m.usage.OutputTokens
-	var usagePercent float64
-	if m.usage.ContextLimit > 0 {
-		usagePercent = (float64(m.usage.ContextLength) / float64(m.usage.ContextLimit)) * 100
+	if len(m.sessionUsage) == 0 {
+		return ""
+	}
+
+	var totalTokens int
+	var totalCost float64
+	for _, usage := range m.sessionUsage {
+		totalTokens += usage.InputTokens + usage.OutputTokens
+		totalCost += usage.Cost
+	}
+
+	agentTotals := make(map[string]*runtime.Usage)
+	for sessionID, usage := range m.sessionUsage {
+		agent := m.sessionAgent[sessionID]
+		if agent == "" {
+			continue
+		}
+		if existing, ok := agentTotals[agent]; ok {
+			existing.InputTokens += usage.InputTokens
+			existing.OutputTokens += usage.OutputTokens
+			existing.Cost += usage.Cost
+		} else {
+			u := *usage
+			agentTotals[agent] = &u
+		}
 	}
 
-	percentageText := styles.MutedStyle.Render(fmt.Sprintf("%.0f%%", usagePercent))
-	totalTokensText := styles.SubtleStyle.Render(fmt.Sprintf("(%s)", formatTokenCount(totalTokens)))
-	costText := styles.MutedStyle.Render(fmt.Sprintf("$%.2f", m.usage.Cost))
+	var b strings.Builder
+	b.WriteString(styles.HighlightStyle.Render("TOTAL USAGE"))
+	b.WriteString("\n  ")
+	line := fmt.Sprintf("Tokens: %s | Cost: $%s", formatTokenCount(totalTokens), formatCost(totalCost))
+	if ctxText, ok := m.contextPercent(); ok {
+		line = fmt.Sprintf("%s | %s", line, ctxText)
+	}
+	b.WriteString(line)
+
+	b.WriteString("\n--------------------------------\n")
+	b.WriteString(styles.HighlightStyle.Render("SESSION BREAKDOWN"))
+
+	agentNames := make([]string, 0, len(agentTotals))
+	for name := range agentTotals {
+		agentNames = append(agentNames, name)
+	}
+	sort.Strings(agentNames)
+
+	if len(agentNames) == 0 {
+		b.WriteString("\n  ")
+		b.WriteString(styles.SubtleStyle.Render("No usage yet"))
+		return b.String()
+	}
+
+	for _, name := range agentNames {
+		usage := agentTotals[name]
+		tokens := usage.InputTokens + usage.OutputTokens
+		b.WriteString(fmt.Sprintf("\n  %s", styles.SubtleStyle.Render(name)))
+		b.WriteString(fmt.Sprintf("\n    Tokens: %s | Cost: $%s", formatTokenCount(tokens), formatCost(usage.Cost)))
+	}
+
+	return b.String()
+}
+
+// tokenUsageSummary returns a single-line summary for horizontal layout.
+func (m *model) tokenUsageSummary() string {
+	if len(m.sessionUsage) == 0 {
+		return ""
+	}
+
+	var totalTokens int
+	var totalCost float64
+	for _, usage := range m.sessionUsage {
+		totalTokens += usage.InputTokens + usage.OutputTokens
+		totalCost += usage.Cost
+	}
+
+	if ctxText, ok := m.contextPercent(); ok {
+		return fmt.Sprintf("Tokens: %s | Cost: $%s | %s", formatTokenCount(totalTokens), formatCost(totalCost), ctxText)
+	}
 
-	return fmt.Sprintf("%s %s %s", percentageText, totalTokensText, costText)
+	return fmt.Sprintf("Tokens: %s | Cost: $%s", formatTokenCount(totalTokens), formatCost(totalCost))
 }
 
 // agentInfo renders the current agent information

pkg/tui/page/chat/chat.go

@@ -247,7 +247,9 @@ func (p *chatPage) Update(msg tea.Msg) (layout.Model, tea.Cmd) {
 		spinnerCmd := p.setWorking(true)
 		cmd := p.messages.AddAssistantMessage()
 		p.startProgressBar()
-		return p, tea.Batch(cmd, spinnerCmd)
+		sidebarModel, sidebarCmd := p.sidebar.Update(msg)
+		p.sidebar = sidebarModel.(sidebar.Model)
+		return p, tea.Batch(cmd, spinnerCmd, sidebarCmd)
 	case *runtime.AgentChoiceEvent:
 		if p.streamCancelled {
 			return p, nil
@@ -261,7 +263,7 @@ func (p *chatPage) Update(msg tea.Msg) (layout.Model, tea.Cmd) {
 		cmd := p.messages.AppendToLastMessage(msg.AgentName, types.MessageTypeAssistantReasoning, msg.Content)
 		return p, cmd
 	case *runtime.TokenUsageEvent:
-		p.sidebar.SetTokenUsage(msg.Usage)
+		p.sidebar.SetTokenUsage(msg)
 	case *runtime.AgentInfoEvent:
 		p.sidebar.SetAgentInfo(msg.AgentName, msg.Model, msg.Description)
 	case *runtime.TeamInfoEvent:
@@ -279,7 +281,13 @@ func (p *chatPage) Update(msg tea.Msg) (layout.Model, tea.Cmd) {
 		}
 		p.streamCancelled = false
 		p.stopProgressBar()
-		return p, tea.Batch(p.messages.ScrollToBottom(), spinnerCmd)
+		sidebarModel, sidebarCmd := p.sidebar.Update(msg)
+		p.sidebar = sidebarModel.(sidebar.Model)
+		return p, tea.Batch(p.messages.ScrollToBottom(), spinnerCmd, sidebarCmd)
+	case *runtime.SessionTitleEvent:
+		sidebarModel, sidebarCmd := p.sidebar.Update(msg)
+		p.sidebar = sidebarModel.(sidebar.Model)
+		return p, sidebarCmd
 	case *runtime.PartialToolCallEvent:
 		// When we first receive a tool call, show it immediately in pending state
 		spinnerCmd := p.setWorking(true)