Enha: save/load message token stats

4 files changed, 30 insertions, 20 deletions

diff --git a/bot.go b/bot.go
index 1367c02..527017f 100644
--- a/bot.go
+++ b/bot.go
@@ -784,7 +784,7 @@ func showSpinner() {
 			case toolRunningMode:
 				textArea.SetTitle(spinners[spin] + " tool")
 			case botRespMode:
-				textArea.SetTitle(spinners[spin] + " " + botPersona)
+				textArea.SetTitle(spinners[spin] + " " + botPersona + " (F6 to interrupt)")
 			default:
 				textArea.SetTitle(spinners[spin] + " input")
 			}
diff --git a/helpfuncs.go b/helpfuncs.go
index 356dd37..f47fdd1 100644
--- a/helpfuncs.go
+++ b/helpfuncs.go
@@ -359,7 +359,7 @@ func makeStatusLine() string {
 	}
 	// Get model color based on load status for local llama.cpp models
 	modelColor := getModelColor()
-	statusLine := fmt.Sprintf(statusLineTempl, boolColors[botRespMode], activeChatName,
+	statusLine := fmt.Sprintf(statusLineTempl, activeChatName,
 		boolColors[cfg.ToolUse], modelColor, chatBody.Model, boolColors[cfg.SkipLLMResp],
 		cfg.CurrentAPI, persona, botPersona)
 	if cfg.STT_ENABLED {
diff --git a/main.go b/main.go
index cab86a2..8347b77 100644
--- a/main.go
+++ b/main.go
@@ -16,7 +16,7 @@ var (
 	shellHistory      []string
 	shellHistoryPos   int = -1
 	thinkingCollapsed     = false
-	statusLineTempl       = "help (F12) | [%s:-:b]llm writes[-:-:-] (F6 to interrupt) | chat: [orange:-:b]%s[-:-:-] (F1) | [%s:-:b]tool use[-:-:-] (ctrl+k) | model: [%s:-:b]%s[-:-:-] (ctrl+l) | [%s:-:b]skip LLM resp[-:-:-] (F10)\nAPI: [orange:-:b]%s[-:-:-] (ctrl+v) | writing as: [orange:-:b]%s[-:-:-] (ctrl+q) | bot will write as [orange:-:b]%s[-:-:-] (ctrl+x)"
+	statusLineTempl       = "help (F12) | chat: [orange:-:b]%s[-:-:-] (F1) | [%s:-:b]tool use[-:-:-] (ctrl+k) | model: [%s:-:b]%s[-:-:-] (ctrl+l) | [%s:-:b]skip LLM resp[-:-:-] (F10)\nAPI: [orange:-:b]%s[-:-:-] (ctrl+v) | writing as: [orange:-:b]%s[-:-:-] (ctrl+q) | bot will write as [orange:-:b]%s[-:-:-] (ctrl+x)"
 	focusSwitcher         = map[tview.Primitive]tview.Primitive{}
 )
 
diff --git a/models/models.go b/models/models.go
index 7ee4c12..5ea85ba 100644
--- a/models/models.go
+++ b/models/models.go
@@ -115,33 +115,39 @@ type RoleMsg struct {
 }
 
 // MarshalJSON implements custom JSON marshaling for RoleMsg
-func (m *RoleMsg) MarshalJSON() ([]byte, error) {
+//
+//nolint:gocritic
+func (m RoleMsg) MarshalJSON() ([]byte, error) {
 	if m.hasContentParts {
 		// Use structured content format
 		aux := struct {
-			Role       string   `json:"role"`
-			Content    []any    `json:"content"`
-			ToolCallID string   `json:"tool_call_id,omitempty"`
-			KnownTo    []string `json:"known_to,omitempty"`
+			Role       string         `json:"role"`
+			Content    []any          `json:"content"`
+			ToolCallID string         `json:"tool_call_id,omitempty"`
+			KnownTo    []string       `json:"known_to,omitempty"`
+			Stats      *ResponseStats `json:"stats,omitempty"`
 		}{
 			Role:       m.Role,
 			Content:    m.ContentParts,
 			ToolCallID: m.ToolCallID,
 			KnownTo:    m.KnownTo,
+			Stats:      m.Stats,
 		}
 		return json.Marshal(aux)
 	} else {
 		// Use simple content format
 		aux := struct {
-			Role       string   `json:"role"`
-			Content    string   `json:"content"`
-			ToolCallID string   `json:"tool_call_id,omitempty"`
-			KnownTo    []string `json:"known_to,omitempty"`
+			Role       string         `json:"role"`
+			Content    string         `json:"content"`
+			ToolCallID string         `json:"tool_call_id,omitempty"`
+			KnownTo    []string       `json:"known_to,omitempty"`
+			Stats      *ResponseStats `json:"stats,omitempty"`
 		}{
 			Role:       m.Role,
 			Content:    m.Content,
 			ToolCallID: m.ToolCallID,
 			KnownTo:    m.KnownTo,
+			Stats:      m.Stats,
 		}
 		return json.Marshal(aux)
 	}
@@ -151,26 +157,29 @@ func (m *RoleMsg) MarshalJSON() ([]byte, error) {
 func (m *RoleMsg) UnmarshalJSON(data []byte) error {
 	// First, try to unmarshal as structured content format
 	var structured struct {
-		Role       string   `json:"role"`
-		Content    []any    `json:"content"`
-		ToolCallID string   `json:"tool_call_id,omitempty"`
-		KnownTo    []string `json:"known_to,omitempty"`
+		Role       string         `json:"role"`
+		Content    []any          `json:"content"`
+		ToolCallID string         `json:"tool_call_id,omitempty"`
+		KnownTo    []string       `json:"known_to,omitempty"`
+		Stats      *ResponseStats `json:"stats,omitempty"`
 	}
 	if err := json.Unmarshal(data, &structured); err == nil && len(structured.Content) > 0 {
 		m.Role = structured.Role
 		m.ContentParts = structured.Content
 		m.ToolCallID = structured.ToolCallID
 		m.KnownTo = structured.KnownTo
+		m.Stats = structured.Stats
 		m.hasContentParts = true
 		return nil
 	}
 
 	// Otherwise, unmarshal as simple content format
 	var simple struct {
-		Role       string   `json:"role"`
-		Content    string   `json:"content"`
-		ToolCallID string   `json:"tool_call_id,omitempty"`
-		KnownTo    []string `json:"known_to,omitempty"`
+		Role       string         `json:"role"`
+		Content    string         `json:"content"`
+		ToolCallID string         `json:"tool_call_id,omitempty"`
+		KnownTo    []string       `json:"known_to,omitempty"`
+		Stats      *ResponseStats `json:"stats,omitempty"`
 	}
 	if err := json.Unmarshal(data, &simple); err != nil {
 		return err
@@ -179,6 +188,7 @@ func (m *RoleMsg) UnmarshalJSON(data []byte) error {
 	m.Content = simple.Content
 	m.ToolCallID = simple.ToolCallID
 	m.KnownTo = simple.KnownTo
+	m.Stats = simple.Stats
 	m.hasContentParts = false
 	return nil
 }