implement nested tool calls and initial setup for result metadata

12 files changed, 549 insertions, 138 deletions

diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index 998dc1551..b01ffec3c 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -305,6 +305,7 @@ func (c *agent) generate(ctx context.Context, sessionID string, content string)
 		assistantMsg, err := c.Messages.Create(sessionID, message.CreateMessageParams{
 			Role:  message.Assistant,
 			Parts: []message.ContentPart{},
+			Model: c.model.ID,
 		})
 		if err != nil {
 			return err
diff --git a/internal/llm/tools/bash.go b/internal/llm/tools/bash.go
index 4e80ae60a..d20afb7f2 100644
--- a/internal/llm/tools/bash.go
+++ b/internal/llm/tools/bash.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/kujtimiihoxha/termai/internal/config"
 	"github.com/kujtimiihoxha/termai/internal/llm/tools/shell"
@@ -21,6 +22,9 @@ type BashPermissionsParams struct {
 	Timeout int    `json:"timeout"`
 }
 
+type BashToolResponseMetadata struct {
+	Took int64 `json:"took"`
+}
 type bashTool struct {
 	permissions permission.Service
 }
@@ -272,11 +276,13 @@ func (b *bashTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
 			return NewTextErrorResponse("permission denied"), nil
 		}
 	}
+	startTime := time.Now()
 	shell := shell.GetPersistentShell(config.WorkingDirectory())
 	stdout, stderr, exitCode, interrupted, err := shell.Exec(ctx, params.Command, params.Timeout)
 	if err != nil {
 		return NewTextErrorResponse(fmt.Sprintf("error executing command: %s", err)), nil
 	}
+	took := time.Since(startTime).Milliseconds()
 
 	stdout = truncateOutput(stdout)
 	stderr = truncateOutput(stderr)
@@ -304,10 +310,13 @@ func (b *bashTool) Run(ctx context.Context, call ToolCall) (ToolResponse, error)
 		stdout += "\n" + errorMessage
 	}
 
+	metadata := BashToolResponseMetadata{
+		Took: took,
+	}
 	if stdout == "" {
-		return NewTextResponse("no output"), nil
+		return WithResponseMetadata(NewTextResponse("no output"), metadata), nil
 	}
