diff --git a/internal/ui/chat/messages.go b/internal/ui/chat/messages.go
index 53c0e5f45bb169b6ddbc92bc92262b4e47e3a509..3e12efefa4db1c9d78cbfe5a44374f4a918290f4 100644
--- a/internal/ui/chat/messages.go
+++ b/internal/ui/chat/messages.go
@@ -163,6 +163,10 @@ func GetMessageItems(sty *styles.Styles, msg *message.Message, toolResults map[s
 	return []MessageItem{}
 }
 
+// shouldRenderAssistantMessage determines if an assistant message should be rendered
+//
+// In some cases the assistant message only has tools so we do not want to render an
+// empty message.
 func shouldRenderAssistantMessage(msg *message.Message) bool {
 	content := strings.TrimSpace(msg.Content().Text)
 	thinking := strings.TrimSpace(msg.ReasoningContent().Thinking)
diff --git a/internal/ui/model/ui.go b/internal/ui/model/ui.go
index 5c738f6835a4666bd87315fb9ce43e498ba29d05..2caa9ac9b910a4439d1d69a62c595dce94239695 100644
--- a/internal/ui/model/ui.go
+++ b/internal/ui/model/ui.go
@@ -384,6 +384,7 @@ func (m *UI) setSessionMessages(msgs []message.Message) tea.Cmd {
 	return tea.Batch(cmds...)
 }
 
+// appendSessionMessage appends a new message to the current session in the chat
 func (m *UI) appendSessionMessage(msg message.Message) tea.Cmd {
 	items := chat.GetMessageItems(m.com.Styles, &msg, nil)
 	var cmds []tea.Cmd
@@ -399,6 +400,8 @@ func (m *UI) appendSessionMessage(msg message.Message) tea.Cmd {
 	return tea.Batch(cmds...)
 }
 
+// updateSessionMessage updates an existing message in the current session in the chat
+// INFO: currently only updates the assistant when I add tools this will get a bit more complex
 func (m *UI) updateSessionMessage(msg message.Message) tea.Cmd {
 	existingItem := m.chat.GetMessageItem(msg.ID)
 	switch msg.Role {