Merge pull request olasunkanmi-SE#372 from olasunkanmi-SE/feature/cost_tracking

Feature/cost tracking

Author: olasunkanmi-SE

src/services/cost-tracking.service.ts

@@ -18,6 +18,30 @@ export interface IConversationCost extends ITokenUsage {
   requestCount: number;
 }
 
+export interface IProviderBreakdown {
+  provider: string;
+  inputTokens: number;
+  outputTokens: number;
+  totalTokens: number;
+  estimatedCostUSD: number;
+  requestCount: number;
+}
+
+export interface ICostTotals {
+  inputTokens: number;
+  outputTokens: number;
+  totalTokens: number;
+  estimatedCostUSD: number;
+  requestCount: number;
+  conversationCount: number;
+}
+
+export interface ICostSummary {
+  totals: ICostTotals;
+  providers: IProviderBreakdown[];
+  conversations: Array<{ threadId: string } & IConversationCost>;
+}
+
 interface IPricing {
   inputPerMillion: number;
   outputPerMillion: number;
@@ -142,6 +166,59 @@ export class CostTrackingService {
     this.conversations.clear();
   }
 
+  /**
+   * Returns an aggregated cost summary across all tracked conversations.
+   */
+  getCostSummary(): ICostSummary {
+    const conversations: ICostSummary["conversations"] = [];
+    const providerMap = new Map<string, IProviderBreakdown>();
+
+    let totalInput = 0;
+    let totalOutput = 0;
+    let totalCost = 0;
+    let totalRequests = 0;
+
+    for (const [threadId, cost] of this.conversations) {
+      conversations.push({ threadId, ...cost });
+
+      totalInput += cost.inputTokens;
+      totalOutput += cost.outputTokens;
+      totalCost += cost.estimatedCostUSD;
+      totalRequests += cost.requestCount;
+
+      const existing = providerMap.get(cost.provider);
+      if (existing) {
+        existing.inputTokens += cost.inputTokens;
+        existing.outputTokens += cost.outputTokens;
+        existing.totalTokens += cost.totalTokens;
+        existing.estimatedCostUSD += cost.estimatedCostUSD;
+        existing.requestCount += cost.requestCount;
+      } else {
+        providerMap.set(cost.provider, {
+          provider: cost.provider,
+          inputTokens: cost.inputTokens,
+          outputTokens: cost.outputTokens,
+          totalTokens: cost.totalTokens,
+          estimatedCostUSD: cost.estimatedCostUSD,
+          requestCount: cost.requestCount,
+        });
+      }
+    }
+
+    return {
+      totals: {
+        inputTokens: totalInput,
+        outputTokens: totalOutput,
+        totalTokens: totalInput + totalOutput,
+        estimatedCostUSD: Math.round(totalCost * 1_000_000) / 1_000_000,
+        requestCount: totalRequests,
+        conversationCount: conversations.length,
+      },
+      providers: [...providerMap.values()],
+      conversations,
+    };
+  }
+
   private getPricing(model: string): IPricing {
     // Try exact match first
     if (MODEL_PRICING[model]) {

src/test/suite/cost-tracking-handler.test.ts

@@ -0,0 +1,104 @@
+import * as assert from "assert";
+import * as sinon from "sinon";
+import { CostTrackingHandler } from "../../webview-providers/handlers/cost-tracking-handler";
+import { HandlerContext } from "../../webview-providers/handlers/types";
+import { CostTrackingService } from "../../services/cost-tracking.service";
+
+suite("CostTrackingHandler", () => {
+  let handler: CostTrackingHandler;
+  let ctx: HandlerContext;
+  let postMessageStub: sinon.SinonStub;
+  let service: CostTrackingService;
+
+  setup(() => {
+    service = CostTrackingService.getInstance();
+    service.resetAll();
+
+    handler = new CostTrackingHandler();
+
+    postMessageStub = sinon.stub().resolves(true);
+    ctx = {
+      webview: { webview: { postMessage: postMessageStub } },
+      logger: { warn: sinon.stub(), info: sinon.stub(), error: sinon.stub() },
+      extensionUri: {} as any,
+      sendResponse: sinon.stub(),
+    } as unknown as HandlerContext;
+  });
+
+  teardown(() => {
+    service.resetAll();
+    sinon.restore();
+  });
+
+  // ── commands ────────────────────────────────────────────────────
+
+  test("registers cost-summary and cost-reset commands", () => {
+    assert.ok(handler.commands.includes("cost-summary"));
+    assert.ok(handler.commands.includes("cost-reset"));
+    assert.strictEqual(handler.commands.length, 2);
+  });
+
+  // ── cost-summary ───────────────────────────────────────────────
+
+  test("cost-summary posts empty result when no data", async () => {
+    await handler.handle({ command: "cost-summary" }, ctx);
+
+    assert.ok(postMessageStub.calledOnce);
+    const msg = postMessageStub.firstCall.args[0];
+    assert.strictEqual(msg.type, "cost-summary-result");
+    assert.strictEqual(msg.totals.conversationCount, 0);
+    assert.strictEqual(msg.totals.totalTokens, 0);
+    assert.deepStrictEqual(msg.providers, []);
+    assert.deepStrictEqual(msg.conversations, []);
+  });
+
+  test("cost-summary returns aggregated data", async () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1000, 500);
+    service.recordUsage("t2", "openai", "gpt-4o", 2000, 1000);
+
+    await handler.handle({ command: "cost-summary" }, ctx);
+
+    assert.ok(postMessageStub.calledOnce);
+    const msg = postMessageStub.firstCall.args[0];
+    assert.strictEqual(msg.type, "cost-summary-result");
+    assert.strictEqual(msg.totals.conversationCount, 2);
+    assert.strictEqual(msg.totals.inputTokens, 3000);
+    assert.strictEqual(msg.totals.outputTokens, 1500);
+    assert.strictEqual(msg.totals.requestCount, 2);
+    assert.strictEqual(msg.providers.length, 2);
+    assert.strictEqual(msg.conversations.length, 2);
+  });
+
+  // ── cost-reset ─────────────────────────────────────────────────
+
+  test("cost-reset clears data and posts null totals", async () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1000, 500);
+
+    await handler.handle({ command: "cost-reset" }, ctx);
+
+    assert.ok(postMessageStub.calledOnce);
+    const msg = postMessageStub.firstCall.args[0];
+    assert.strictEqual(msg.type, "cost-summary-result");
+    assert.strictEqual(msg.totals, null);
+    assert.deepStrictEqual(msg.providers, []);
+    assert.deepStrictEqual(msg.conversations, []);
+
+    // Service should be cleared
+    assert.strictEqual(service.getConversationCost("t1"), null);
+  });
+
+  // ── unknown messages ───────────────────────────────────────────
+
+  test("ignores unrelated commands", async () => {
+    await handler.handle({ command: "some-other-command" }, ctx);
+    assert.ok(postMessageStub.notCalled);
+  });
+
+  test("ignores malformed messages", async () => {
+    await handler.handle({} as any, ctx);
+    assert.ok(postMessageStub.notCalled);
+
+    await handler.handle({ command: 123 } as any, ctx);
+    assert.ok(postMessageStub.notCalled);
+  });
+});

