@@ -1,16 +1,21 @@
11package chat
22
33import (
4+ "encoding/json"
45 "fmt"
5− "regexp"
6− "strconv"
6+ "math"
77 "strings"
88
9+ "github.com/charmbracelet/bubbles/spinner"
910 "github.com/charmbracelet/bubbles/viewport"
1011 tea "github.com/charmbracelet/bubbletea"
1112 "github.com/charmbracelet/glamour"
1213 "github.com/charmbracelet/lipgloss"
14+ "github.com/charmbracelet/x/ansi"
1315 "github.com/kujtimiihoxha/termai/internal/app"
16+ "github.com/kujtimiihoxha/termai/internal/llm/agent"
17+ "github.com/kujtimiihoxha/termai/internal/llm/models"
18+ "github.com/kujtimiihoxha/termai/internal/llm/tools"
1419 "github.com/kujtimiihoxha/termai/internal/message"
1520 "github.com/kujtimiihoxha/termai/internal/pubsub"
1621 "github.com/kujtimiihoxha/termai/internal/session"
@@ -18,10 +23,20 @@ import (
1823 "github.com/kujtimiihoxha/termai/internal/tui/util"
1924)
2025
26+type uiMessageType int
27+
28+const (
29+ userMessageType uiMessageType = iota
30+ assistantMessageType
31+ toolMessageType
32+)
33+
2134type uiMessage struct {
22− position int
23− height int
24− content string
35+ ID string
36+ messageType uiMessageType
37+ position int
38+ height int
39+ content string
2540}
2641
2742type messagesCmp struct {
@@ -32,141 +47,116 @@ type messagesCmp struct {
3247 session session.Session
3348 messages []message.Message
3449 uiMessages []uiMessage
35− currentIndex int
50+ currentMsgID string
3651 renderer *glamour.TermRenderer
3752 focusRenderer *glamour.TermRenderer
3853 cachedContent map[string]string
54+ agentWorking bool
55+ spinner spinner.Model
56+ needsRerender bool
57+ lastViewport string
3958}
4059
4160func (m *messagesCmp) Init() tea.Cmd {
42− return m.viewport.Init()
43−}
44−
45−var ansiEscape = regexp.MustCompile("\x1b\\[[0-9;]*m")
46−
47−func hexToBgSGR(hex string) (string, error) {
48− hex = strings.TrimPrefix(hex, "#")
49− if len(hex) != 6 {
50− return "", fmt.Errorf("invalid hex color: must be 6 hexadecimal digits")
51− }
52−
53− // Parse RGB components in one block
54− rgb := make([]uint64, 3)
55− for i := 0; i < 3; i++ {
56− val, err := strconv.ParseUint(hex[i*2:i*2+2], 16, 8)
57− if err != nil {
58− return "", err
59− }
60− rgb[i] = val
61− }
62−
63− return fmt.Sprintf("48;2;%d;%d;%d", rgb[0], rgb[1], rgb[2]), nil
64−}
65−
66−func forceReplaceBackgroundColors(input string, newBg string) string {
67− return ansiEscape.ReplaceAllStringFunc(input, func(seq string) string {
68− // Extract content between "\x1b[" and "m"
69− content := seq[2 : len(seq)-1]
70− tokens := strings.Split(content, ";")
71− var newTokens []string
72−
73− // Skip background color tokens
74− for i := 0; i < len(tokens); i++ {
75− if tokens[i] == "" {
76− continue
77− }
78−
79− val, err := strconv.Atoi(tokens[i])
80− if err != nil {
81− newTokens = append(newTokens, tokens[i])
82− continue
83− }
84−
85− // Skip background color tokens
86− if val == 48 {
87− // Skip "48;5;N" or "48;2;R;G;B" sequences
88− if i+1 < len(tokens) {
89− if nextVal, err := strconv.Atoi(tokens[i+1]); err == nil {
90− switch nextVal {
91− case 5:
92− i += 2 // Skip "5" and color index
93− case 2:
94− i += 4 // Skip "2" and RGB components
95− }
96− }
97− }
98− } else if (val < 40 || val > 47) && (val < 100 || val > 107) && val != 49 {
99− // Keep non-background tokens
100− newTokens = append(newTokens, tokens[i])
101− }
102− }
103−
104− // Add new background if provided
105− if newBg != "" {
106− newTokens = append(newTokens, strings.Split(newBg, ";")...)
107− }
108−
109− if len(newTokens) == 0 {
110− return ""
111− }
112−
113− return "\x1b[" + strings.Join(newTokens, ";") + "m"
114− })
61+ return tea.Batch(m.viewport.Init())
11562}
11663
11764func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
65+ var cmds []tea.Cmd
11866 switch msg := msg.(type) {
67+ case AgentWorkingMsg:
68+ m.agentWorking = bool(msg)
69+ if m.agentWorking {
70+ cmds = append(cmds, m.spinner.Tick)
71+ }
11972 case EditorFocusMsg:
12073 m.writingMode = bool(msg)
12174 case SessionSelectedMsg:
12275 if msg.ID != m.session.ID {
12376 cmd := m.SetSession(msg)
77+ m.needsRerender = true
12478 return m, cmd
12579 }
12680 return m, nil
81+ case SessionClearedMsg:
82+ m.session = session.Session{}
83+ m.messages = make([]message.Message, 0)
84+ m.currentMsgID = ""
85+ m.needsRerender = true
86+ return m, nil
87+
88+ case tea.KeyMsg:
89+ if m.writingMode {
90+ return m, nil
91+ }
12792 case pubsub.Event[message.Message]:
12893 if msg.Type == pubsub.CreatedEvent {
12994 if msg.Payload.SessionID == m.session.ID {
13095 // check if message exists
96+
97+ messageExists := false
13198 for _, v := range m.messages {
13299 if v.ID == msg.Payload.ID {
133− return m, nil
100+ messageExists = true
101+ break
134102 }
135103 }
136104
137− m.messages = append(m.messages, msg.Payload)
138− m.renderView()
139− m.viewport.GotoBottom()
105+ if !messageExists {
106+ m.messages = append(m.messages, msg.Payload)
107+ delete(m.cachedContent, m.currentMsgID)
108+ m.currentMsgID = msg.Payload.ID
109+ m.needsRerender = true
110+ }
140111 }
141112 for _, v := range m.messages {
142113 for _, c := range v.ToolCalls() {
143114 // the message is being added to the session of a tool called
144115 if c.ID == msg.Payload.SessionID {
145− m.renderView()
146− m.viewport.GotoBottom()
116+ m.needsRerender = true
147117 }
148118 }
149119 }
150120 } else if msg.Type == pubsub.UpdatedEvent && msg.Payload.SessionID == m.session.ID {
151121 for i, v := range m.messages {
152122 if v.ID == msg.Payload.ID {
123+ if !m.messages[i].IsFinished() && msg.Payload.IsFinished() && msg.Payload.FinishReason() == "end_turn" || msg.Payload.FinishReason() == "canceled" {
124+ cmds = append(cmds, util.CmdHandler(AgentWorkingMsg(false)))
125+ }
153126 m.messages[i] = msg.Payload
154127 delete(m.cachedContent, msg.Payload.ID)
155− m.renderView()
156− if i == len(m.messages)-1 {
157− m.viewport.GotoBottom()
158− }
128+ m.needsRerender = true
159129 break
160130 }
161131 }
162132 }
163133 }
134+ if m.agentWorking {
135+ u, cmd := m.spinner.Update(msg)
136+ m.spinner = u
137+ cmds = append(cmds, cmd)
138+ }
139+ oldPos := m.viewport.YPosition
164140 u, cmd := m.viewport.Update(msg)
165141 m.viewport = u
166− return m, cmd
142+ m.needsRerender = m.needsRerender || m.viewport.YPosition != oldPos
143+ cmds = append(cmds, cmd)
144+ if m.needsRerender {
145+ m.renderView()
146+ if len(m.messages) > 0 {
147+ if msg, ok := msg.(pubsub.Event[message.Message]); ok {
148+ if (msg.Type == pubsub.CreatedEvent) ||
149+ (msg.Type == pubsub.UpdatedEvent && msg.Payload.ID == m.messages[len(m.messages)-1].ID) {
150+ m.viewport.GotoBottom()
151+ }
152+ }
153+ }
154+ m.needsRerender = false
155+ }
156+ return m, tea.Batch(cmds...)
167157}
168158
169−func (m *messagesCmp) renderUserMessage(inx int, msg message.Message) string {
159+func (m *messagesCmp) renderSimpleMessage(msg message.Message, info ...string) string {
170160 if v, ok := m.cachedContent[msg.ID]; ok {
171161 return v
172162 }
@@ -178,7 +168,7 @@ func (m *messagesCmp) renderUserMessage(inx int, msg message.Message) string {
178168 BorderStyle(lipgloss.ThickBorder())
179169
180170 renderer := m.renderer
181− if inx == m.currentIndex {
171+ if msg.ID == m.currentMsgID {
182172 style = style.
183173 Foreground(styles.Forground).
184174 BorderForeground(styles.Blue).
@@ -186,33 +176,269 @@ func (m *messagesCmp) renderUserMessage(inx int, msg message.Message) string {
186176 renderer = m.focusRenderer
187177 }
188178 c, _ := renderer.Render(msg.Content().String())
189− col, _ := hexToBgSGR(styles.Background.Dark)
190− rendered := style.Render(forceReplaceBackgroundColors(c, col))
179+ parts := []string{
180+ styles.ForceReplaceBackgroundWithLipgloss(c, styles.Background),
181+ }
182+ // remove newline at the end
183+ parts[0] = strings.TrimSuffix(parts[0], "\n")
184+ if len(info) > 0 {
185+ parts = append(parts, info...)
186+ }
187+ rendered := style.Render(
188+ lipgloss.JoinVertical(
189+ lipgloss.Left,
190+ parts...,
191+ ),
192+ )
191193 m.cachedContent[msg.ID] = rendered
192194 return rendered
193195}
194196
197+func formatTimeDifference(unixTime1, unixTime2 int64) string {
198+ diffSeconds := float64(math.Abs(float64(unixTime2 - unixTime1)))
199+
200+ if diffSeconds < 60 {
201+ return fmt.Sprintf("%.1fs", diffSeconds)
202+ }
203+
204+ minutes := int(diffSeconds / 60)
205+ seconds := int(diffSeconds) % 60
206+ return fmt.Sprintf("%dm%ds", minutes, seconds)
207+}
208+
209+func (m *messagesCmp) renderToolCall(toolCall message.ToolCall, isNested bool) string {
210+ key := ""
211+ value := ""
212+ switch toolCall.Name {
213+ // TODO: add result data to the tools
214+ case agent.AgentToolName:
215+ key = "Task"
216+ var params agent.AgentParams
217+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
218+ value = params.Prompt
219+ // TODO: handle nested calls
220+ case tools.BashToolName:
221+ key = "Bash"
222+ var params tools.BashParams
223+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
224+ value = params.Command
225+ case tools.EditToolName:
226+ key = "Edit"
227+ var params tools.EditParams
228+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
229+ value = params.FilePath
230+ case tools.FetchToolName:
231+ key = "Fetch"
232+ var params tools.FetchParams
233+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
234+ value = params.URL
235+ case tools.GlobToolName:
236+ key = "Glob"
237+ var params tools.GlobParams
238+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
239+ if params.Path == "" {
240+ params.Path = "."
241+ }
242+ value = fmt.Sprintf("%s (%s)", params.Pattern, params.Path)
243+ case tools.GrepToolName:
244+ key = "Grep"
245+ var params tools.GrepParams
246+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
247+ if params.Path == "" {
248+ params.Path = "."
249+ }
250+ value = fmt.Sprintf("%s (%s)", params.Pattern, params.Path)
251+ case tools.LSToolName:
252+ key = "Ls"
253+ var params tools.LSParams
254+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
255+ if params.Path == "" {
256+ params.Path = "."
257+ }
258+ value = params.Path
259+ case tools.SourcegraphToolName:
260+ key = "Sourcegraph"
261+ var params tools.SourcegraphParams
262+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
263+ value = params.Query
264+ case tools.ViewToolName:
265+ key = "View"
266+ var params tools.ViewParams
267+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
268+ value = params.FilePath
269+ case tools.WriteToolName:
270+ key = "Write"
271+ var params tools.WriteParams
272+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
273+ value = params.FilePath
274+ default:
275+ key = toolCall.Name
276+ var params map[string]any
277+ json.Unmarshal([]byte(toolCall.Input), ¶ms)
278+ jsonData, _ := json.Marshal(params)
279+ value = string(jsonData)
280+ }
281+
282+ style := styles.BaseStyle.
283+ Width(m.width).
284+ BorderLeft(true).
285+ BorderStyle(lipgloss.ThickBorder()).
286+ PaddingLeft(1).
287+ BorderForeground(styles.Yellow)
288+
289+ keyStyle := styles.BaseStyle.
290+ Foreground(styles.ForgroundDim)
291+ valyeStyle := styles.BaseStyle.
292+ Foreground(styles.Forground)
293+
294+ if isNested {
295+ valyeStyle = valyeStyle.Foreground(styles.ForgroundMid)
296+ }
297+ keyValye := keyStyle.Render(
298+ fmt.Sprintf("%s: ", key),
299+ )
300+ if !isNested {
301+ value = valyeStyle.
302+ Width(m.width - lipgloss.Width(keyValye) - 2).
303+ Render(
304+ ansi.Truncate(
305+ value,
306+ m.width-lipgloss.Width(keyValye)-2,
307+ "...",
308+ ),
309+ )
310+ } else {
311+ keyValye = keyStyle.Render(
312+ fmt.Sprintf(" └ %s: ", key),
313+ )
314+ value = valyeStyle.
315+ Width(m.width - lipgloss.Width(keyValye) - 2).
316+ Render(
317+ ansi.Truncate(
318+ value,
319+ m.width-lipgloss.Width(keyValye)-2,
320+ "...",
321+ ),
322+ )
323+ }
324+
325+ innerToolCalls := make([]string, 0)
326+ if toolCall.Name == agent.AgentToolName {
327+ messages, _ := m.app.Messages.List(toolCall.ID)
328+ toolCalls := make([]message.ToolCall, 0)
329+ for _, v := range messages {
330+ toolCalls = append(toolCalls, v.ToolCalls()...)
331+ }
332+ for _, v := range toolCalls {
333+ call := m.renderToolCall(v, true)
334+ innerToolCalls = append(innerToolCalls, call)
335+ }
336+ }
337+
338+ if isNested {
339+ return lipgloss.JoinHorizontal(
340+ lipgloss.Left,
341+ keyValye,
342+ value,
343+ )
344+ }
345+ callContent := lipgloss.JoinHorizontal(
346+ lipgloss.Left,
347+ keyValye,
348+ value,
349+ )
350+ callContent = strings.ReplaceAll(callContent, "\n", "")
351+ if len(innerToolCalls) > 0 {
352+ callContent = lipgloss.JoinVertical(
353+ lipgloss.Left,
354+ callContent,
355+ lipgloss.JoinVertical(
356+ lipgloss.Left,
357+ innerToolCalls...,
358+ ),
359+ )
360+ }
361+ return style.Render(callContent)
362+}
363+
364+func (m *messagesCmp) renderAssistantMessage(msg message.Message) []uiMessage {
365+ // find the user message that is before this assistant message
366+ var userMsg message.Message
367+ for i := len(m.messages) - 1; i >= 0; i-- {
368+ if m.messages[i].Role == message.User {
369+ userMsg = m.messages[i]
370+ break
371+ }
372+ }
373+ messages := make([]uiMessage, 0)
374+ if msg.Content().String() != "" {
375+ info := make([]string, 0)
376+ if msg.IsFinished() && msg.FinishReason() == "end_turn" {
377+ finish := msg.FinishPart()
378+ took := formatTimeDifference(userMsg.CreatedAt, finish.Time)
379+
380+ info = append(info, styles.BaseStyle.Width(m.width-1).Foreground(styles.ForgroundDim).Render(
381+ fmt.Sprintf(" %s (%s)", models.SupportedModels[msg.Model].Name, took),
382+ ))
383+ }
384+ content := m.renderSimpleMessage(msg, info...)
385+ messages = append(messages, uiMessage{
386+ messageType: assistantMessageType,
387+ position: 0, // gets updated in renderView
388+ height: lipgloss.Height(content),
389+ content: content,
390+ })
391+ }
392+ for _, v := range msg.ToolCalls() {
393+ content := m.renderToolCall(v, false)
394+ messages = append(messages,
395+ uiMessage{
396+ messageType: toolMessageType,
397+ position: 0, // gets updated in renderView
398+ height: lipgloss.Height(content),
399+ content: content,
400+ },
401+ )
402+ }
403+
404+ return messages
405+}
406+
195407func (m *messagesCmp) renderView() {
196408 m.uiMessages = make([]uiMessage, 0)
197409 pos := 0
198410
199411 for _, v := range m.messages {
200− content := ""
201412 switch v.Role {
202413 case message.User:
203− content = m.renderUserMessage(pos, v)
414+ content := m.renderSimpleMessage(v)
415+ m.uiMessages = append(m.uiMessages, uiMessage{
416+ messageType: userMessageType,
417+ position: pos,
418+ height: lipgloss.Height(content),
419+ content: content,
420+ })
421+ pos += lipgloss.Height(content) + 1 // + 1 for spacing
422+ case message.Assistant:
423+ assistantMessages := m.renderAssistantMessage(v)
424+ for _, msg := range assistantMessages {
425+ msg.position = pos
426+ m.uiMessages = append(m.uiMessages, msg)
427+ pos += msg.height + 1 // + 1 for spacing
428+ }
429+
204430 }
205− m.uiMessages = append(m.uiMessages, uiMessage{
206− position: pos,
207− height: lipgloss.Height(content),
208− content: content,
209− })
210− pos += lipgloss.Height(content) + 1 // + 1 for spacing
211431 }
212432
213433 messages := make([]string, 0)
214434 for _, v := range m.uiMessages {
215− messages = append(messages, v.content)
435+ messages = append(messages, v.content,
436+ styles.BaseStyle.
437+ Width(m.width).
438+ Render(
439+ "",
440+ ),
441+ )
216442 }
217443 m.viewport.SetContent(
218444 styles.BaseStyle.
@@ -246,7 +472,6 @@ func (m *messagesCmp) View() string {
246472 )
247473 }
248474
249− m.renderView()
250475 return styles.BaseStyle.
251476 Width(m.width).
252477 Render(
@@ -260,15 +485,21 @@ func (m *messagesCmp) View() string {
260485
261486func (m *messagesCmp) help() string {
262487 text := ""
488+
489+ if m.agentWorking {
490+ text += styles.BaseStyle.Foreground(styles.PrimaryColor).Bold(true).Render(
491+ fmt.Sprintf("%s %s ", m.spinner.View(), "Generating..."),
492+ )
493+ }
263494 if m.writingMode {
264− text = lipgloss.JoinHorizontal(
495+ text += lipgloss.JoinHorizontal(
265496 lipgloss.Left,
266497 styles.BaseStyle.Foreground(styles.ForgroundDim).Bold(true).Render("press "),
267498 styles.BaseStyle.Foreground(styles.Forground).Bold(true).Render("esc"),
268499 styles.BaseStyle.Foreground(styles.ForgroundDim).Bold(true).Render(" to exit writing mode"),
269500 )
270501 } else {
271− text = lipgloss.JoinHorizontal(
502+ text += lipgloss.JoinHorizontal(
272503 lipgloss.Left,
273504 styles.BaseStyle.Foreground(styles.ForgroundDim).Bold(true).Render("press "),
274505 styles.BaseStyle.Foreground(styles.Forground).Bold(true).Render("i"),
@@ -306,7 +537,15 @@ func (m *messagesCmp) SetSize(width, height int) {
306537 glamour.WithWordWrap(width-1),
307538 )
308539 m.focusRenderer = focusRenderer
540+ // clear the cached content
541+ for k := range m.cachedContent {
542+ delete(m.cachedContent, k)
543+ }
309544 m.renderer = renderer
545+ if len(m.messages) > 0 {
546+ m.renderView()
547+ m.viewport.GotoBottom()
548+ }
310549}
311550
312551func (m *messagesCmp) GetSize() (int, int) {
@@ -320,7 +559,8 @@ func (m *messagesCmp) SetSession(session session.Session) tea.Cmd {
320559 return util.ReportError(err)
321560 }
322561 m.messages = messages
323− m.messages = append(m.messages, m.messages[0])
562+ m.currentMsgID = m.messages[len(m.messages)-1].ID
563+ m.needsRerender = true
324564 return nil
325565}
326566
@@ -333,6 +573,9 @@ func NewMessagesCmp(app *app.App) tea.Model {
333573 glamour.WithStyles(styles.MarkdownTheme(false)),
334574 glamour.WithWordWrap(80),
335575 )
576+
577+ s := spinner.New()
578+ s.Spinner = spinner.Pulse
336579 return &messagesCmp{
337580 app: app,
338581 writingMode: true,
@@ -340,5 +583,6 @@ func NewMessagesCmp(app *app.App) tea.Model {
340583 viewport: viewport.New(0, 0),
341584 focusRenderer: focusRenderer,
342585 renderer: renderer,
586+ spinner: s,
343587 }
344588}