rework llm

8 files changed, 520 insertions, 276 deletions

diff --git a/internal/tui/components/core/dialog.go b/internal/tui/components/core/dialog.go
index 954e26e87..a8fef2e86 100644
--- a/internal/tui/components/core/dialog.go
+++ b/internal/tui/components/core/dialog.go
@@ -14,7 +14,12 @@ type SizeableModel interface {
 }
 
 type DialogMsg struct {
-	Content SizeableModel
+	Content     SizeableModel
+	WidthRatio  float64
+	HeightRatio float64
+
+	MinWidth  int
+	MinHeight int
 }
 
 type DialogCloseMsg struct{}
@@ -36,7 +41,18 @@ type DialogCmp interface {
 }
 
 type dialogCmp struct {
-	content SizeableModel
+	content      SizeableModel
+	screenWidth  int
+	screenHeight int
+
+	widthRatio  float64
+	heightRatio float64
+
+	minWidth  int
+	minHeight int
+
+	width  int
+	height int
 }
 
 func (d *dialogCmp) Init() tea.Cmd {
@@ -45,8 +61,26 @@ func (d *dialogCmp) Init() tea.Cmd {
 
 func (d *dialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		d.screenWidth = msg.Width
+		d.screenHeight = msg.Height
+		d.width = max(int(float64(d.screenWidth)*d.widthRatio), d.minWidth)
+		d.height = max(int(float64(d.screenHeight)*d.heightRatio), d.minHeight)
+		if d.content != nil {
+			d.content.SetSize(d.width, d.height)
+		}
+		return d, nil
 	case DialogMsg:
 		d.content = msg.Content
+		d.widthRatio = msg.WidthRatio
+		d.heightRatio = msg.HeightRatio
+		d.minWidth = msg.MinWidth
+		d.minHeight = msg.MinHeight
+		d.width = max(int(float64(d.screenWidth)*d.widthRatio), d.minWidth)
+		d.height = max(int(float64(d.screenHeight)*d.heightRatio), d.minHeight)
+		if d.content != nil {
+			d.content.SetSize(d.width, d.height)
+		}
 	case DialogCloseMsg:
 		d.content = nil
 		return d, nil
@@ -75,8 +109,7 @@ func (d *dialogCmp) BindingKeys() []key.Binding {
 }
 
 func (d *dialogCmp) View() string {
-	w, h := d.content.GetSize()
-	return lipgloss.NewStyle().Width(w).Height(h).Render(d.content.View())
+	return lipgloss.NewStyle().Width(d.width).Height(d.height).Render(d.content.View())
 }
 
 func NewDialogCmp() DialogCmp {
diff --git a/internal/tui/components/core/status.go b/internal/tui/components/core/status.go
index 8c9b3f60b..cd9dc55a0 100644
--- a/internal/tui/components/core/status.go
+++ b/internal/tui/components/core/status.go
@@ -3,6 +3,8 @@ package core
 import (
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
+	"github.com/kujtimiihoxha/termai/internal/config"
+	"github.com/kujtimiihoxha/termai/internal/llm/models"
 	"github.com/kujtimiihoxha/termai/internal/tui/styles"
 	"github.com/kujtimiihoxha/termai/internal/tui/util"
 	"github.com/kujtimiihoxha/termai/internal/version"
@@ -57,14 +59,19 @@ func (m statusCmp) View() string {
 			Width(m.availableFooterMsgWidth()).
 			Render(m.info)
 	}
-
+	status += m.model()
 	status += versionWidget
 	return status
 }
 
 func (m statusCmp) availableFooterMsgWidth() int {
 	// -2 to accommodate padding
-	return max(0, m.width-lipgloss.Width(helpWidget)-lipgloss.Width(versionWidget))
+	return max(0, m.width-lipgloss.Width(helpWidget)-lipgloss.Width(versionWidget)-lipgloss.Width(m.model()))
+}
+
+func (m statusCmp) model() string {
+	model := models.SupportedModels[config.Get().Model.Coder]
+	return styles.Padded.Background(styles.Grey).Foreground(styles.Text).Render(model.Name)
 }
 
 func NewStatusCmp() tea.Model {
diff --git a/internal/tui/components/dialog/permission.go b/internal/tui/components/dialog/permission.go
index f7581330a..7dcca9bae 100644
--- a/internal/tui/components/dialog/permission.go
+++ b/internal/tui/components/dialog/permission.go
@@ -1,9 +1,14 @@
 package dialog
 
 import (
+	"fmt"
+
 	"github.com/charmbracelet/bubbles/key"
+	"github.com/charmbracelet/bubbles/viewport"
 	tea "github.com/charmbracelet/bubbletea"
+	"github.com/charmbracelet/glamour"
 	"github.com/charmbracelet/lipgloss"
+	"github.com/kujtimiihoxha/termai/internal/llm/tools"
 	"github.com/kujtimiihoxha/termai/internal/permission"
 	"github.com/kujtimiihoxha/termai/internal/tui/components/core"
 	"github.com/kujtimiihoxha/termai/internal/tui/layout"
@@ -28,12 +33,6 @@ type PermissionResponseMsg struct {
 	Action     PermissionAction
 }
 
-// Width and height constants for the dialog
-var (
-	permissionWidth  = 60
-	permissionHeight = 10
-)
-
 // PermissionDialog interface for permission dialog component
 type PermissionDialog interface {
 	tea.Model
@@ -41,13 +40,28 @@ type PermissionDialog interface {
 	layout.Bindings
 }
 
+type keyMap struct {
+	ChangeFocus key.Binding
+}
+
+var keyMapValue = keyMap{
+	ChangeFocus: key.NewBinding(
+		key.WithKeys("tab"),
+		key.WithHelp("tab", "change focus"),
+	),
+}
+
 // permissionDialogCmp is the implementation of PermissionDialog
 type permissionDialogCmp struct {
-	form       *huh.Form
-	content    string
-	width      int
-	height     int
-	permission permission.PermissionRequest
+	form            *huh.Form
+	width           int
+	height          int
+	permission      permission.PermissionRequest
+	windowSize      tea.WindowSizeMsg
+	r               *glamour.TermRenderer
+	contentViewPort viewport.Model
+	isViewportFocus bool
+	selectOption    *huh.Select[string]
 }
 
 func (p *permissionDialogCmp) Init() tea.Cmd {
@@ -57,41 +71,101 @@ func (p *permissionDialogCmp) Init() tea.Cmd {
 func (p *permissionDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	var cmds []tea.Cmd
 
-	// Process the form
-	form, cmd := p.form.Update(msg)
-	if f, ok := form.(*huh.Form); ok {
-		p.form = f
-		cmds = append(cmds, cmd)
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		p.windowSize = msg
+	case tea.KeyMsg:
+		if key.Matches(msg, keyMapValue.ChangeFocus) {
+			p.isViewportFocus = !p.isViewportFocus
+			if p.isViewportFocus {
+				p.selectOption.Blur()
+			} else {
+				p.selectOption.Focus()
+			}
+			return p, nil
+		}
 	}
 
-	if p.form.State == huh.StateCompleted {
-		// Get the selected action
-		action := p.form.GetString("action")
-
-		// Close the dialog and return the response
-		return p, tea.Batch(
-			util.CmdHandler(core.DialogCloseMsg{}),
-			util.CmdHandler(PermissionResponseMsg{Action: PermissionAction(action), Permission: p.permission}),
-		)
+	if p.isViewportFocus {
+		viewPort, cmd := p.contentViewPort.Update(msg)
+		p.contentViewPort = viewPort
+		cmds = append(cmds, cmd)
+	} else {
+		form, cmd := p.form.Update(msg)
+		if f, ok := form.(*huh.Form); ok {
+			p.form = f
+			cmds = append(cmds, cmd)
+		}
+
+		if p.form.State == huh.StateCompleted {
+			// Get the selected action
+			action := p.form.GetString("action")
+
+			// Close the dialog and return the response
+			return p, tea.Batch(
+				util.CmdHandler(core.DialogCloseMsg{}),
+				util.CmdHandler(PermissionResponseMsg{Action: PermissionAction(action), Permission: p.permission}),
+			)
+		}
 	}
-
 	return p, tea.Batch(cmds...)
 }
 
-func (p *permissionDialogCmp) View() string {
-	contentStyle := lipgloss.NewStyle().
-		Width(p.width).
-		Padding(1, 0).
-		Foreground(styles.Text).
-		Align(lipgloss.Center)
+func (p *permissionDialogCmp) render() string {
+	form := p.form.View()
+	keyStyle := lipgloss.NewStyle().Bold(true).Foreground(styles.Rosewater)
+	valueStyle := lipgloss.NewStyle().Foreground(styles.Peach)
+
+	headerParts := []string{
+		lipgloss.JoinHorizontal(lipgloss.Left, keyStyle.Render("Tool:"), " ", valueStyle.Render(p.permission.ToolName)),
+		" ",
+		lipgloss.JoinHorizontal(lipgloss.Left, keyStyle.Render("Path:"), " ", valueStyle.Render(p.permission.Path)),
+		" ",
+	}
+	r, _ := glamour.NewTermRenderer(
+		glamour.WithStyles(styles.CatppuccinMarkdownStyle()),
+		glamour.WithWordWrap(p.width-10),
+		glamour.WithEmoji(),
+	)
+	content := ""
+	switch p.permission.ToolName {
+	case tools.BashToolName:
+		pr := p.permission.Params.(tools.BashPermissionsParams)
+		headerParts = append(headerParts, keyStyle.Render("Command:"))
+		content, _ = r.Render(fmt.Sprintf("```bash\n%s\n```", pr.Command))
+	case tools.EditToolName:
+		pr := p.permission.Params.(tools.EditPermissionsParams)
+		headerParts = append(headerParts, keyStyle.Render("Update:"))
+		content, _ = r.Render(fmt.Sprintf("```diff\n%s\n```", pr.Diff))
+	case tools.WriteToolName:
+		pr := p.permission.Params.(tools.WritePermissionsParams)
+		headerParts = append(headerParts, keyStyle.Render("Content:"))
+		content, _ = r.Render(fmt.Sprintf("```diff\n%s\n```", pr.Content))
+	default:
+		content, _ = r.Render(p.permission.Description)
+	}
+	headerContent := lipgloss.NewStyle().Padding(0, 1).Render(lipgloss.JoinVertical(lipgloss.Left, headerParts...))
+	p.contentViewPort.Width = p.width - 2 - 2
+	p.contentViewPort.Height = p.height - lipgloss.Height(headerContent) - lipgloss.Height(form) - 2 - 2 - 1
+	p.contentViewPort.SetContent(content)
+	contentBorder := lipgloss.RoundedBorder()
+	if p.isViewportFocus {
+		contentBorder = lipgloss.DoubleBorder()
+	}
+	cotentStyle := lipgloss.NewStyle().MarginTop(1).Padding(0, 1).Border(contentBorder).BorderForeground(styles.Flamingo)
 
 	return lipgloss.JoinVertical(
-		lipgloss.Center,
-		contentStyle.Render(p.content),
-		p.form.View(),
+		lipgloss.Top,
+		headerContent,
+		cotentStyle.Render(p.contentViewPort.View()),
+		form,
 	)
 }
 
+func (p *permissionDialogCmp) View() string {
+	return p.render()
+}
+
 func (p *permissionDialogCmp) GetSize() (int, int) {
 	return p.width, p.height
 }
@@ -99,13 +173,14 @@ func (p *permissionDialogCmp) GetSize() (int, int) {
 func (p *permissionDialogCmp) SetSize(width int, height int) {
 	p.width = width
 	p.height = height
+	p.form = p.form.WithWidth(width)
 }
 
 func (p *permissionDialogCmp) BindingKeys() []key.Binding {
 	return p.form.KeyBinds()
 }
 
-func newPermissionDialogCmp(permission permission.PermissionRequest, content string) PermissionDialog {
+func newPermissionDialogCmp(permission permission.PermissionRequest) PermissionDialog {
 	// Create a note field for displaying the content
 
 	// Create select field for the permission options
@@ -116,14 +191,13 @@ func newPermissionDialogCmp(permission permission.PermissionRequest, content str
 			huh.NewOption("Allow for this session", string(PermissionAllowForSession)),
 			huh.NewOption("Deny", string(PermissionDeny)),
 		).
-		Title("Permission Request")
+		Title("Select an action")
 
 	// Apply theme
 	theme := styles.HuhTheme()
 
 	// Setup form width and height
 	form := huh.NewForm(huh.NewGroup(selectOption)).
-		WithWidth(permissionWidth - 2).
 		WithShowHelp(false).
 		WithTheme(theme).
 		WithShowErrors(false)
@@ -132,25 +206,22 @@ func newPermissionDialogCmp(permission permission.PermissionRequest, content str
 	selectOption.Focus()
 
 	return &permissionDialogCmp{
-		permission: permission,
-		form:       form,
-		content:    content,
-		width:      permissionWidth,
-		height:     permissionHeight,
+		permission:   permission,
+		form:         form,
+		selectOption: selectOption,
 	}
 }
 
 // NewPermissionDialogCmd creates a new permission dialog command
-func NewPermissionDialogCmd(permission permission.PermissionRequest, content string) tea.Cmd {
-	permDialog := newPermissionDialogCmp(permission, content)
+func NewPermissionDialogCmd(permission permission.PermissionRequest) tea.Cmd {
+	permDialog := newPermissionDialogCmp(permission)
 
 	// Create the dialog layout
 	dialogPane := layout.NewSinglePane(
 		permDialog.(*permissionDialogCmp),
-		layout.WithSignlePaneSize(permissionWidth+2, permissionHeight+2),
 		layout.WithSinglePaneBordered(true),
 		layout.WithSinglePaneFocusable(true),
-		layout.WithSinglePaneActiveColor(styles.Blue),
+		layout.WithSinglePaneActiveColor(styles.Warning),
 		layout.WithSignlePaneBorderText(map[layout.BorderPosition]string{
 			layout.TopMiddleBorder: " Permission Required ",
 		}),
@@ -158,10 +229,24 @@ func NewPermissionDialogCmd(permission permission.PermissionRequest, content str
 
 	// Focus the dialog
 	dialogPane.Focus()
-
+	widthRatio := 0.7
+	heightRatio := 0.6
+	minWidth := 100
+	minHeight := 30
+
+	switch permission.ToolName {
+	case tools.BashToolName:
+		widthRatio = 0.5
+		heightRatio = 0.3
+		minWidth = 80
+		minHeight = 20
+	}
 	// Return the dialog command
 	return util.CmdHandler(core.DialogMsg{
-		Content: dialogPane,
+		Content:     dialogPane,
+		WidthRatio:  widthRatio,
+		HeightRatio: heightRatio,
+		MinWidth:    minWidth,
+		MinHeight:   minHeight,
 	})
 }
-
diff --git a/internal/tui/components/dialog/quit.go b/internal/tui/components/dialog/quit.go
index e73550ab5..60c1fc0d2 100644
--- a/internal/tui/components/dialog/quit.go
+++ b/internal/tui/components/dialog/quit.go
@@ -3,7 +3,6 @@ package dialog
 import (
 	"github.com/charmbracelet/bubbles/key"
 	tea "github.com/charmbracelet/bubbletea"
-	"github.com/charmbracelet/lipgloss"
 	"github.com/kujtimiihoxha/termai/internal/tui/components/core"
 	"github.com/kujtimiihoxha/termai/internal/tui/layout"
 	"github.com/kujtimiihoxha/termai/internal/tui/styles"
@@ -14,11 +13,6 @@ import (
 
 const question = "Are you sure you want to quit?"
 
-var (
-	width  = lipgloss.Width(question) + 6
-	height = 3
-)
-
 type QuitDialog interface {
 	tea.Model
 	layout.Sizeable
@@ -37,8 +31,6 @@ func (q *quitDialogCmp) Init() tea.Cmd {
 
 func (q *quitDialogCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	var cmds []tea.Cmd
-
-	// Process the form
 	form, cmd := q.form.Update(msg)
 	if f, ok := form.(*huh.Form); ok {
 		q.form = f
@@ -67,6 +59,7 @@ func (q *quitDialogCmp) GetSize() (int, int) {
 func (q *quitDialogCmp) SetSize(width int, height int) {
 	q.width = width
 	q.height = height
+	q.form = q.form.WithWidth(width).WithHeight(height)
 }
 
 func (q *quitDialogCmp) BindingKeys() []key.Binding {
@@ -84,28 +77,30 @@ func newQuitDialogCmp() QuitDialog {
 	theme.Focused.FocusedButton = theme.Focused.FocusedButton.Background(styles.Warning)
 	theme.Blurred.FocusedButton = theme.Blurred.FocusedButton.Background(styles.Warning)
 	form := huh.NewForm(huh.NewGroup(confirm)).
-		WithWidth(width).
-		WithHeight(height).
 		WithShowHelp(false).
+		WithWidth(0).
+		WithHeight(0).
 		WithTheme(theme).
 		WithShowErrors(false)
 	confirm.Focus()
 	return &quitDialogCmp{
-		form:  form,
-		width: width,
+		form: form,
 	}
 }
 
 func NewQuitDialogCmd() tea.Cmd {
 	content := layout.NewSinglePane(
 		newQuitDialogCmp().(*quitDialogCmp),
-		layout.WithSignlePaneSize(width+2, height+2),
 		layout.WithSinglePaneBordered(true),
 		layout.WithSinglePaneFocusable(true),
 		layout.WithSinglePaneActiveColor(styles.Warning),
 	)
 	content.Focus()
 	return util.CmdHandler(core.DialogMsg{
-		Content: content,
+		Content:     content,
+		WidthRatio:  0.2,
+		HeightRatio: 0.1,
+		MinWidth:    40,
+		MinHeight:   5,
 	})
 }
diff --git a/internal/tui/components/messages/message.go b/internal/tui/components/messages/message.go
deleted file mode 100644
index 619950850..000000000
--- a/internal/tui/components/messages/message.go
+++ /dev/null
@@ -1,108 +0,0 @@
-package messages
-
-import (
-	"fmt"
-
-	tea "github.com/charmbracelet/bubbletea"
-	"github.com/charmbracelet/lipgloss"
-	"github.com/cloudwego/eino/schema"
-	"github.com/kujtimiihoxha/termai/internal/message"
-	"github.com/kujtimiihoxha/termai/internal/tui/layout"
-	"github.com/kujtimiihoxha/termai/internal/tui/styles"
-)
-
-const (
-	maxHeight = 10
-)
-
-type MessagesCmp interface {
-	tea.Model
-	layout.Focusable
-	layout.Bordered
-	layout.Sizeable
-}
-
-type messageCmp struct {
-	message  message.Message
-	width    int
-	height   int
-	focused  bool
-	expanded bool
-}
-
-func (m *messageCmp) Init() tea.Cmd {
-	return nil
-}
-
-func (m *messageCmp) Update(tea.Msg) (tea.Model, tea.Cmd) {
-	return m, nil
-}
-
-func (m *messageCmp) View() string {
-	wrapper := layout.NewSinglePane(
-		m,
-		layout.WithSinglePaneBordered(true),
-		layout.WithSinglePaneFocusable(true),
-		layout.WithSinglePanePadding(1),
-		layout.WithSinglePaneActiveColor(m.borderColor()),
-	)
-	if m.focused {
-		wrapper.Focus()
-	}
-	wrapper.SetSize(m.width, m.height)
-	return wrapper.View()
-}
-
-func (m *messageCmp) Blur() tea.Cmd {
-	m.focused = false
-	return nil
-}
-
-func (m *messageCmp) borderColor() lipgloss.TerminalColor {
-	switch m.message.MessageData.Role {
-	case schema.Assistant:
-		return styles.Mauve
-	case schema.User:
-		return styles.Flamingo
-	}
-	return styles.Blue
-}
-
-func (m *messageCmp) BorderText() map[layout.BorderPosition]string {
-	role := ""
-	icon := ""
-	switch m.message.MessageData.Role {
-	case schema.Assistant:
-		role = "Assistant"
-		icon = styles.BotIcon
-	case schema.User:
-		role = "User"
-		icon = styles.UserIcon
-	}
-	return map[layout.BorderPosition]string{
-		layout.TopLeftBorder: fmt.Sprintf("%s %s ", role, icon),
-	}
-}
-
-func (m *messageCmp) Focus() tea.Cmd {
-	m.focused = true
-	return nil
-}
-
-func (m *messageCmp) IsFocused() bool {
-	return m.focused
-}
-
-func (m *messageCmp) GetSize() (int, int) {
-	return m.width, 0
-}
-
-func (m *messageCmp) SetSize(width int, height int) {
-	m.width = width
-}
-
-func NewMessageCmp(msg message.Message) MessagesCmp {
-	return &messageCmp{
-		message: msg,
-	}
-}
diff --git a/internal/tui/components/repl/editor.go b/internal/tui/components/repl/editor.go
index d0af8d2c5..cce966ca7 100644
--- a/internal/tui/components/repl/editor.go
+++ b/internal/tui/components/repl/editor.go
@@ -6,10 +6,11 @@ import (
 	"github.com/charmbracelet/bubbles/key"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
-	"github.com/cloudwego/eino/schema"
 	"github.com/kujtimiihoxha/termai/internal/app"
+	"github.com/kujtimiihoxha/termai/internal/llm/agent"
 	"github.com/kujtimiihoxha/termai/internal/tui/layout"
 	"github.com/kujtimiihoxha/termai/internal/tui/styles"
+	"github.com/kujtimiihoxha/termai/internal/tui/util"
 	"github.com/kujtimiihoxha/vimtea"
 )
 
@@ -112,7 +113,7 @@ func (m *editorCmp) BorderText() map[layout.BorderPosition]string {
 		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
 	}
 	return map[layout.BorderPosition]string{
-		layout.TopLeftBorder: title,
+		layout.BottomLeftBorder: title,
 	}
 }
 
@@ -137,9 +138,15 @@ func (m *editorCmp) SetSize(width int, height int) {
 
 func (m *editorCmp) Send() tea.Cmd {
 	return func() tea.Msg {
+		messages, _ := m.app.Messages.List(m.sessionID)
+		if hasUnfinishedMessages(messages) {
+			return util.InfoMsg("Assistant is still working on the previous message")
+		}
+		a, _ := agent.NewCoderAgent(m.app)
+
 		content := strings.Join(m.editor.GetBuffer().Lines(), "\n")
-		m.app.Messages.Create(m.sessionID, *schema.UserMessage(content))
-		m.app.LLM.SendRequest(m.sessionID, content)
+		go a.Generate(m.sessionID, content)
+
 		return m.editor.Reset()
 	}
 }
@@ -153,10 +160,11 @@ func (m *editorCmp) BindingKeys() []key.Binding {
 }
 
 func NewEditorCmp(app *app.App) EditorCmp {
+	editor := vimtea.NewEditor(
+		vimtea.WithFileName("message.md"),
+	)
 	return &editorCmp{
-		app: app,
-		editor: vimtea.NewEditor(
-			vimtea.WithFileName("message.md"),
-		),
+		app:    app,
+		editor: editor,
 	}
 }
diff --git a/internal/tui/components/repl/messages.go b/internal/tui/components/repl/messages.go
index 9b3c5bde8..7956867ce 100644
--- a/internal/tui/components/repl/messages.go
+++ b/internal/tui/components/repl/messages.go
@@ -1,8 +1,9 @@
 package repl
 
 import (
+	"encoding/json"
 	"fmt"
-	"slices"
+	"sort"
 	"strings"
 
 	"github.com/charmbracelet/bubbles/key"
@@ -10,8 +11,8 @@ import (
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/glamour"
 	"github.com/charmbracelet/lipgloss"
-	"github.com/cloudwego/eino/schema"
 	"github.com/kujtimiihoxha/termai/internal/app"
+	"github.com/kujtimiihoxha/termai/internal/llm/agent"
 	"github.com/kujtimiihoxha/termai/internal/message"
 	"github.com/kujtimiihoxha/termai/internal/pubsub"
 	"github.com/kujtimiihoxha/termai/internal/session"
@@ -28,30 +29,50 @@ type MessagesCmp interface {
 }
 
 type messagesCmp struct {
-	app        *app.App
-	messages   []message.Message
-	session    session.Session
-	viewport   viewport.Model
-	mdRenderer *glamour.TermRenderer
-	width      int
-	height     int
-	focused    bool
-	cachedView string
+	app            *app.App
+	messages       []message.Message
+	selectedMsgIdx int // Index of the selected message
+	session        session.Session
+	viewport       viewport.Model
+	mdRenderer     *glamour.TermRenderer
+	width          int
+	height         int
+	focused        bool
+	cachedView     string
 }
 
 func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	switch msg := msg.(type) {
 	case pubsub.Event[message.Message]:
 		if msg.Type == pubsub.CreatedEvent {
-			m.messages = append(m.messages, msg.Payload)
-			m.renderView()
-			m.viewport.GotoBottom()
+			if msg.Payload.SessionID == m.session.ID {
+				m.messages = append(m.messages, msg.Payload)
+				m.renderView()
+				m.viewport.GotoBottom()
+			}
+			for _, v := range m.messages {
+				for _, c := range v.ToolCalls {
+					if c.ID == msg.Payload.SessionID {
+						m.renderView()
+						m.viewport.GotoBottom()
+					}
+				}
+			}
+		} else if msg.Type == pubsub.UpdatedEvent && msg.Payload.SessionID == m.session.ID {
+			for i, v := range m.messages {
+				if v.ID == msg.Payload.ID {
+					m.messages[i] = msg.Payload
+					m.renderView()
+					if i == len(m.messages)-1 {
+						m.viewport.GotoBottom()
+					}
+					break
+				}
+			}
 		}
 	case pubsub.Event[session.Session]:
-		if msg.Type == pubsub.UpdatedEvent {
-			if m.session.ID == msg.Payload.ID {
-				m.session = msg.Payload
-			}
+		if msg.Type == pubsub.UpdatedEvent && m.session.ID == msg.Payload.ID {
+			m.session = msg.Payload
 		}
 	case SelectedSessionMsg:
 		m.session, _ = m.app.Sessions.Get(msg.SessionID)
@@ -67,26 +88,24 @@ func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	return m, nil
 }
 
-func borderColor(role schema.RoleType) lipgloss.TerminalColor {
+func borderColor(role message.MessageRole) lipgloss.TerminalColor {
 	switch role {
-	case schema.Assistant:
+	case message.Assistant:
 		return styles.Mauve
-	case schema.User:
+	case message.User:
 		return styles.Rosewater
-	case schema.Tool:
-		return styles.Peach
 	}
 	return styles.Blue
 }
 
-func borderText(msgRole schema.RoleType, currentMessage int) map[layout.BorderPosition]string {
+func borderText(msgRole message.MessageRole, currentMessage int) map[layout.BorderPosition]string {
 	role := ""
 	icon := ""
 	switch msgRole {
-	case schema.Assistant:
+	case message.Assistant:
 		role = "Assistant"
 		icon = styles.BotIcon
-	case schema.User:
+	case message.User:
 		role = "User"
 		icon = styles.UserIcon
 	}
@@ -106,81 +125,259 @@ func borderText(msgRole schema.RoleType, currentMessage int) map[layout.BorderPo
 	}
 }
 
+func hasUnfinishedMessages(messages []message.Message) bool {
+	if len(messages) == 0 {
+		return false
+	}
+	for _, msg := range messages {
+		if !msg.Finished {
+			return true
+		}
+	}
+	lastMessage := messages[len(messages)-1]
+	return lastMessage.Role != message.Assistant
+}
+
+func (m *messagesCmp) renderMessageWithToolCall(content string, tools []message.ToolCall, futureMessages []message.Message) string {
+	allParts := []string{content}
+
+	leftPaddingValue := 4
+	connectorStyle := lipgloss.NewStyle().
+		Foreground(styles.Peach).
+		Bold(true)
+
+	toolCallStyle := lipgloss.NewStyle().
+		Border(lipgloss.RoundedBorder()).
+		BorderForeground(styles.Peach).
+		Width(m.width-leftPaddingValue-5).
+		Padding(0, 1)
+
+	toolResultStyle := lipgloss.NewStyle().
+		Border(lipgloss.RoundedBorder()).
+		BorderForeground(styles.Green).
+		Width(m.width-leftPaddingValue-5).
+		Padding(0, 1)
+
+	leftPadding := lipgloss.NewStyle().Padding(0, 0, 0, leftPaddingValue)
+
+	runningStyle := lipgloss.NewStyle().
+		Foreground(styles.Peach).
+		Bold(true)
+
+	renderTool := func(toolCall message.ToolCall) string {
+		toolHeader := lipgloss.NewStyle().
+			Bold(true).
+			Foreground(styles.Blue).
+			Render(fmt.Sprintf("%s %s", styles.ToolIcon, toolCall.Name))
+
+		var paramLines []string
+		var args map[string]interface{}
+		var paramOrder []string
+
+		json.Unmarshal([]byte(toolCall.Input), &args)
+
+		for key := range args {
+			paramOrder = append(paramOrder, key)
+		}
+		sort.Strings(paramOrder)
+
+		for _, name := range paramOrder {
+			value := args[name]
+			paramName := lipgloss.NewStyle().
+				Foreground(styles.Peach).
+				Bold(true).
+				Render(name)
+
+			truncate := m.width - leftPaddingValue*2 - 10
+			if len(fmt.Sprintf("%v", value)) > truncate {
+				value = fmt.Sprintf("%v", value)[:truncate] + lipgloss.NewStyle().Foreground(styles.Blue).Render("... (truncated)")
+			}
+			paramValue := fmt.Sprintf("%v", value)
+			paramLines = append(paramLines, fmt.Sprintf("  %s: %s", paramName, paramValue))
+		}
+
+		paramBlock := lipgloss.JoinVertical(lipgloss.Left, paramLines...)
+
+		toolContent := lipgloss.JoinVertical(lipgloss.Left, toolHeader, paramBlock)
+		return toolCallStyle.Render(toolContent)
+	}
+
+	findToolResult := func(toolCallID string, messages []message.Message) *message.ToolResult {
+		for _, msg := range messages {
+			if msg.Role == message.Tool {
+				for _, result := range msg.ToolResults {
+					if result.ToolCallID == toolCallID {
+						return &result
+					}
+				}
+			}
+		}
+		return nil
+	}
+
+	renderToolResult := func(result message.ToolResult) string {
+		resultHeader := lipgloss.NewStyle().
+			Bold(true).
+			Foreground(styles.Green).
+			Render(fmt.Sprintf("%s Result", styles.CheckIcon))
+		if result.IsError {
+			resultHeader = lipgloss.NewStyle().
+				Bold(true).
+				Foreground(styles.Red).
+				Render(fmt.Sprintf("%s Error", styles.ErrorIcon))
+		}
+
+		truncate := 200
+		content := result.Content
+		if len(content) > truncate {
+			content = content[:truncate] + lipgloss.NewStyle().Foreground(styles.Blue).Render("... (truncated)")
+		}
+
+		resultContent := lipgloss.JoinVertical(lipgloss.Left, resultHeader, content)
+		return toolResultStyle.Render(resultContent)
+	}
+
+	connector := connectorStyle.Render("└─> Tool Calls:")
+	allParts = append(allParts, connector)
+
+	for _, toolCall := range tools {
+		toolOutput := renderTool(toolCall)
+		allParts = append(allParts, leftPadding.Render(toolOutput))
+
+		result := findToolResult(toolCall.ID, futureMessages)
+		if result != nil {
+
+			resultOutput := renderToolResult(*result)
+			allParts = append(allParts, leftPadding.Render(resultOutput))
+
+		} else if toolCall.Name == agent.AgentToolName {
+
+			runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
+			allParts = append(allParts, leftPadding.Render(runningIndicator))
+			taskSessionMessages, _ := m.app.Messages.List(toolCall.ID)
+			for _, msg := range taskSessionMessages {
+				if msg.Role == message.Assistant {
+					for _, toolCall := range msg.ToolCalls {
+						toolHeader := lipgloss.NewStyle().
+							Bold(true).
+							Foreground(styles.Blue).
+							Render(fmt.Sprintf("%s %s", styles.ToolIcon, toolCall.Name))
+
+						var paramLines []string
+						var args map[string]interface{}
+						var paramOrder []string
+
+						json.Unmarshal([]byte(toolCall.Input), &args)
+
+						for key := range args {
+							paramOrder = append(paramOrder, key)
+						}
+						sort.Strings(paramOrder)
+
+						for _, name := range paramOrder {
+							value := args[name]
+							paramName := lipgloss.NewStyle().
+								Foreground(styles.Peach).
+								Bold(true).
+								Render(name)
+
+							truncate := 50
+							if len(fmt.Sprintf("%v", value)) > truncate {
+								value = fmt.Sprintf("%v", value)[:truncate] + lipgloss.NewStyle().Foreground(styles.Blue).Render("... (truncated)")
+							}
+							paramValue := fmt.Sprintf("%v", value)
+							paramLines = append(paramLines, fmt.Sprintf("  %s: %s", paramName, paramValue))
+						}
+
+						paramBlock := lipgloss.JoinVertical(lipgloss.Left, paramLines...)
+						toolContent := lipgloss.JoinVertical(lipgloss.Left, toolHeader, paramBlock)
+						toolOutput := toolCallStyle.BorderForeground(styles.Teal).MaxWidth(m.width - leftPaddingValue*2 - 2).Render(toolContent)
+						allParts = append(allParts, lipgloss.NewStyle().Padding(0, 0, 0, leftPaddingValue*2).Render(toolOutput))
+					}
+				}
+			}
+
+		} else {
+			runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
+			allParts = append(allParts, "    "+runningIndicator)
+		}
+	}
+
+	for _, msg := range futureMessages {
+		if msg.Content != "" {
+			break
+		}
+
+		for _, toolCall := range msg.ToolCalls {
+			toolOutput := renderTool(toolCall)
+			allParts = append(allParts, "    "+strings.ReplaceAll(toolOutput, "\n", "\n    "))
+
+			result := findToolResult(toolCall.ID, futureMessages)
+			if result != nil {
+				resultOutput := renderToolResult(*result)
+				allParts = append(allParts, "    "+strings.ReplaceAll(resultOutput, "\n", "\n    "))
+			} else {
+				runningIndicator := runningStyle.Render(fmt.Sprintf("%s Running...", styles.SpinnerIcon))
+				allParts = append(allParts, "    "+runningIndicator)
+			}
+		}
+	}
+
+	return lipgloss.JoinVertical(lipgloss.Left, allParts...)
+}
+
 func (m *messagesCmp) renderView() {
 	stringMessages := make([]string, 0)
 	r, _ := glamour.NewTermRenderer(
 		glamour.WithStyles(styles.CatppuccinMarkdownStyle()),
-		glamour.WithWordWrap(m.width-10),
+		glamour.WithWordWrap(m.width-20),
 		glamour.WithEmoji(),
 	)
 	textStyle := lipgloss.NewStyle().Width(m.width - 4)
 	currentMessage := 1
-	for _, msg := range m.messages {
-		if msg.MessageData.Role == schema.Tool {
-			continue
-		}
-		content := msg.MessageData.Content
-		if content != "" {
-			content, _ = r.Render(msg.MessageData.Content)
-			stringMessages = append(stringMessages, layout.Borderize(
-				textStyle.Render(content),
-				layout.BorderOptions{
-					InactiveBorder: lipgloss.DoubleBorder(),
-					ActiveBorder:   lipgloss.DoubleBorder(),
-					ActiveColor:    borderColor(msg.MessageData.Role),
-					InactiveColor:  borderColor(msg.MessageData.Role),
-					EmbeddedText:   borderText(msg.MessageData.Role, currentMessage),
-				},
-			))
-			currentMessage++
-		}
-		for _, toolCall := range msg.MessageData.ToolCalls {
-			resultInx := slices.IndexFunc(m.messages, func(m message.Message) bool {
-				return m.MessageData.ToolCallID == toolCall.ID
-			})
-			content := fmt.Sprintf("**Arguments**\n```json\n%s\n```\n", toolCall.Function.Arguments)
-			if resultInx == -1 {
-				content += "Running..."
-			} else {
-				result := m.messages[resultInx].MessageData.Content
-				if result != "" {
-					lines := strings.Split(result, "\n")
-					if len(lines) > 15 {
-						result = strings.Join(lines[:15], "\n")
-					}
-					content += fmt.Sprintf("**Result**\n```\n%s\n```\n", result)
-					if len(lines) > 15 {
-						content += fmt.Sprintf("\n\n *...%d lines are truncated* ", len(lines)-15)
-					}
-				}
+	displayedMsgCount := 0 // Track the actual displayed messages count
+
+	prevMessageWasUser := false
+	for inx, msg := range m.messages {
+		content := msg.Content
+		if content != "" || prevMessageWasUser {
+			if msg.Thinking != "" && content == "" {
+				content = msg.Thinking
+			} else if content == "" {
+				content = "..."
 			}
 			content, _ = r.Render(content)
-			stringMessages = append(stringMessages, layout.Borderize(
+
+			isSelected := inx == m.selectedMsgIdx
+
+			border := lipgloss.DoubleBorder()
+			activeColor := borderColor(msg.Role)
+
+			if isSelected {
+				activeColor = styles.Primary // Use primary color for selected message
+			}
+
+			content = layout.Borderize(
 				textStyle.Render(content),
 				layout.BorderOptions{
-					InactiveBorder: lipgloss.DoubleBorder(),
-					ActiveBorder:   lipgloss.DoubleBorder(),
-					ActiveColor:    borderColor(schema.Tool),
-					InactiveColor:  borderColor(schema.Tool),
-					EmbeddedText: map[layout.BorderPosition]string{
-						layout.TopLeftBorder: lipgloss.NewStyle().
-							Padding(0, 1).
-							Bold(true).
-							Foreground(styles.Crust).
-							Background(borderColor(schema.Tool)).
-							Render(
-								fmt.Sprintf("Tool [%s] %s ", toolCall.Function.Name, styles.ToolIcon),
-							),
-						layout.TopRightBorder: lipgloss.NewStyle().
-							Padding(0, 1).
-							Bold(true).
-							Foreground(styles.Crust).
-							Background(borderColor(schema.Tool)).
-							Render(fmt.Sprintf("#%d ", currentMessage)),
-					},
+					InactiveBorder: border,
+					ActiveBorder:   border,
+					ActiveColor:    activeColor,
+					InactiveColor:  borderColor(msg.Role),
+					EmbeddedText:   borderText(msg.Role, currentMessage),
 				},
-			))
+			)
+			if len(msg.ToolCalls) > 0 {
+				content = m.renderMessageWithToolCall(content, msg.ToolCalls, m.messages[inx+1:])
+			}
+			stringMessages = append(stringMessages, content)
 			currentMessage++
+			displayedMsgCount++
+		}
+		if msg.Role == message.User && msg.Content != "" {
+			prevMessageWasUser = true
+		} else {
+			prevMessageWasUser = false
 		}
 	}
 	m.viewport.SetContent(lipgloss.JoinVertical(lipgloss.Top, stringMessages...))
@@ -191,7 +388,9 @@ func (m *messagesCmp) View() string {
 }
 
 func (m *messagesCmp) BindingKeys() []key.Binding {
-	return layout.KeyMapToSlice(m.viewport.KeyMap)
+	keys := layout.KeyMapToSlice(m.viewport.KeyMap)
+
+	return keys
 }
 
 func (m *messagesCmp) Blur() tea.Cmd {
@@ -208,10 +407,17 @@ func (m *messagesCmp) BorderText() map[layout.BorderPosition]string {
 	if m.focused {
 		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
 	}
-	return map[layout.BorderPosition]string{
+	borderTest := map[layout.BorderPosition]string{
 		layout.TopLeftBorder:     title,
 		layout.BottomRightBorder: formatTokensAndCost(m.session.CompletionTokens+m.session.PromptTokens, m.session.Cost),
 	}
+	if hasUnfinishedMessages(m.messages) {
+		borderTest[layout.BottomLeftBorder] = lipgloss.NewStyle().Foreground(styles.Peach).Render("Thinking...")
+	} else {
+		borderTest[layout.BottomLeftBorder] = lipgloss.NewStyle().Foreground(styles.Text).Render("Sleeping " + styles.SleepIcon + " ")
+	}
+
+	return borderTest
 }
 
 func (m *messagesCmp) Focus() tea.Cmd {
@@ -232,6 +438,7 @@ func (m *messagesCmp) SetSize(width int, height int) {
 	m.height = height
 	m.viewport.Width = width - 2   // padding
 	m.viewport.Height = height - 2 // padding
+	m.renderView()
 }
 
 func (m *messagesCmp) Init() tea.Cmd {
diff --git a/internal/tui/components/repl/sessions.go b/internal/tui/components/repl/sessions.go
index 0f208ced9..a1302baae 100644
--- a/internal/tui/components/repl/sessions.go
+++ b/internal/tui/components/repl/sessions.go
@@ -89,7 +89,23 @@ func (i *sessionsCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		}
 		return i, i.list.SetItems(items)
 	case pubsub.Event[session.Session]:
-		if msg.Type == pubsub.UpdatedEvent {
+		if msg.Type == pubsub.CreatedEvent && msg.Payload.ParentSessionID == "" {
+			// Check if the session is already in the list
+			items := i.list.Items()
+			for _, item := range items {
+				s := item.(listItem)
+				if s.id == msg.Payload.ID {
+					return i, nil
+				}
+			}
+			// insert the new session at the top of the list
+			items = append([]list.Item{listItem{
+				id:    msg.Payload.ID,
+				title: msg.Payload.Title,
+				desc:  formatTokensAndCost(msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost),
+			}}, items...)
+			return i, i.list.SetItems(items)
+		} else if msg.Type == pubsub.UpdatedEvent {
 			// update the session in the list
 			items := i.list.Items()
 			for idx, item := range items {
@@ -229,3 +245,4 @@ func NewSessionsCmp(app *app.App) SessionsCmp {
 		focused: false,
 	}
 }
+