fix logs and add cancellation

6 files changed, 50 insertions, 21 deletions

diff --git a/internal/tui/components/core/status.go b/internal/tui/components/core/status.go
index d67bdb7a6..93ba34507 100644
--- a/internal/tui/components/core/status.go
+++ b/internal/tui/components/core/status.go
@@ -13,7 +13,7 @@ import (
 )
 
 type statusCmp struct {
-	info       *util.InfoMsg
+	info       util.InfoMsg
 	width      int
 	messageTTL time.Duration
 }
@@ -35,14 +35,14 @@ func (m statusCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.width = msg.Width
 		return m, nil
 	case util.InfoMsg:
-		m.info = &msg
+		m.info = msg
 		ttl := msg.TTL
 		if ttl == 0 {
 			ttl = m.messageTTL
 		}
 		return m, m.clearMessageCmd(ttl)
 	case util.ClearStatusMsg:
-		m.info = nil
+		m.info = util.InfoMsg{}
 	}
 	return m, nil
 }
@@ -54,7 +54,7 @@ var (
 
 func (m statusCmp) View() string {
 	status := styles.Padded.Background(styles.Grey).Foreground(styles.Text).Render("? help")
-	if m.info != nil {
+	if m.info.Msg != "" {
 		infoStyle := styles.Padded.
 			Foreground(styles.Base).
 			Width(m.availableFooterMsgWidth())
diff --git a/internal/tui/components/logs/details.go b/internal/tui/components/logs/details.go
index e02dfd9cc..dbace5508 100644
--- a/internal/tui/components/logs/details.go
+++ b/internal/tui/components/logs/details.go
@@ -30,7 +30,7 @@ type detailCmp struct {
 }
 
 func (i *detailCmp) Init() tea.Cmd {
-	messages := logging.Get().List()
+	messages := logging.List()
 	if len(messages) == 0 {
 		return nil
 	}
diff --git a/internal/tui/components/logs/table.go b/internal/tui/components/logs/table.go
index 0e90d2d2a..9500059b1 100644
--- a/internal/tui/components/logs/table.go
+++ b/internal/tui/components/logs/table.go
@@ -22,8 +22,6 @@ type TableComponent interface {
 	layout.Bordered
 }
 
-var logger = logging.Get()
-
 type tableCmp struct {
 	table table.Model
 }
@@ -57,7 +55,7 @@ func (i *tableCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		if selectedRow != nil {
 			if prevSelectedRow == nil || selectedRow[0] == prevSelectedRow[0] {
 				var log logging.LogMessage
-				for _, row := range logging.Get().List() {
+				for _, row := range logging.List() {
 					if row.ID == selectedRow[0] {
 						log = row
 						break
@@ -112,7 +110,7 @@ func (i *tableCmp) BindingKeys() []key.Binding {
 func (i *tableCmp) setRows() {
 	rows := []table.Row{}
 
-	logs := logger.List()
+	logs := logging.List()
 	slices.SortFunc(logs, func(a, b logging.LogMessage) int {
 		if a.Time.Before(b.Time) {
 			return 1
diff --git a/internal/tui/components/repl/editor.go b/internal/tui/components/repl/editor.go
index 0228996b8..37ac275e3 100644
--- a/internal/tui/components/repl/editor.go
+++ b/internal/tui/components/repl/editor.go
@@ -12,6 +12,7 @@ import (
 	"github.com/kujtimiihoxha/termai/internal/tui/styles"
 	"github.com/kujtimiihoxha/termai/internal/tui/util"
 	"github.com/kujtimiihoxha/vimtea"
+	"golang.org/x/net/context"
 )
 
 type EditorCmp interface {
@@ -23,18 +24,20 @@ type EditorCmp interface {
 }
 
 type editorCmp struct {
-	app        *app.App
-	editor     vimtea.Editor
-	editorMode vimtea.EditorMode
-	sessionID  string
-	focused    bool
-	width      int
-	height     int
+	app           *app.App
+	editor        vimtea.Editor
+	editorMode    vimtea.EditorMode
+	sessionID     string
+	focused       bool
+	width         int
+	height        int
+	cancelMessage context.CancelFunc
 }
 
 type editorKeyMap struct {
 	SendMessage    key.Binding
 	SendMessageI   key.Binding
+	CancelMessage  key.Binding
 	InsertMode     key.Binding
 	NormaMode      key.Binding
 	VisualMode     key.Binding
@@ -50,6 +53,10 @@ var editorKeyMapValue = editorKeyMap{
 		key.WithKeys("ctrl+s"),
 		key.WithHelp("ctrl+s", "send message insert mode"),
 	),
+	CancelMessage: key.NewBinding(
+		key.WithKeys("ctrl+x"),
+		key.WithHelp("ctrl+x", "cancel current message"),
+	),
 	InsertMode: key.NewBinding(
 		key.WithKeys("i"),
 		key.WithHelp("i", "insert mode"),
@@ -93,6 +100,8 @@ func (m *editorCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 				if m.editorMode == vimtea.ModeInsert {
 					return m, m.Send()
 				}
+			case key.Matches(msg, editorKeyMapValue.CancelMessage):
+				return m, m.Cancel()
 			}
 		}
 		u, cmd := m.editor.Update(msg)
@@ -136,6 +145,16 @@ func (m *editorCmp) SetSize(width int, height int) {
 	m.editor.SetSize(width, height)
 }
 
+func (m *editorCmp) Cancel() tea.Cmd {
+	if m.cancelMessage == nil {
+		return util.ReportWarn("No message to cancel")
+	}
+
+	m.cancelMessage()
+	m.cancelMessage = nil
+	return util.ReportWarn("Message cancelled")
+}
+
 func (m *editorCmp) Send() tea.Cmd {
 	return func() tea.Msg {
 		messages, err := m.app.Messages.List(m.sessionID)
@@ -151,7 +170,13 @@ func (m *editorCmp) Send() tea.Cmd {
 		}
 
 		content := strings.Join(m.editor.GetBuffer().Lines(), "\n")
-		go a.Generate(m.sessionID, content)
+		ctx, cancel := context.WithCancel(m.app.Context)
+		m.cancelMessage = cancel
+		go func() {
+			defer cancel()
+			a.Generate(ctx, m.sessionID, content)
+			m.cancelMessage = nil
+		}()
 
 		return m.editor.Reset()
 	}
diff --git a/internal/tui/components/repl/messages.go b/internal/tui/components/repl/messages.go
index 481783a2e..57a55c579 100644
--- a/internal/tui/components/repl/messages.go
+++ b/internal/tui/components/repl/messages.go
@@ -309,7 +309,7 @@ func (m *messagesCmp) renderMessageWithToolCall(content string, tools []message.
 	}
 
 	for _, msg := range futureMessages {
-		if msg.Content().String() != "" {
+		if msg.Content().String() != "" || msg.FinishReason() == "canceled" {
 			break
 		}
 
@@ -345,13 +345,18 @@ func (m *messagesCmp) renderView() {
 	prevMessageWasUser := false
 	for inx, msg := range m.messages {
 		content := msg.Content().String()
-		if content != "" || prevMessageWasUser {
+		if content != "" || prevMessageWasUser || msg.FinishReason() == "canceled" {
 			if msg.ReasoningContent().String() != "" && content == "" {
 				content = msg.ReasoningContent().String()
 			} else if content == "" {
 				content = "..."
 			}
-			content, _ = r.Render(content)
+			if msg.FinishReason() == "canceled" {
+				content, _ = r.Render(content)
+				content += lipgloss.NewStyle().Padding(1, 0, 0, 1).Foreground(styles.Error).Render(styles.ErrorIcon + " Canceled")
+			} else {
+				content, _ = r.Render(content)
+			}
 
 			isSelected := inx == m.selectedMsgIdx
 
diff --git a/internal/tui/tui.go b/internal/tui/tui.go
index 233b189dd..9e863d2ac 100644
--- a/internal/tui/tui.go
+++ b/internal/tui/tui.go
@@ -101,7 +101,8 @@ func (a appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 	// Status
 	case util.InfoMsg:
 		a.status, cmd = a.status.Update(msg)
-		return a, cmd
+		cmds = append(cmds, cmd)
+		return a, tea.Batch(cmds...)
 	case pubsub.Event[logging.LogMessage]:
 		if msg.Payload.Persist {
 			switch msg.Payload.Level {