🤖 fix: attribute session usage sources and dedupe abort cleanup

src/browser/components/RightSidebar/CostsTab.tsx

@@ -2,9 +2,11 @@ import React from "react";
 import { useWorkspaceUsage, useWorkspaceConsumers } from "@/browser/stores/WorkspaceStore";
 import { getModelStatsResolved } from "@/common/utils/tokens/modelStats";
 import {
+  getTotalCost,
   sumUsageHistory,
   formatCostWithDollar,
   type ChatUsageDisplay,
+  type SessionUsageSource,
 } from "@/common/utils/tokens/usageAggregator";
 import { usePersistedState } from "@/browser/hooks/usePersistedState";
 import { PREFERRED_COMPACTION_MODEL_KEY } from "@/common/constants/storage";
@@ -52,6 +54,23 @@ const VIEW_MODE_OPTIONS: Array<ToggleOption<ViewMode>> = [
   { value: "last-request", label: "Last Request" },
 ];
 
+const SOURCE_LABELS: Record<SessionUsageSource, string> = {
+  main: "Main",
+  system1: "System 1",
+  plan: "Plan",
+  subagent: "Sub-agent",
+};
+
+function getUsageTotalTokens(usage: ChatUsageDisplay): number {
+  return (
+    usage.input.tokens +
+    usage.cached.tokens +
+    usage.cacheCreate.tokens +
+    usage.output.tokens +
+    usage.reasoning.tokens
+  );
+}
+
 interface CostsTabProps {
   workspaceId: string;
 }
@@ -120,6 +139,22 @@ const CostsTabComponent: React.FC<CostsTabProps> = ({ workspaceId }) => {
   // Cost and Details table use viewMode
   const displayUsage = viewMode === "last-request" ? lastRequestUsage : sessionUsage;
 
+  const sourceRows =
+    viewMode === "session" && usage.sourceTotals
+      ? Object.entries(usage.sourceTotals)
+          .map(([source, sourceUsage]) => {
+            const typedSource = source as SessionUsageSource;
+            return {
+              source,
+              label: SOURCE_LABELS[typedSource] ?? source,
+              tokens: getUsageTotalTokens(sourceUsage),
+              cost: getTotalCost(sourceUsage),
+            };
+          })
+          .filter((row) => row.tokens > 0)
+          .sort((a, b) => b.tokens - a.tokens)
+      : [];
+
   return (
     <div className="text-light font-primary text-[13px] leading-relaxed">
       {hasUsageData && (
@@ -423,6 +458,36 @@ const CostsTabComponent: React.FC<CostsTabProps> = ({ workspaceId }) => {
                       })}
                     </tbody>
                   </table>
+
+                  {sourceRows.length > 0 && (
+                    <div data-testid="cost-source-breakdown" className="mt-3">
+                      <h3 className="text-subtle m-0 mb-1 text-[11px] font-semibold tracking-wide uppercase">
+                        Source Breakdown
+                      </h3>
+                      <table className="w-full border-collapse text-[11px]">
+                        <thead>
+                          <tr className="border-border-light border-b">
+                            <th className="text-muted py-1 pr-2 text-left font-medium">Source</th>
+                            <th className="text-muted py-1 pr-2 text-right font-medium">Tokens</th>
+                            <th className="text-muted py-1 text-right font-medium">Cost</th>
+                          </tr>
+                        </thead>
+                        <tbody>
+                          {sourceRows.map((row) => (
+                            <tr key={row.source}>
+                              <td className="text-foreground py-1 pr-2">{row.label}</td>
+                              <td className="text-foreground py-1 pr-2 text-right">
+                                {formatTokens(row.tokens)}
+                              </td>
+                              <td className="text-foreground py-1 text-right">
+                                {formatCostWithDollar(row.cost)}
+                              </td>
+                            </tr>
+                          ))}
+                        </tbody>
+                      </table>
+                    </div>
+                  )}
                 </>
               );
             })()}

src/browser/stores/WorkspaceStore.test.ts

@@ -15,7 +15,7 @@ interface LoadMoreResponse {
 
 // Mock client
 // eslint-disable-next-line require-yield
-const mockOnChat = mock(async function* (
+const defaultOnChatImplementation = async function* (
   _input?: { workspaceId: string; mode?: unknown },
   options?: { signal?: AbortSignal }
 ): AsyncGenerator<WorkspaceChatMessage, void, unknown> {
@@ -27,7 +27,9 @@ const mockOnChat = mock(async function* (
     }
     options.signal.addEventListener("abort", () => resolve(), { once: true });
   });
-});
+};
+
+const mockOnChat = mock(defaultOnChatImplementation);
 
 const mockGetSessionUsage = mock((_input: { workspaceId: string }) =>
   Promise.resolve<unknown>(undefined)
@@ -178,6 +180,16 @@ function createHistoryMessageEvent(id: string, historySequence: number): Workspa
   };
 }
 
