feat: massive overhaul to manage costs (#103)

### Explanation _No response_ ### Issue Closes #102 ### Attestations - [ ] I have read and agree to the [Code of Conduct](https://docs.nhcarrigan.com/community/coc/) - [ ] I have read and agree to the [Community Guidelines](https://docs.nhcarrigan.com/community/guide/). - [ ] My contribution complies with the [Contributor Covenant](https://docs.nhcarrigan.com/dev/covenant/). ### Dependencies - [ ] I have pinned the dependencies to a specific patch version. ### Style - [ ] I have run the linter and resolved any errors. - [ ] My pull request uses an appropriate title, matching the conventional commit standards. - [ ] My scope of feat/fix/chore/etc. correctly matches the nature of changes in my pull request. ### Tests - [ ] My contribution adds new code, and I have added tests to cover it. - [ ] My contribution modifies existing code, and I have updated the tests to reflect these changes. - [ ] All new and existing tests pass locally with my changes. - [ ] Code coverage remains at or above the configured threshold. ### Documentation _No response_ ### Versioning _No response_ Reviewed-on: #103 Co-authored-by: Naomi Carrigan <commits@nhcarrigan.com> Co-committed-by: Naomi Carrigan <commits@nhcarrigan.com>
2026-02-04 19:58:43 -08:00
parent daedbfd865
commit 1c45507cdf
30 changed files with 4024 additions and 103 deletions
@@ -1,7 +1,25 @@
 import { describe, it, expect, beforeEach, vi } from "vitest";
 import { get } from "svelte/store";
-import { stats, formattedStats, resetSessionStats } from "./stats";
-import type { UsageStats } from "./stats";
+import {
+  stats,
+  formattedStats,
+  resetSessionStats,
+  contextWarning,
+  getContextWarningMessage,
+  estimateMessageCost,
+  formatTokenCount,
+  MODEL_PRICING,
+} from "./stats";
+import type { UsageStats, ToolTokenStats } from "./stats";
+
+// Helper function to create ToolTokenStats for tests
+function toolStats(callCount: number, inputTokens = 0, outputTokens = 0): ToolTokenStats {
+  return {
+    call_count: callCount,
+    estimated_input_tokens: inputTokens,
+    estimated_output_tokens: outputTokens,
+  };
+}

 // Mock Tauri APIs
 vi.mock("@tauri-apps/api/event", () => ({
@@ -34,6 +52,11 @@ describe("stats store", () => {
      tools_usage: {},
      session_tools_usage: {},
      session_duration_seconds: 0,
+      context_tokens_used: 0,
+      context_window_limit: 200000,
+      context_utilisation_percent: 0,
+      potential_cache_hits: 0,
+      potential_cache_savings_tokens: 0,
    });
  });

@@ -63,9 +86,14 @@ describe("stats store", () => {
        session_files_edited: 2,
        files_created: 1,
        session_files_created: 1,
-        tools_usage: { Read: 5, Edit: 3 },
-        session_tools_usage: { Read: 2, Edit: 1 },
+        tools_usage: { Read: toolStats(5), Edit: toolStats(3) },
+        session_tools_usage: { Read: toolStats(2), Edit: toolStats(1) },
        session_duration_seconds: 300,
+        context_tokens_used: 500,
+        context_window_limit: 200000,
+        context_utilisation_percent: 0.25,
+        potential_cache_hits: 0,
+        potential_cache_savings_tokens: 0,
      };

      stats.set(newStats);
@@ -74,7 +102,8 @@ describe("stats store", () => {
      expect(currentStats.total_input_tokens).toBe(1000);
      expect(currentStats.total_output_tokens).toBe(2000);
      expect(currentStats.model).toBe("claude-sonnet-4");
-      expect(currentStats.tools_usage).toEqual({ Read: 5, Edit: 3 });
+      expect(currentStats.tools_usage.Read?.call_count).toBe(5);
+      expect(currentStats.tools_usage.Edit?.call_count).toBe(3);
    });

    it("can be updated with update function", () => {
@@ -109,9 +138,14 @@ describe("stats store", () => {
        session_files_edited: 2,
        files_created: 1,
        session_files_created: 1,
-        tools_usage: { Read: 5, Edit: 3 },
-        session_tools_usage: { Read: 2, Edit: 1 },
+        tools_usage: { Read: toolStats(5), Edit: toolStats(3) },
+        session_tools_usage: { Read: toolStats(2), Edit: toolStats(1) },
        session_duration_seconds: 300,
+        context_tokens_used: 500,
+        context_window_limit: 200000,
+        context_utilisation_percent: 0.25,
+        potential_cache_hits: 0,
+        potential_cache_savings_tokens: 0,
      });

      // Reset session stats
@@ -127,7 +161,8 @@ describe("stats store", () => {
      expect(currentStats.code_blocks_generated).toBe(3);
      expect(currentStats.files_edited).toBe(5);
      expect(currentStats.files_created).toBe(1);
-      expect(currentStats.tools_usage).toEqual({ Read: 5, Edit: 3 });
+      expect(currentStats.tools_usage.Read?.call_count).toBe(5);
+      expect(currentStats.tools_usage.Edit?.call_count).toBe(3);
      expect(currentStats.model).toBe("claude-sonnet-4");

      // Session stats should be reset
@@ -277,8 +312,8 @@ describe("stats store", () => {
    });

    it("exposes tools usage directly", () => {
-      const toolsUsage = { Read: 10, Edit: 5, Write: 3 };
-      const sessionToolsUsage = { Read: 2, Edit: 1 };
+      const toolsUsage = { Read: toolStats(10), Edit: toolStats(5), Write: toolStats(3) };
+      const sessionToolsUsage = { Read: toolStats(2), Edit: toolStats(1) };

      stats.update((current) => ({
        ...current,
@@ -331,9 +366,14 @@ describe("stats store", () => {
        session_files_edited: 1,
        files_created: 1,
        session_files_created: 0,
-        tools_usage: { Read: 3 },
-        session_tools_usage: { Read: 1 },
+        tools_usage: { Read: toolStats(3) },
+        session_tools_usage: { Read: toolStats(1) },
        session_duration_seconds: 60,
+        context_tokens_used: 50,
+        context_window_limit: 200000,
+        context_utilisation_percent: 0.025,
+        potential_cache_hits: 0,
+        potential_cache_savings_tokens: 0,
      };

      stats.set(fullStats);
@@ -343,4 +383,221 @@ describe("stats store", () => {
      expect(currentStats).toEqual(fullStats);
    });
  });
+
+  describe("context window tracking", () => {
+    it("tracks context tokens used", () => {
+      stats.update((current) => ({
+        ...current,
+        context_tokens_used: 100000,
+        context_window_limit: 200000,
+        context_utilisation_percent: 50.0,
+      }));
+
+      const currentStats = get(stats);
+      expect(currentStats.context_tokens_used).toBe(100000);
+      expect(currentStats.context_window_limit).toBe(200000);
+      expect(currentStats.context_utilisation_percent).toBe(50.0);
+    });
+
+    it("formats context stats correctly", () => {
+      stats.update((current) => ({
+        ...current,
+        context_tokens_used: 150000,
+        context_window_limit: 200000,
+        context_utilisation_percent: 75.5,
+      }));
+
+      const formatted = get(formattedStats);
+      expect(formatted.contextUsed).toBe("150,000");
+      expect(formatted.contextLimit).toBe("200,000");
+      expect(formatted.contextRemaining).toBe("50,000");
+      expect(formatted.contextUtilisation).toBe("75.5%");
+    });
+
+    it("calculates remaining tokens correctly at limit", () => {
+      stats.update((current) => ({
+        ...current,
+        context_tokens_used: 200000,
+        context_window_limit: 200000,
+        context_utilisation_percent: 100.0,
+      }));
+
+      const formatted = get(formattedStats);
+      expect(formatted.contextRemaining).toBe("0");
+    });
+
+    it("handles over-limit gracefully", () => {
+      stats.update((current) => ({
+        ...current,
+        context_tokens_used: 250000,
+        context_window_limit: 200000,
+        context_utilisation_percent: 125.0,
+      }));
+
+      const formatted = get(formattedStats);
+      expect(formatted.contextRemaining).toBe("0");
+    });
+  });
+
+  describe("contextWarning derived store", () => {
+    it("returns null when under 50%", () => {
+      stats.update((current) => ({
+        ...current,
+        context_utilisation_percent: 40.0,
+      }));
+
+      const warning = get(contextWarning);
+      expect(warning).toBeNull();
+    });
+
+    it("returns moderate when between 50-74%", () => {
+      stats.update((current) => ({
+        ...current,
+        context_utilisation_percent: 60.0,
+      }));
+
+      const warning = get(contextWarning);
+      expect(warning).toBe("moderate");
+    });
+
+    it("returns high when between 75-89%", () => {
+      stats.update((current) => ({
+        ...current,
+        context_utilisation_percent: 80.0,
+      }));
+
+      const warning = get(contextWarning);
+      expect(warning).toBe("high");
+    });
+
+    it("returns critical when 90%+", () => {
+      stats.update((current) => ({
+        ...current,
+        context_utilisation_percent: 95.0,
+      }));
+
+      const warning = get(contextWarning);
+      expect(warning).toBe("critical");
+    });
+  });
+
+  describe("getContextWarningMessage", () => {
+    it("returns correct message for moderate warning", () => {
+      const message = getContextWarningMessage("moderate");
+      expect(message).toContain("50%+");
+      expect(message).toContain("Consider starting a new conversation");
+    });
+
+    it("returns correct message for high warning", () => {
+      const message = getContextWarningMessage("high");
+      expect(message).toContain("75%+");
+      expect(message).toContain("Responses may degrade");
+    });
+
+    it("returns correct message for critical warning", () => {
+      const message = getContextWarningMessage("critical");
+      expect(message).toContain("90%+");
+      expect(message).toContain("Start a new conversation");
+    });
+  });
+
+  describe("formatTokenCount", () => {
+    it("formats small numbers directly", () => {
+      expect(formatTokenCount(0)).toBe("0");
+      expect(formatTokenCount(100)).toBe("100");
+      expect(formatTokenCount(999)).toBe("999");
+    });
+
+    it("formats thousands with K suffix", () => {
+      expect(formatTokenCount(1000)).toBe("1.0K");
+      expect(formatTokenCount(1500)).toBe("1.5K");
+      expect(formatTokenCount(10000)).toBe("10.0K");
+      expect(formatTokenCount(999999)).toBe("1000.0K");
+    });
+
+    it("formats millions with M suffix", () => {
+      expect(formatTokenCount(1000000)).toBe("1.0M");
+      expect(formatTokenCount(1500000)).toBe("1.5M");
+      expect(formatTokenCount(10000000)).toBe("10.0M");
+    });
+  });
+
+  describe("estimateMessageCost", () => {
+    it("estimates tokens at ~4 chars per token", () => {
+      const result = estimateMessageCost("test", 0, null); // 4 chars = 1 token
+      expect(result.messageTokens).toBe(1);
+    });
+
+    it("rounds up partial tokens", () => {
+      const result = estimateMessageCost("a", 0, null); // 1 char rounds up to 1 token
+      expect(result.messageTokens).toBe(1);
+
+      const result2 = estimateMessageCost("abcde", 0, null); // 5 chars = 2 tokens
+      expect(result2.messageTokens).toBe(2);
+    });
+
+    it("returns 0 tokens for empty string", () => {
+      const result = estimateMessageCost("", 0, null);
+      expect(result.messageTokens).toBe(0);
+      expect(result.estimatedCost).toBe(0);
+    });
+
+    it("adds context tokens to total", () => {
+      const result = estimateMessageCost("test", 1000, null); // 1 token + 1000 context
+      expect(result.messageTokens).toBe(1);
+      expect(result.totalInputTokens).toBe(1001);
+    });
+
+    it("calculates cost using Sonnet pricing by default", () => {
+      // 100 chars = 25 tokens, $3 per million input tokens
+      const result = estimateMessageCost("a".repeat(100), 0, null);
+      expect(result.messageTokens).toBe(25);
+      const expectedCost = (25 / 1_000_000) * 3.0;
+      expect(result.estimatedCost).toBeCloseTo(expectedCost, 8);
+    });
+
+    it("uses Opus pricing for Opus models", () => {
+      const result = estimateMessageCost("a".repeat(100), 0, "claude-opus-4-5-20251101");
+      expect(result.messageTokens).toBe(25);
+      const expectedCost = (25 / 1_000_000) * 5.0; // Opus 4.5: $5 per million input
+      expect(result.estimatedCost).toBeCloseTo(expectedCost, 8);
+    });
+
+    it("uses Haiku pricing for Haiku models", () => {
+      const result = estimateMessageCost("a".repeat(100), 0, "claude-3-5-haiku-20241022");
+      expect(result.messageTokens).toBe(25);
+      const expectedCost = (25 / 1_000_000) * 1.0; // Haiku: $1 per million
+      expect(result.estimatedCost).toBeCloseTo(expectedCost, 8);
+    });
+
+    it("falls back to Sonnet pricing for unknown models", () => {
+      const result = estimateMessageCost("a".repeat(100), 0, "unknown-model");
+      expect(result.messageTokens).toBe(25);
+      const expectedCost = (25 / 1_000_000) * 3.0; // Default Sonnet: $3 per million
+      expect(result.estimatedCost).toBeCloseTo(expectedCost, 8);
+    });
+  });
+
+  describe("MODEL_PRICING", () => {
+    it("contains expected Opus pricing", () => {
+      // Opus 4.5 has reduced pricing
+      expect(MODEL_PRICING["claude-opus-4-5-20251101"]).toEqual({ input: 5.0, output: 25.0 });
+      // Previous Opus models have higher pricing
+      expect(MODEL_PRICING["claude-opus-4-1-20250805"]).toEqual({ input: 15.0, output: 75.0 });
+      expect(MODEL_PRICING["claude-opus-4-20250514"]).toEqual({ input: 15.0, output: 75.0 });
+    });
+
+    it("contains expected Sonnet pricing", () => {
+      expect(MODEL_PRICING["claude-sonnet-4-5-20250929"]).toEqual({ input: 3.0, output: 15.0 });
+      expect(MODEL_PRICING["claude-sonnet-4-20250514"]).toEqual({ input: 3.0, output: 15.0 });
+      expect(MODEL_PRICING["claude-3-7-sonnet-20250219"]).toEqual({ input: 3.0, output: 15.0 });
+      expect(MODEL_PRICING["claude-3-5-sonnet-20241022"]).toEqual({ input: 3.0, output: 15.0 });
+    });
+
+    it("contains expected Haiku pricing", () => {
+      expect(MODEL_PRICING["claude-haiku-4-5-20251001"]).toEqual({ input: 1.0, output: 5.0 });
+      expect(MODEL_PRICING["claude-3-5-haiku-20241022"]).toEqual({ input: 1.0, output: 5.0 });
+      expect(MODEL_PRICING["claude-3-haiku-20240307"]).toEqual({ input: 0.25, output: 1.25 });
+    });
+  });
 });