From 1d9b33327db51d117e747550c19fdef34f84fd6f Mon Sep 17 00:00:00 2001
From: joshellis625 <4002542+joshellis625@users.noreply.github.com>
Date: Mon, 9 Feb 2026 07:47:21 -0500
Subject: [PATCH] fix: clamp chat totals when removing message usage

---
 JChat/Chat.swift                          | 14 ++++++++++++++
 JChat/Documentation/CHANGELOG_INTERNAL.md |  1 +
 JChat/Views/ChatViewModel.swift           | 16 ++++------------
 JChatTests/JChatTests.swift               | 10 ++++++++++
 4 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/JChat/Chat.swift b/JChat/Chat.swift
index dc145ba..7c3c9a5 100644
--- a/JChat/Chat.swift
+++ b/JChat/Chat.swift
@@ -139,6 +139,20 @@ final class Chat {
         if verbosityOverride != nil { count += 1 }
         return count
     }
+
+    // MARK: - Usage Totals
+
+    func addUsage(promptTokens: Int, completionTokens: Int, cost: Double) {
+        totalPromptTokens += promptTokens
+        totalCompletionTokens += completionTokens
+        totalCost += cost
+    }
+
+    func removeUsage(promptTokens: Int, completionTokens: Int, cost: Double) {
+        totalPromptTokens = max(0, totalPromptTokens - promptTokens)
+        totalCompletionTokens = max(0, totalCompletionTokens - completionTokens)
+        totalCost = max(0, totalCost - cost)
+    }
 }
 
 @Model
diff --git a/JChat/Documentation/CHANGELOG_INTERNAL.md b/JChat/Documentation/CHANGELOG_INTERNAL.md
index 8959cbb..b452ec6 100644
--- a/JChat/Documentation/CHANGELOG_INTERNAL.md
+++ b/JChat/Documentation/CHANGELOG_INTERNAL.md
@@ -10,3 +10,4 @@
 - Added targeted unit tests for core chat model behavior.
 - Improved Markdown code block readability with stronger contrast, consistent header labels, and always-visible copy action.
 - Replaced abrupt setup blocking with a first-run readiness checklist for API key and default model in `ContentView`.
+- Hardened chat usage accounting by clamping token/cost totals at zero during delete/regenerate paths.
diff --git a/JChat/Views/ChatViewModel.swift b/JChat/Views/ChatViewModel.swift
index 6094c13..1ffbd03 100644
--- a/JChat/Views/ChatViewModel.swift
+++ b/JChat/Views/ChatViewModel.swift
@@ -161,9 +161,7 @@ class ChatViewModel {
                         if let cachedModel = try? context.fetch(modelDescriptor).first {
                             assistantMessage.cost = cachedModel.calculateCost(promptTokens: prompt, completionTokens: completion)
                         }
-                        chat.totalPromptTokens += prompt
-                        chat.totalCompletionTokens += completion
-                        chat.totalCost += assistantMessage.cost
+                        chat.addUsage(promptTokens: prompt, completionTokens: completion, cost: assistantMessage.cost)
                     case .modelID(let id):
                         assistantMessage.modelID = id
                     case .done:
@@ -214,9 +212,7 @@ class ChatViewModel {
               sorted[index - 1].role == .user else { return }
 
         // Subtract old assistant message stats
-        chat.totalPromptTokens -= message.promptTokens
-        chat.totalCompletionTokens -= message.completionTokens
-        chat.totalCost -= message.cost
+        chat.removeUsage(promptTokens: message.promptTokens, completionTokens: message.completionTokens, cost: message.cost)
 
         // Delete old assistant message
         chat.messages.removeAll { $0.id == message.id }
@@ -304,9 +300,7 @@ class ChatViewModel {
                         if let cachedModel = try? context.fetch(modelDescriptor).first {
                             assistantMessage.cost = cachedModel.calculateCost(promptTokens: prompt, completionTokens: completion)
                         }
-                        chat.totalPromptTokens += prompt
-                        chat.totalCompletionTokens += completion
-                        chat.totalCost += assistantMessage.cost
+                        chat.addUsage(promptTokens: prompt, completionTokens: completion, cost: assistantMessage.cost)
                     case .modelID(let id):
                         assistantMessage.modelID = id
                     case .done:
@@ -335,9 +329,7 @@ class ChatViewModel {
 
     func deleteMessage(_ message: Message, in context: ModelContext) {
         guard let chat = selectedChat else { return }
-        chat.totalPromptTokens -= message.promptTokens
-        chat.totalCompletionTokens -= message.completionTokens
-        chat.totalCost -= message.cost
+        chat.removeUsage(promptTokens: message.promptTokens, completionTokens: message.completionTokens, cost: message.cost)
         chat.messages.removeAll { $0.id == message.id }
         context.delete(message)
         try? context.save()
diff --git a/JChatTests/JChatTests.swift b/JChatTests/JChatTests.swift
index af2994b..fe5bc6e 100644
--- a/JChatTests/JChatTests.swift
+++ b/JChatTests/JChatTests.swift
@@ -64,4 +64,14 @@ struct JChatTests {
         #expect(chat.overrideCount == 0)
     }
 
+    @Test func usageRemovalClampsAtZero() {
+        let chat = Chat()
+        chat.addUsage(promptTokens: 10, completionTokens: 5, cost: 0.12)
+        chat.removeUsage(promptTokens: 99, completionTokens: 99, cost: 1.0)
+
+        #expect(chat.totalPromptTokens == 0)
+        #expect(chat.totalCompletionTokens == 0)
+        #expect(chat.totalCost == 0)
+    }
+
 }