-	return NewTextResponse(stdout), nil
+	return WithResponseMetadata(NewTextResponse(stdout), metadata), nil
 }
 
 func truncateOutput(content string) string {
diff --git a/internal/llm/tools/tools.go b/internal/llm/tools/tools.go
index e15c1c31f..6bb528686 100644
--- a/internal/llm/tools/tools.go
+++ b/internal/llm/tools/tools.go
@@ -1,6 +1,9 @@
 package tools
 
-import "context"
+import (
+	"context"
+	"encoding/json"
+)
 
 type ToolInfo struct {
 	Name        string
@@ -17,9 +20,10 @@ const (
 )
 
 type ToolResponse struct {
-	Type    toolResponseType `json:"type"`
-	Content string           `json:"content"`
-	IsError bool             `json:"is_error"`
+	Type     toolResponseType `json:"type"`
+	Content  string           `json:"content"`
+	Metadata string           `json:"metadata,omitempty"`
+	IsError  bool             `json:"is_error"`
 }
 
 func NewTextResponse(content string) ToolResponse {
@@ -29,6 +33,17 @@ func NewTextResponse(content string) ToolResponse {
 	}
 }
 
+func WithResponseMetadata(response ToolResponse, metadata any) ToolResponse {
+	if metadata != nil {
+		metadataBytes, err := json.Marshal(metadata)
+		if err != nil {
+			return response
+		}
+		response.Metadata = string(metadataBytes)
+	}
+	return response
+}
+
 func NewTextErrorResponse(content string) ToolResponse {
 	return ToolResponse{
 		Type:    ToolResponseTypeText,
diff --git a/internal/message/content.go b/internal/message/content.go
index cd263798b..422c04f52 100644
--- a/internal/message/content.go
+++ b/internal/message/content.go
@@ -3,6 +3,8 @@ package message
 import (
 	"encoding/base64"
 	"time"
+
+	"github.com/kujtimiihoxha/termai/internal/llm/models"
 )
 
 type MessageRole string
@@ -65,7 +67,6 @@ type ToolCall struct {
 	Name     string `json:"name"`
 	Input    string `json:"input"`
 	Type     string `json:"type"`
-	Metadata any    `json:"metadata"`
 	Finished bool   `json:"finished"`
 }
 
@@ -75,6 +76,7 @@ type ToolResult struct {
 	ToolCallID string `json:"tool_call_id"`
 	Name       string `json:"name"`
 	Content    string `json:"content"`
+	Metadata   string `json:"metadata"`
 	IsError    bool   `json:"is_error"`
 }
 
@@ -92,6 +94,7 @@ type Message struct {
 	Role      MessageRole
 	SessionID string
 	Parts     []ContentPart
+	Model     models.ModelID
 
 	CreatedAt int64
 	UpdatedAt int64
diff --git a/internal/message/message.go b/internal/message/message.go
index eeeb83ed2..06dae13a5 100644
--- a/internal/message/message.go
+++ b/internal/message/message.go
@@ -155,6 +155,7 @@ func (s *service) fromDBItem(item db.Message) (Message, error) {
 		SessionID: item.SessionID,
 		Role:      MessageRole(item.Role),
 		Parts:     parts,
+		Model:     models.ModelID(item.Model.String),
 		CreatedAt: item.CreatedAt,
 		UpdatedAt: item.UpdatedAt,
 	}, nil
diff --git a/internal/tui/components/chat/editor.go b/internal/tui/components/chat/editor.go
index df336818c..e87f1ffae 100644
--- a/internal/tui/components/chat/editor.go
+++ b/internal/tui/components/chat/editor.go
@@ -77,21 +77,20 @@ func (m *editorCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	case AgentWorkingMsg:
 		m.agentWorking = bool(msg)
 	case tea.KeyMsg:
-		if key.Matches(msg, focusedKeyMaps.Send) {
+		// if the key does not match any binding, return
+		if m.textarea.Focused() && key.Matches(msg, focusedKeyMaps.Send) {
 			return m, m.send()
 		}
-		if key.Matches(msg, bluredKeyMaps.Send) {
+		if !m.textarea.Focused() && key.Matches(msg, bluredKeyMaps.Send) {
 			return m, m.send()
 		}
-		if key.Matches(msg, focusedKeyMaps.Blur) {
+		if m.textarea.Focused() && key.Matches(msg, focusedKeyMaps.Blur) {
 			m.textarea.Blur()
 			return m, util.CmdHandler(EditorFocusMsg(false))
 		}
-		if key.Matches(msg, bluredKeyMaps.Focus) {
-			if !m.textarea.Focused() {
-				m.textarea.Focus()
-				return m, tea.Batch(textarea.Blink, util.CmdHandler(EditorFocusMsg(true)))
-			}
+		if !m.textarea.Focused() && key.Matches(msg, bluredKeyMaps.Focus) {
+			m.textarea.Focus()
+			return m, tea.Batch(textarea.Blink, util.CmdHandler(EditorFocusMsg(true)))
 		}
 	}
 	m.textarea, cmd = m.textarea.Update(msg)
diff --git a/internal/tui/components/chat/messages.go b/internal/tui/components/chat/messages.go
index 0a7e6e2a4..b5a361392 100644
--- a/internal/tui/components/chat/messages.go
+++ b/internal/tui/components/chat/messages.go
@@ -1,16 +1,21 @@
 package chat
 
 import (
+	"encoding/json"
 	"fmt"
-	"regexp"
-	"strconv"
+	"math"
 	"strings"
 
+	"github.com/charmbracelet/bubbles/spinner"
 	"github.com/charmbracelet/bubbles/viewport"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/glamour"
 	"github.com/charmbracelet/lipgloss"
+	"github.com/charmbracelet/x/ansi"
 	"github.com/kujtimiihoxha/termai/internal/app"
+	"github.com/kujtimiihoxha/termai/internal/llm/agent"
+	"github.com/kujtimiihoxha/termai/internal/llm/models"
+	"github.com/kujtimiihoxha/termai/internal/llm/tools"
 	"github.com/kujtimiihoxha/termai/internal/message"
 	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/session"
@@ -18,10 +23,20 @@ import (
 	"github.com/kujtimiihoxha/termai/internal/tui/util"
 )
 
+type uiMessageType int
+
+const (
+	userMessageType uiMessageType = iota
+	assistantMessageType
+	toolMessageType
+)
+
 type uiMessage struct {
-	position int
-	height   int
-	content  string
+	ID          string
+	messageType uiMessageType
+	position    int
+	height      int
+	content     string
 }
 
 type messagesCmp struct {
@@ -32,141 +47,116 @@ type messagesCmp struct {
 	session       session.Session
 	messages      []message.Message
 	uiMessages    []uiMessage
-	currentIndex  int
+	currentMsgID  string
 	renderer      *glamour.TermRenderer
 	focusRenderer *glamour.TermRenderer
 	cachedContent map[string]string
+	agentWorking  bool
+	spinner       spinner.Model
+	needsRerender bool
+	lastViewport  string
 }
 
 func (m *messagesCmp) Init() tea.Cmd {
-	return m.viewport.Init()
-}
-
-var ansiEscape = regexp.MustCompile("\x1b\\[[0-9;]*m")
-
-func hexToBgSGR(hex string) (string, error) {
-	hex = strings.TrimPrefix(hex, "#")
-	if len(hex) != 6 {
-		return "", fmt.Errorf("invalid hex color: must be 6 hexadecimal digits")
-	}
-
-	// Parse RGB components in one block
-	rgb := make([]uint64, 3)
-	for i := 0; i < 3; i++ {
-		val, err := strconv.ParseUint(hex[i*2:i*2+2], 16, 8)
-		if err != nil {
-			return "", err
-		}
-		rgb[i] = val
-	}
-
-	return fmt.Sprintf("48;2;%d;%d;%d", rgb[0], rgb[1], rgb[2]), nil
-}
-
-func forceReplaceBackgroundColors(input string, newBg string) string {
-	return ansiEscape.ReplaceAllStringFunc(input, func(seq string) string {
-		// Extract content between "\x1b[" and "m"
-		content := seq[2 : len(seq)-1]
-		tokens := strings.Split(content, ";")
-		var newTokens []string
-
-		// Skip background color tokens
-		for i := 0; i < len(tokens); i++ {
-			if tokens[i] == "" {
-				continue
-			}
-
-			val, err := strconv.Atoi(tokens[i])
-			if err != nil {
-				newTokens = append(newTokens, tokens[i])
-				continue
-			}
-
-			// Skip background color tokens
-			if val == 48 {
-				// Skip "48;5;N" or "48;2;R;G;B" sequences
-				if i+1 < len(tokens) {
-					if nextVal, err := strconv.Atoi(tokens[i+1]); err == nil {
-						switch nextVal {
-						case 5:
-							i += 2 // Skip "5" and color index
-						case 2:
-							i += 4 // Skip "2" and RGB components
-						}
-					}
-				}
-			} else if (val < 40 || val > 47) && (val < 100 || val > 107) && val != 49 {
-				// Keep non-background tokens
-				newTokens = append(newTokens, tokens[i])
-			}
-		}
-
-		// Add new background if provided
-		if newBg != "" {
-			newTokens = append(newTokens, strings.Split(newBg, ";")...)
-		}
-
-		if len(newTokens) == 0 {
-			return ""
-		}
-
-		return "\x1b[" + strings.Join(newTokens, ";") + "m"
-	})
+	return tea.Batch(m.viewport.Init())
 }
 
 func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	var cmds []tea.Cmd
 	switch msg := msg.(type) {
+	case AgentWorkingMsg:
+		m.agentWorking = bool(msg)
+		if m.agentWorking {
+			cmds = append(cmds, m.spinner.Tick)
+		}
 	case EditorFocusMsg:
 		m.writingMode = bool(msg)
 	case SessionSelectedMsg:
 		if msg.ID != m.session.ID {
 			cmd := m.SetSession(msg)
+			m.needsRerender = true
 			return m, cmd
 		}
 		return m, nil
+	case SessionClearedMsg:
+		m.session = session.Session{}
+		m.messages = make([]message.Message, 0)
+		m.currentMsgID = ""
+		m.needsRerender = true
+		return m, nil
+
+	case tea.KeyMsg:
+		if m.writingMode {
+			return m, nil
+		}
 	case pubsub.Event[message.Message]:
 		if msg.Type == pubsub.CreatedEvent {
 			if msg.Payload.SessionID == m.session.ID {
 				// check if message exists
+
+				messageExists := false
 				for _, v := range m.messages {
 					if v.ID == msg.Payload.ID {
-						return m, nil
+						messageExists = true
+						break
 					}
 				}
 
-				m.messages = append(m.messages, msg.Payload)
-				m.renderView()
-				m.viewport.GotoBottom()
+				if !messageExists {
+					m.messages = append(m.messages, msg.Payload)
+					delete(m.cachedContent, m.currentMsgID)
+					m.currentMsgID = msg.Payload.ID
+					m.needsRerender = true
+				}
 			}
 			for _, v := range m.messages {
 				for _, c := range v.ToolCalls() {
 					// the message is being added to the session of a tool called
 					if c.ID == msg.Payload.SessionID {
-						m.renderView()
-						m.viewport.GotoBottom()
+						m.needsRerender = true
 					}
 				}
 			}
 		} else if msg.Type == pubsub.UpdatedEvent && msg.Payload.SessionID == m.session.ID {
 			for i, v := range m.messages {
 				if v.ID == msg.Payload.ID {
+					if !m.messages[i].IsFinished() && msg.Payload.IsFinished() && msg.Payload.FinishReason() == "end_turn" || msg.Payload.FinishReason() == "canceled" {
+						cmds = append(cmds, util.CmdHandler(AgentWorkingMsg(false)))
+					}
 					m.messages[i] = msg.Payload
 					delete(m.cachedContent, msg.Payload.ID)
-					m.renderView()
-					if i == len(m.messages)-1 {
-						m.viewport.GotoBottom()
-					}
+					m.needsRerender = true
 					break
 				}
 			}
 		}
 	}
+	if m.agentWorking {
+		u, cmd := m.spinner.Update(msg)
+		m.spinner = u
+		cmds = append(cmds, cmd)
+	}
+	oldPos := m.viewport.YPosition
 	u, cmd := m.viewport.Update(msg)
 	m.viewport = u
-	return m, cmd
+	m.needsRerender = m.needsRerender || m.viewport.YPosition != oldPos
+	cmds = append(cmds, cmd)
+	if m.needsRerender {
+		m.renderView()
+		if len(m.messages) > 0 {
+			if msg, ok := msg.(pubsub.Event[message.Message]); ok {
+				if (msg.Type == pubsub.CreatedEvent) ||
+					(msg.Type == pubsub.UpdatedEvent && msg.Payload.ID == m.messages[len(m.messages)-1].ID) {
+					m.viewport.GotoBottom()
+				}
+			}
+		}
+		m.needsRerender = false
+	}
+	return m, tea.Batch(cmds...)
 }
 
-func (m *messagesCmp) renderUserMessage(inx int, msg message.Message) string {
+func (m *messagesCmp) renderSimpleMessage(msg message.Message, info ...string) string {
 	if v, ok := m.cachedContent[msg.ID]; ok {
 		return v
 	}
@@ -178,7 +168,7 @@ func (m *messagesCmp) renderUserMessage(inx int, msg message.Message) string {
 		BorderStyle(lipgloss.ThickBorder())
 
 	renderer := m.renderer
-	if inx == m.currentIndex {
+	if msg.ID == m.currentMsgID {
 		style = style.
 			Foreground(styles.Forground).
 			BorderForeground(styles.Blue).
@@ -186,33 +176,269 @@ func (m *messagesCmp) renderUserMessage(inx int, msg message.Message) string {
 		renderer = m.focusRenderer
 	}
 	c, _ := renderer.Render(msg.Content().String())
-	col, _ := hexToBgSGR(styles.Background.Dark)
-	rendered := style.Render(forceReplaceBackgroundColors(c, col))
+	parts := []string{
+		styles.ForceReplaceBackgroundWithLipgloss(c, styles.Background),
+	}
+	// remove newline at the end
+	parts[0] = strings.TrimSuffix(parts[0], "\n")
+	if len(info) > 0 {
+		parts = append(parts, info...)
+	}
+	rendered := style.Render(
+		lipgloss.JoinVertical(
+			lipgloss.Left,
+			parts...,
+		),
+	)
 	m.cachedContent[msg.ID] = rendered
 	return rendered
 }
 
+func formatTimeDifference(unixTime1, unixTime2 int64) string {
+	diffSeconds := float64(math.Abs(float64(unixTime2 - unixTime1)))
+
+	if diffSeconds < 60 {
+		return fmt.Sprintf("%.1fs", diffSeconds)
+	}
+
+	minutes := int(diffSeconds / 60)
+	seconds := int(diffSeconds) % 60
+	return fmt.Sprintf("%dm%ds", minutes, seconds)
+}
+
+func (m *messagesCmp) renderToolCall(toolCall message.ToolCall, isNested bool) string {
+	key := ""
+	value := ""
+	switch toolCall.Name {
+	// TODO: add result data to the tools
+	case agent.AgentToolName:
+		key = "Task"
+		var params agent.AgentParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		value = params.Prompt
+	// TODO: handle nested calls
+	case tools.BashToolName:
+		key = "Bash"
+		var params tools.BashParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		value = params.Command
+	case tools.EditToolName:
+		key = "Edit"
+		var params tools.EditParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		value = params.FilePath
+	case tools.FetchToolName:
+		key = "Fetch"
+		var params tools.FetchParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		value = params.URL
+	case tools.GlobToolName:
+		key = "Glob"
+		var params tools.GlobParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		if params.Path == "" {
+			params.Path = "."
+		}
+		value = fmt.Sprintf("%s (%s)", params.Pattern, params.Path)
+	case tools.GrepToolName:
+		key = "Grep"
+		var params tools.GrepParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		if params.Path == "" {
+			params.Path = "."
+		}
+		value = fmt.Sprintf("%s (%s)", params.Pattern, params.Path)
+	case tools.LSToolName:
+		key = "Ls"
+		var params tools.LSParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		if params.Path == "" {
+			params.Path = "."
+		}
+		value = params.Path
+	case tools.SourcegraphToolName:
+		key = "Sourcegraph"
+		var params tools.SourcegraphParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		value = params.Query
+	case tools.ViewToolName:
+		key = "View"
+		var params tools.ViewParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		value = params.FilePath
+	case tools.WriteToolName:
+		key = "Write"
+		var params tools.WriteParams
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		value = params.FilePath
+	default:
+		key = toolCall.Name
+		var params map[string]any
+		json.Unmarshal([]byte(toolCall.Input), &params)
+		jsonData, _ := json.Marshal(params)
+		value = string(jsonData)
+	}
+
+	style := styles.BaseStyle.
+		Width(m.width).
+		BorderLeft(true).
+		BorderStyle(lipgloss.ThickBorder()).
+		PaddingLeft(1).
+		BorderForeground(styles.Yellow)
+
+	keyStyle := styles.BaseStyle.
+		Foreground(styles.ForgroundDim)
+	valyeStyle := styles.BaseStyle.
+		Foreground(styles.Forground)
+
+	if isNested {
+		valyeStyle = valyeStyle.Foreground(styles.ForgroundMid)
+	}
+	keyValye := keyStyle.Render(
+		fmt.Sprintf("%s: ", key),
+	)
+	if !isNested {
+		value = valyeStyle.
+			Width(m.width - lipgloss.Width(keyValye) - 2).
+			Render(
+				ansi.Truncate(
+					value,
+					m.width-lipgloss.Width(keyValye)-2,
+					"...",
+				),
+			)
+	} else {
+		keyValye = keyStyle.Render(
+			fmt.Sprintf(" └ %s: ", key),
+		)
+		value = valyeStyle.
+			Width(m.width - lipgloss.Width(keyValye) - 2).
+			Render(
+				ansi.Truncate(
+					value,
+					m.width-lipgloss.Width(keyValye)-2,
+					"...",
+				),
+			)
+	}
+
+	innerToolCalls := make([]string, 0)
+	if toolCall.Name == agent.AgentToolName {
+		messages, _ := m.app.Messages.List(toolCall.ID)
+		toolCalls := make([]message.ToolCall, 0)
+		for _, v := range messages {
+			toolCalls = append(toolCalls, v.ToolCalls()...)
+		}
+		for _, v := range toolCalls {
+			call := m.renderToolCall(v, true)
+			innerToolCalls = append(innerToolCalls, call)
+		}
+	}
+
+	if isNested {
+		return lipgloss.JoinHorizontal(
+			lipgloss.Left,
+			keyValye,
+			value,
+		)
+	}
+	callContent := lipgloss.JoinHorizontal(
+		lipgloss.Left,
+		keyValye,
+		value,
+	)
+	callContent = strings.ReplaceAll(callContent, "\n", "")
+	if len(innerToolCalls) > 0 {
+		callContent = lipgloss.JoinVertical(
+			lipgloss.Left,
+			callContent,
+			lipgloss.JoinVertical(
+				lipgloss.Left,
+				innerToolCalls...,
+			),
+		)
+	}
+	return style.Render(callContent)
+}
+
+func (m *messagesCmp) renderAssistantMessage(msg message.Message) []uiMessage {
+	// find the user message that is before this assistant message
+	var userMsg message.Message
+	for i := len(m.messages) - 1; i >= 0; i-- {
+		if m.messages[i].Role == message.User {
+			userMsg = m.messages[i]
+			break
+		}
+	}
+	messages := make([]uiMessage, 0)
+	if msg.Content().String() != "" {
+		info := make([]string, 0)
+		if msg.IsFinished() && msg.FinishReason() == "end_turn" {
+			finish := msg.FinishPart()
+			took := formatTimeDifference(userMsg.CreatedAt, finish.Time)
+
+			info = append(info, styles.BaseStyle.Width(m.width-1).Foreground(styles.ForgroundDim).Render(
+				fmt.Sprintf(" %s (%s)", models.SupportedModels[msg.Model].Name, took),
+			))
+		}
+		content := m.renderSimpleMessage(msg, info...)
+		messages = append(messages, uiMessage{
+			messageType: assistantMessageType,
+			position:    0, // gets updated in renderView
+			height:      lipgloss.Height(content),
+			content:     content,
+		})
+	}
+	for _, v := range msg.ToolCalls() {
+		content := m.renderToolCall(v, false)
+		messages = append(messages,
+			uiMessage{
+				messageType: toolMessageType,
+				position:    0, // gets updated in renderView
+				height:      lipgloss.Height(content),
+				content:     content,
+			},
+		)
+	}
+
+	return messages
+}
+
 func (m *messagesCmp) renderView() {
 	m.uiMessages = make([]uiMessage, 0)
 	pos := 0
 
 	for _, v := range m.messages {
-		content := ""
 		switch v.Role {
 		case message.User:
-			content = m.renderUserMessage(pos, v)
+			content := m.renderSimpleMessage(v)
+			m.uiMessages = append(m.uiMessages, uiMessage{
+				messageType: userMessageType,
+				position:    pos,
+				height:      lipgloss.Height(content),
+				content:     content,
+			})
+			pos += lipgloss.Height(content) + 1 // + 1 for spacing
+		case message.Assistant:
+			assistantMessages := m.renderAssistantMessage(v)
+			for _, msg := range assistantMessages {
+				msg.position = pos
+				m.uiMessages = append(m.uiMessages, msg)
+				pos += msg.height + 1 // + 1 for spacing
+			}
+
 		}
-		m.uiMessages = append(m.uiMessages, uiMessage{
-			position: pos,
-			height:   lipgloss.Height(content),
-			content:  content,
-		})
-		pos += lipgloss.Height(content) + 1 // + 1 for spacing
 	}
 
 	messages := make([]string, 0)
 	for _, v := range m.uiMessages {
-		messages = append(messages, v.content)
+		messages = append(messages, v.content,
+			styles.BaseStyle.
+				Width(m.width).
+				Render(
+					"",
+				),
+		)
 	}
 	m.viewport.SetContent(
 		styles.BaseStyle.
@@ -246,7 +472,6 @@ func (m *messagesCmp) View() string {
 			)
 	}
 
-	m.renderView()
 	return styles.BaseStyle.
 		Width(m.width).
 		Render(
@@ -260,15 +485,21 @@ func (m *messagesCmp) View() string {
 
 func (m *messagesCmp) help() string {
 	text := ""
+
+	if m.agentWorking {
+		text += styles.BaseStyle.Foreground(styles.PrimaryColor).Bold(true).Render(
+			fmt.Sprintf("%s %s ", m.spinner.View(), "Generating..."),
+		)
+	}
 	if m.writingMode {
-		text = lipgloss.JoinHorizontal(
+		text += lipgloss.JoinHorizontal(
 			lipgloss.Left,
 			styles.BaseStyle.Foreground(styles.ForgroundDim).Bold(true).Render("press "),
 			styles.BaseStyle.Foreground(styles.Forground).Bold(true).Render("esc"),
 			styles.BaseStyle.Foreground(styles.ForgroundDim).Bold(true).Render(" to exit writing mode"),
 		)
 	} else {
-		text = lipgloss.JoinHorizontal(
+		text += lipgloss.JoinHorizontal(
 			lipgloss.Left,
 			styles.BaseStyle.Foreground(styles.ForgroundDim).Bold(true).Render("press "),
 			styles.BaseStyle.Foreground(styles.Forground).Bold(true).Render("i"),
@@ -306,7 +537,15 @@ func (m *messagesCmp) SetSize(width, height int) {
 		glamour.WithWordWrap(width-1),
 	)
 	m.focusRenderer = focusRenderer
+	// clear the cached content
+	for k := range m.cachedContent {
+		delete(m.cachedContent, k)
+	}
 	m.renderer = renderer
+	if len(m.messages) > 0 {
+		m.renderView()
+		m.viewport.GotoBottom()
+	}
 }
 
 func (m *messagesCmp) GetSize() (int, int) {
@@ -320,7 +559,8 @@ func (m *messagesCmp) SetSession(session session.Session) tea.Cmd {
 		return util.ReportError(err)
 	}
 	m.messages = messages
-	m.messages = append(m.messages, m.messages[0])
+	m.currentMsgID = m.messages[len(m.messages)-1].ID
+	m.needsRerender = true
 	return nil
 }
 
@@ -333,6 +573,9 @@ func NewMessagesCmp(app *app.App) tea.Model {
 		glamour.WithStyles(styles.MarkdownTheme(false)),
 		glamour.WithWordWrap(80),
 	)
+
+	s := spinner.New()
+	s.Spinner = spinner.Pulse
 	return &messagesCmp{
 		app:           app,
 		writingMode:   true,
@@ -340,5 +583,6 @@ func NewMessagesCmp(app *app.App) tea.Model {
 		viewport:      viewport.New(0, 0),
 		focusRenderer: focusRenderer,
 		renderer:      renderer,
+		spinner:       s,
 	}
 }
diff --git a/internal/tui/components/chat/sidebar.go b/internal/tui/components/chat/sidebar.go
index 65c06f4a1..51192cf9a 100644
--- a/internal/tui/components/chat/sidebar.go
+++ b/internal/tui/components/chat/sidebar.go
@@ -5,6 +5,7 @@ import (
 
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
+	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/session"
 	"github.com/kujtimiihoxha/termai/internal/tui/styles"
 )
@@ -19,6 +20,14 @@ func (m *sidebarCmp) Init() tea.Cmd {
 }
 
 func (m *sidebarCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case pubsub.Event[session.Session]:
+		if msg.Type == pubsub.UpdatedEvent {
+			if m.session.ID == msg.Payload.ID {
+				m.session = msg.Payload
+			}
+		}
+	}
 	return m, nil
 }
 
@@ -45,7 +54,7 @@ func (m *sidebarCmp) sessionSection() string {
 	sessionValue := styles.BaseStyle.
 		Foreground(styles.Forground).
 		Width(m.width - lipgloss.Width(sessionKey)).
-		Render(": New Session")
+		Render(fmt.Sprintf(": %s", m.session.Title))
 	return lipgloss.JoinHorizontal(
 		lipgloss.Left,
 		sessionKey,
diff --git a/internal/tui/page/chat.go b/internal/tui/page/chat.go
index 7ac0d2293..a7a51bb84 100644
--- a/internal/tui/page/chat.go
+++ b/internal/tui/page/chat.go
@@ -1,9 +1,10 @@
 package page
 
 import (
+	"github.com/charmbracelet/bubbles/key"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/kujtimiihoxha/termai/internal/app"
-	"github.com/kujtimiihoxha/termai/internal/message"
+	"github.com/kujtimiihoxha/termai/internal/llm/agent"
 	"github.com/kujtimiihoxha/termai/internal/session"
 	"github.com/kujtimiihoxha/termai/internal/tui/components/chat"
 	"github.com/kujtimiihoxha/termai/internal/tui/layout"
@@ -18,8 +19,32 @@ type chatPage struct {
 	session session.Session
 }
 
+type ChatKeyMap struct {
+	NewSession key.Binding
+}
+
+var keyMap = ChatKeyMap{
+	NewSession: key.NewBinding(
+		key.WithKeys("ctrl+n"),
+		key.WithHelp("ctrl+n", "new session"),
+	),
+}
+
 func (p *chatPage) Init() tea.Cmd {
-	return p.layout.Init()
+	// TODO: remove
+	cmds := []tea.Cmd{
+		p.layout.Init(),
+	}
+
+	sessions, _ := p.app.Sessions.List()
+	if len(sessions) > 0 {
+		p.session = sessions[0]
+		cmd := p.setSidebar()
+		cmds = append(cmds, util.CmdHandler(chat.SessionSelectedMsg(p.session)), cmd)
+	}
+	return tea.Batch(
+		cmds...,
+	)
 }
 
 func (p *chatPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
@@ -31,6 +56,13 @@ func (p *chatPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		if cmd != nil {
 			return p, cmd
 		}
+	case tea.KeyMsg:
+		switch {
+		case key.Matches(msg, keyMap.NewSession):
+			p.session = session.Session{}
+			p.clearSidebar()
+			return p, util.CmdHandler(chat.SessionClearedMsg{})
+		}
 	}
 	u, cmd := p.layout.Update(msg)
 	p.layout = u.(layout.SplitPaneLayout)
@@ -51,6 +83,12 @@ func (p *chatPage) setSidebar() tea.Cmd {
 	return sidebarContainer.Init()
 }
 
+func (p *chatPage) clearSidebar() {
+	p.layout.SetRightPanel(nil)
+	width, height := p.layout.GetSize()
+	p.layout.SetSize(width, height)
+}
+
 func (p *chatPage) sendMessage(text string) tea.Cmd {
 	var cmds []tea.Cmd
 	if p.session.ID == "" {
@@ -66,15 +104,15 @@ func (p *chatPage) sendMessage(text string) tea.Cmd {
 		}
 		cmds = append(cmds, util.CmdHandler(chat.SessionSelectedMsg(session)))
 	}
-	// TODO: actually call agent
-	p.app.Messages.Create(p.session.ID, message.CreateMessageParams{
-		Role: message.User,
-		Parts: []message.ContentPart{
-			message.TextContent{
-				Text: text,
-			},
-		},
-	})
+	// TODO: move this to a service
+	a, err := agent.NewCoderAgent(p.app)
+	if err != nil {
+		return util.ReportError(err)
+	}
+	go func() {
+		a.Generate(p.app.Context, p.session.ID, text)
+	}()
+
 	return tea.Batch(cmds...)
 }
 
@@ -85,7 +123,7 @@ func (p *chatPage) View() string {
 func NewChatPage(app *app.App) tea.Model {
 	messagesContainer := layout.NewContainer(
 		chat.NewMessagesCmp(app),
-		layout.WithPadding(1, 1, 1, 1),
+		layout.WithPadding(1, 1, 0, 1),
 	)
 
 	editorContainer := layout.NewContainer(
diff --git a/internal/tui/styles/background.go b/internal/tui/styles/background.go
new file mode 100644
index 000000000..bf6cbc105
--- /dev/null
+++ b/internal/tui/styles/background.go
@@ -0,0 +1,81 @@
+package styles
+
+import (
+	"fmt"
+	"regexp"
+	"strconv"
+	"strings"
+
+	"github.com/charmbracelet/lipgloss"
+)
+
+var ansiEscape = regexp.MustCompile("\x1b\\[[0-9;]*m")
+
+func getColorRGB(c lipgloss.TerminalColor) (uint8, uint8, uint8) {
+	r, g, b, a := c.RGBA()
+
+	// Un-premultiply alpha if needed
+	if a > 0 && a < 0xffff {
+		r = (r * 0xffff) / a
+		g = (g * 0xffff) / a
+		b = (b * 0xffff) / a
+	}
+
+	// Convert from 16-bit to 8-bit color
+	return uint8(r >> 8), uint8(g >> 8), uint8(b >> 8)
+}
+
+func ForceReplaceBackgroundWithLipgloss(input string, newBgColor lipgloss.TerminalColor) string {
+	r, g, b := getColorRGB(newBgColor)
+
+	newBg := fmt.Sprintf("48;2;%d;%d;%d", r, g, b)
+
+	return ansiEscape.ReplaceAllStringFunc(input, func(seq string) string {
+		// Extract content between "\x1b[" and "m"
+		content := seq[2 : len(seq)-1]
+		tokens := strings.Split(content, ";")
+		var newTokens []string
+
+		// Skip background color tokens
+		for i := 0; i < len(tokens); i++ {
+			if tokens[i] == "" {
+				continue
+			}
+
+			val, err := strconv.Atoi(tokens[i])
+			if err != nil {
+				newTokens = append(newTokens, tokens[i])
+				continue
+			}
+
+			// Skip background color tokens
+			if val == 48 {
+				// Skip "48;5;N" or "48;2;R;G;B" sequences
+				if i+1 < len(tokens) {
+					if nextVal, err := strconv.Atoi(tokens[i+1]); err == nil {
+						switch nextVal {
+						case 5:
+							i += 2 // Skip "5" and color index
+						case 2:
+							i += 4 // Skip "2" and RGB components
+						}
+					}
+				}
+			} else if (val < 40 || val > 47) && (val < 100 || val > 107) && val != 49 {
+				// Keep non-background tokens
+				newTokens = append(newTokens, tokens[i])
+			}
+		}
+
+		// Add new background if provided
+		if newBg != "" {
+			newTokens = append(newTokens, strings.Split(newBg, ";")...)
+		}
+
+		if len(newTokens) == 0 {
+			return ""
+		}
+
+		return "\x1b[" + strings.Join(newTokens, ";") + "m"
+	})
+}
diff --git a/internal/tui/styles/markdown.go b/internal/tui/styles/markdown.go
index b4e71c51e..52816eab3 100644
--- a/internal/tui/styles/markdown.go
+++ b/internal/tui/styles/markdown.go
@@ -515,6 +515,7 @@ var ASCIIStyleConfig = ansi.StyleConfig{
 	Document: ansi.StyleBlock{
 		StylePrimitive: ansi.StylePrimitive{
 			BackgroundColor: stringPtr(Background.Dark),
+			Color:           stringPtr(ForgroundDim.Dark),
 		},
 		Indent:      uintPtr(1),
 		IndentToken: stringPtr(BaseStyle.Render(" ")),
@@ -688,7 +689,7 @@ var DraculaStyleConfig = ansi.StyleConfig{
 	Heading: ansi.StyleBlock{
 		StylePrimitive: ansi.StylePrimitive{
 			BlockSuffix:     "\n",
-			Color:           stringPtr("#bd93f9"),
+			Color:           stringPtr(PrimaryColor.Dark),
 			Bold:            boolPtr(true),
 			BackgroundColor: stringPtr(Background.Dark),
 		},
@@ -740,7 +741,7 @@ var DraculaStyleConfig = ansi.StyleConfig{
 	},
 	Strong: ansi.StylePrimitive{
 		Bold:            boolPtr(true),
-		Color:           stringPtr("#ffb86c"),
+		Color:           stringPtr(Blue.Dark),
 		BackgroundColor: stringPtr(Background.Dark),
 	},
 	HorizontalRule: ansi.StylePrimitive{
@@ -796,7 +797,7 @@ var DraculaStyleConfig = ansi.StyleConfig{
 	CodeBlock: ansi.StyleCodeBlock{
 		StyleBlock: ansi.StyleBlock{
 			StylePrimitive: ansi.StylePrimitive{
-				Color:           stringPtr("#ffb86c"),
+				Color:           stringPtr(Blue.Dark),
 				BackgroundColor: stringPtr(Background.Dark),
 			},
 			Margin: uintPtr(defaultMargin),
diff --git a/internal/tui/styles/styles.go b/internal/tui/styles/styles.go
index 41863cf1b..476339b57 100644
--- a/internal/tui/styles/styles.go
+++ b/internal/tui/styles/styles.go
@@ -34,6 +34,11 @@ var (
 		Light: "#d3d3d3",
 	}
 
+	ForgroundMid = lipgloss.AdaptiveColor{
+		Dark:  "#a0a0a0",
+		Light: "#a0a0a0",
+	}
+
 	ForgroundDim = lipgloss.AdaptiveColor{
 		Dark:  "#737373",
 		Light: "#737373",
@@ -159,6 +164,11 @@ var (
 		Light: light.Peach().Hex,
 	}
 
+	Yellow = lipgloss.AdaptiveColor{
+		Dark:  dark.Yellow().Hex,
+		Light: light.Yellow().Hex,
+	}
+
 	Primary   = Blue
 	Secondary = Mauve