src/test/suite/cost-tracking.service.test.ts

@@ -0,0 +1,174 @@
+import * as assert from "assert";
+import { CostTrackingService } from "../../services/cost-tracking.service";
+
+suite("CostTrackingService", () => {
+  let service: CostTrackingService;
+
+  setup(() => {
+    // Reset singleton state between tests
+    service = CostTrackingService.getInstance();
+    service.resetAll();
+  });
+
+  teardown(() => {
+    service.resetAll();
+  });
+
+  // ── recordUsage ─────────────────────────────────────────────────
+
+  test("recordUsage tracks a single turn", () => {
+    const result = service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1000, 500);
+
+    assert.strictEqual(result.inputTokens, 1000);
+    assert.strictEqual(result.outputTokens, 500);
+    assert.strictEqual(result.totalTokens, 1500);
+    assert.strictEqual(result.requestCount, 1);
+    assert.strictEqual(result.provider, "anthropic");
+    assert.strictEqual(result.model, "claude-sonnet-4-6");
+    assert.ok(result.estimatedCostUSD > 0);
+  });
+
+  test("recordUsage accumulates across multiple turns", () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1000, 500);
+    const result = service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 2000, 1000);
+
+    assert.strictEqual(result.inputTokens, 3000);
+    assert.strictEqual(result.outputTokens, 1500);
+    assert.strictEqual(result.totalTokens, 4500);
+    assert.strictEqual(result.requestCount, 2);
+  });
+
+  test("recordUsage tracks separate conversations independently", () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1000, 500);
+    service.recordUsage("t2", "openai", "gpt-4o", 2000, 1000);
+
+    const c1 = service.getConversationCost("t1");
+    const c2 = service.getConversationCost("t2");
+
+    assert.strictEqual(c1!.inputTokens, 1000);
+    assert.strictEqual(c2!.inputTokens, 2000);
+    assert.strictEqual(c1!.provider, "anthropic");
+    assert.strictEqual(c2!.provider, "openai");
+  });
+
+  test("recordUsage uses default pricing for unknown models", () => {
+    const result = service.recordUsage("t1", "custom", "some-unknown-model", 1000000, 0);
+    // Default pricing: 3/M input → $3.00 for 1M tokens
+    assert.strictEqual(result.estimatedCostUSD, 3);
+  });
+
+  // ── getConversationCost ─────────────────────────────────────────
+
+  test("getConversationCost returns null for untracked thread", () => {
+    assert.strictEqual(service.getConversationCost("nonexistent"), null);
+  });
+
+  test("getConversationCost returns data for tracked thread", () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 100, 50);
+    const cost = service.getConversationCost("t1");
+
+    assert.ok(cost !== null);
+    assert.strictEqual(cost!.inputTokens, 100);
+    assert.strictEqual(cost!.outputTokens, 50);
+  });
+
+  // ── resetConversation ───────────────────────────────────────────
+
+  test("resetConversation clears a single conversation", () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 100, 50);
+    service.recordUsage("t2", "openai", "gpt-4o", 200, 100);
+
+    service.resetConversation("t1");
+
+    assert.strictEqual(service.getConversationCost("t1"), null);
+    assert.ok(service.getConversationCost("t2") !== null);
+  });
+
+  // ── resetAll ────────────────────────────────────────────────────
+
+  test("resetAll clears all conversations", () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 100, 50);
+    service.recordUsage("t2", "openai", "gpt-4o", 200, 100);
+
+    service.resetAll();
+
+    assert.strictEqual(service.getConversationCost("t1"), null);
+    assert.strictEqual(service.getConversationCost("t2"), null);
+  });
+
+  // ── getCostSummary ──────────────────────────────────────────────
+
+  test("getCostSummary returns zero totals when empty", () => {
+    const summary = service.getCostSummary();
+
+    assert.strictEqual(summary.totals.inputTokens, 0);
+    assert.strictEqual(summary.totals.outputTokens, 0);
+    assert.strictEqual(summary.totals.totalTokens, 0);
+    assert.strictEqual(summary.totals.estimatedCostUSD, 0);
+    assert.strictEqual(summary.totals.requestCount, 0);
+    assert.strictEqual(summary.totals.conversationCount, 0);
+    assert.deepStrictEqual(summary.providers, []);
+    assert.deepStrictEqual(summary.conversations, []);
+  });
+
+  test("getCostSummary aggregates a single conversation", () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1000, 500);
+
+    const summary = service.getCostSummary();
+
+    assert.strictEqual(summary.totals.inputTokens, 1000);
+    assert.strictEqual(summary.totals.outputTokens, 500);
+    assert.strictEqual(summary.totals.totalTokens, 1500);
+    assert.strictEqual(summary.totals.requestCount, 1);
+    assert.strictEqual(summary.totals.conversationCount, 1);
+    assert.strictEqual(summary.providers.length, 1);
+    assert.strictEqual(summary.providers[0].provider, "anthropic");
+    assert.strictEqual(summary.conversations.length, 1);
+    assert.strictEqual(summary.conversations[0].threadId, "t1");
+  });
+
+  test("getCostSummary aggregates multiple conversations per provider", () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1000, 500);
+    service.recordUsage("t2", "anthropic", "claude-sonnet-4-6", 2000, 1000);
+
+    const summary = service.getCostSummary();
+
+    assert.strictEqual(summary.totals.conversationCount, 2);
+    assert.strictEqual(summary.totals.inputTokens, 3000);
+    assert.strictEqual(summary.totals.outputTokens, 1500);
+    assert.strictEqual(summary.providers.length, 1);
+    assert.strictEqual(summary.providers[0].inputTokens, 3000);
+    assert.strictEqual(summary.providers[0].requestCount, 2);
+  });
+
+  test("getCostSummary groups by provider correctly", () => {
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1000, 500);
+    service.recordUsage("t2", "openai", "gpt-4o", 2000, 1000);
+    service.recordUsage("t3", "anthropic", "claude-3-5-haiku-20241022", 500, 250);
+
+    const summary = service.getCostSummary();
+
+    assert.strictEqual(summary.totals.conversationCount, 3);
+    assert.strictEqual(summary.providers.length, 2);
+
+    const anthropic = summary.providers.find((p) => p.provider === "anthropic");
+    const openai = summary.providers.find((p) => p.provider === "openai");
+
+    assert.ok(anthropic);
+    assert.ok(openai);
+    assert.strictEqual(anthropic!.inputTokens, 1500);
+    assert.strictEqual(openai!.inputTokens, 2000);
+    assert.strictEqual(anthropic!.requestCount, 2);
+    assert.strictEqual(openai!.requestCount, 1);
+  });
+
+  test("getCostSummary rounds cost to 6 decimal places", () => {
+    // Use a known model with known pricing to verify rounding
+    service.recordUsage("t1", "anthropic", "claude-sonnet-4-6", 1, 1);
+
+    const summary = service.getCostSummary();
+    const costStr = summary.totals.estimatedCostUSD.toString();
+    const decimals = costStr.includes(".") ? costStr.split(".")[1].length : 0;
+    assert.ok(decimals <= 6, `Cost should have at most 6 decimals, got ${decimals}`);
+  });
+});

src/webview-providers/base.ts

@@ -55,6 +55,7 @@ import {
   BrowserHandler,
   ComposerHandler,
   ConnectorHandler,
+  CostTrackingHandler,
   SkillHandler,
   DiffReviewHandler,
   CheckpointHandler,
@@ -270,6 +271,7 @@ export abstract class BaseWebViewProvider implements vscode.Disposable {
     this.handlerRegistry.register(new ComposerHandler());
     this.handlerRegistry.register(new StandupHandler());
     this.handlerRegistry.register(new TeamGraphHandler());
+    this.handlerRegistry.register(new CostTrackingHandler());
     this.handlerRegistry.register(new DoctorHandler());
     this.handlerRegistry.register(new OnboardingHandler());
     this.handlerRegistry.register(