+function createUsageDisplay(inputTokens: number, outputTokens: number) {
+  return {
+    input: { tokens: inputTokens },
+    cached: { tokens: 0 },
+    cacheCreate: { tokens: 0 },
+    output: { tokens: outputTokens },
+    reasoning: { tokens: 0 },
+  };
+}
+
 async function waitForAbortSignal(signal?: AbortSignal): Promise<void> {
   await new Promise<void>((resolve) => {
     if (!signal) {
@@ -194,6 +206,7 @@ describe("WorkspaceStore", () => {
 
   beforeEach(() => {
     mockOnChat.mockClear();
+    mockOnChat.mockImplementation(defaultOnChatImplementation);
     mockGetSessionUsage.mockClear();
     mockHistoryLoadMore.mockClear();
     mockActivityList.mockClear();
@@ -725,6 +738,96 @@ describe("WorkspaceStore", () => {
       expect(usage.sessionTotal!.input.tokens).toBe(1000);
     });
 
+    it("hydrates source totals from persisted session usage", async () => {
+      const workspaceId = "workspace-source-totals";
+      const sessionUsageData = {
+        byModel: {
+          "claude-sonnet-4": createUsageDisplay(1000, 100),
+        },
+        bySource: {
+          main: createUsageDisplay(900, 90),
+          system1: createUsageDisplay(100, 10),
+        },
+        version: 1 as const,
+      };
+
+      mockGetSessionUsage.mockImplementation(
+        ({ workspaceId: requestedWorkspaceId }: { workspaceId: string }) => {
+          if (requestedWorkspaceId === workspaceId) {
+            return Promise.resolve(sessionUsageData);
+          }
+          return Promise.resolve(undefined);
+        }
+      );
+
+      createAndAddWorkspace(store, workspaceId, {}, false);
+      store.setActiveWorkspaceId(workspaceId);
+      await new Promise((resolve) => setTimeout(resolve, 10));
+
+      const usage = store.getWorkspaceUsage(workspaceId);
+      expect(usage.sourceTotals?.main?.input.tokens).toBe(900);
+      expect(usage.sourceTotals?.system1?.output.tokens).toBe(10);
+    });
+
+    it("merges source totals when session-usage-delta events arrive", async () => {
+      const workspaceId = "workspace-source-delta";
+      mockGetSessionUsage.mockImplementation(
+        ({ workspaceId: requestedWorkspaceId }: { workspaceId: string }) => {
+          if (requestedWorkspaceId === workspaceId) {
+            return Promise.resolve({
+              byModel: {
+                "claude-sonnet-4": createUsageDisplay(100, 50),
+              },
+              bySource: {
+                main: createUsageDisplay(100, 50),
+              },
+              version: 1 as const,
+            });
+          }
+          return Promise.resolve(undefined);
+        }
+      );
+
+      mockOnChat.mockImplementation(async function* (
+        _input?: { workspaceId: string; mode?: unknown },
+        options?: { signal?: AbortSignal }
+      ): AsyncGenerator<WorkspaceChatMessage, void, unknown> {
+        yield { type: "caught-up" };
+        yield {
+          type: "session-usage-delta",
+          workspaceId,
+          sourceWorkspaceId: "child-subagent",
+          byModelDelta: {
+            "claude-sonnet-4": createUsageDisplay(20, 10),
+          },
+          bySourceDelta: {
+            main: createUsageDisplay(5, 2),
+            subagent: createUsageDisplay(15, 8),
+          },
+          timestamp: Date.now(),
+        };
+
+        await waitForAbortSignal(options?.signal);
+      });
+
+      createAndAddWorkspace(store, workspaceId);
+
+      const deadline = Date.now() + 1_000;
+      while (Date.now() < deadline) {
+        const usage = store.getWorkspaceUsage(workspaceId);
+        if (usage.sourceTotals?.subagent?.input.tokens === 15) {
+          break;
+        }
+        await new Promise((resolve) => setTimeout(resolve, 10));
+      }
+
+      const usage = store.getWorkspaceUsage(workspaceId);
+      expect(usage.sourceTotals?.main?.input.tokens).toBe(105);
+      expect(usage.sourceTotals?.main?.output.tokens).toBe(52);
+      expect(usage.sourceTotals?.subagent?.input.tokens).toBe(15);
+      expect(usage.sourceTotals?.subagent?.output.tokens).toBe(8);
+    });
+
     it("ignores stale session-usage fetch when a newer refresh supersedes it", async () => {
       let resolveFirst!: (value: unknown) => void;
       const firstFetch = new Promise((resolve) => {