From a315021c755dc3f6dc13ce8f1b9cd47f69d9526c Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Fri, 10 Apr 2026 19:14:31 +0700
Subject: [PATCH] Fix desktop badge workspace parity

---
 packages/app/src/hooks/use-favicon-status.ts  | 36 ++++++++++---------
 .../app/src/utils/desktop-badge-state.test.ts | 34 ++++++++++++++++++
 packages/app/src/utils/desktop-badge-state.ts | 16 +++++++++
 3 files changed, 69 insertions(+), 17 deletions(-)
 create mode 100644 packages/app/src/utils/desktop-badge-state.test.ts
 create mode 100644 packages/app/src/utils/desktop-badge-state.ts

diff --git a/packages/app/src/hooks/use-favicon-status.ts b/packages/app/src/hooks/use-favicon-status.ts
index ebe47890e..5818a495a 100644
--- a/packages/app/src/hooks/use-favicon-status.ts
+++ b/packages/app/src/hooks/use-favicon-status.ts
@@ -1,8 +1,14 @@
 import { useEffect, useRef, useState } from "react";
 import { Platform } from "react-native";
+import { useShallow } from "zustand/shallow";
 import { getIsElectronRuntimeMac } from "@/constants/layout";
 import { useAggregatedAgents } from "./use-aggregated-agents";
 import { getDesktopHost } from "@/desktop/host";
+import { useSessionStore } from "@/stores/session-store";
+import {
+  deriveMacDockBadgeCountFromWorkspaceStatuses,
+  type DesktopBadgeWorkspaceStatus,
+} from "@/utils/desktop-badge-state";
 
 type FaviconStatus = "none" | "running" | "attention";
 type ColorScheme = "dark" | "light";
@@ -37,21 +43,6 @@ function deriveFaviconStatus(
   return "none";
 }
 
-function deriveMacDockBadgeCount(
-  agents: ReturnType<typeof useAggregatedAgents>["agents"],
-): number | undefined {
-  const attentionCount = agents.filter(
-    (agent) =>
-      (agent.pendingPermissionCount ?? 0) > 0 ||
-      (agent.requiresAttention && agent.attentionReason === "finished"),
-  ).length;
-  if (attentionCount > 0) {
-    return attentionCount;
-  }
-
-  return undefined;
-}
-
 function getFaviconUri(status: FaviconStatus, colorScheme: ColorScheme): string {
   const image = FAVICON_IMAGES[colorScheme][status];
   if (typeof image === "object" && "uri" in image) {
@@ -112,6 +103,17 @@ async function updateMacDockBadge(count?: number) {
 
 export function useFaviconStatus() {
   const { agents } = useAggregatedAgents();
+  const workspaceStatuses = useSessionStore(
+    useShallow((state) => {
+      const statuses: DesktopBadgeWorkspaceStatus[] = [];
+      for (const session of Object.values(state.sessions)) {
+        for (const workspace of session.workspaces.values()) {
+          statuses.push(workspace.status);
+        }
+      }
+      return statuses;
+    }),
+  );
   const [colorScheme, setColorScheme] = useState<ColorScheme>(getSystemColorScheme);
   const lastDockBadgeCountRef = useRef<number | undefined>(undefined);
 
@@ -135,10 +137,10 @@ export function useFaviconStatus() {
     const status = deriveFaviconStatus(agents);
     updateFavicon(status, colorScheme);
 
-    const dockBadgeCount = deriveMacDockBadgeCount(agents);
+    const dockBadgeCount = deriveMacDockBadgeCountFromWorkspaceStatuses(workspaceStatuses);
     if (dockBadgeCount !== lastDockBadgeCountRef.current) {
       lastDockBadgeCountRef.current = dockBadgeCount;
       void updateMacDockBadge(dockBadgeCount);
     }
-  }, [agents, colorScheme]);
+  }, [agents, colorScheme, workspaceStatuses]);
 }
diff --git a/packages/app/src/utils/desktop-badge-state.test.ts b/packages/app/src/utils/desktop-badge-state.test.ts
new file mode 100644
index 000000000..9bcdd39bc
--- /dev/null
+++ b/packages/app/src/utils/desktop-badge-state.test.ts
@@ -0,0 +1,34 @@
+import { describe, expect, it } from "vitest";
+import {
+  deriveMacDockBadgeCountFromWorkspaceStatuses,
+  isWorkspaceActionableForDesktopBadge,
+} from "./desktop-badge-state";
+
+describe("desktop-badge-state", () => {
+  it("treats attention-requiring workspace statuses as actionable", () => {
+    expect(isWorkspaceActionableForDesktopBadge("attention")).toBe(true);
+    expect(isWorkspaceActionableForDesktopBadge("needs_input")).toBe(true);
+    expect(isWorkspaceActionableForDesktopBadge("failed")).toBe(true);
+  });
+
+  it("ignores running and done workspace statuses", () => {
+    expect(isWorkspaceActionableForDesktopBadge("running")).toBe(false);
+    expect(isWorkspaceActionableForDesktopBadge("done")).toBe(false);
+  });
+
+  it("returns undefined when no visible workspaces need attention", () => {
+    expect(deriveMacDockBadgeCountFromWorkspaceStatuses(["done", "running"])).toBeUndefined();
+  });
+
+  it("counts only actionable visible workspaces", () => {
+    expect(
+      deriveMacDockBadgeCountFromWorkspaceStatuses([
+        "done",
+        "attention",
+        "running",
+        "needs_input",
+        "failed",
+      ]),
+    ).toBe(3);
+  });
+});
diff --git a/packages/app/src/utils/desktop-badge-state.ts b/packages/app/src/utils/desktop-badge-state.ts
new file mode 100644
index 000000000..9669492c5
--- /dev/null
+++ b/packages/app/src/utils/desktop-badge-state.ts
@@ -0,0 +1,16 @@
+import type { WorkspaceDescriptor } from "@/stores/session-store";
+
+export type DesktopBadgeWorkspaceStatus = WorkspaceDescriptor["status"];
+
+export function isWorkspaceActionableForDesktopBadge(
+  status: DesktopBadgeWorkspaceStatus,
+): boolean {
+  return status === "attention" || status === "needs_input" || status === "failed";
+}
+
+export function deriveMacDockBadgeCountFromWorkspaceStatuses(
+  statuses: readonly DesktopBadgeWorkspaceStatus[],
+): number | undefined {
+  const actionableCount = statuses.filter(isWorkspaceActionableForDesktopBadge).length;
+  return actionableCount > 0 ? actionableCount : undefined;
+}