diff --git a/apps/mobile/src/components/kilo-chat/hooks/mark-read-operation.ts b/apps/mobile/src/components/kilo-chat/hooks/mark-read-operation.ts index 8467bd3369..d9abb6344c 100644 --- a/apps/mobile/src/components/kilo-chat/hooks/mark-read-operation.ts +++ b/apps/mobile/src/components/kilo-chat/hooks/mark-read-operation.ts @@ -22,16 +22,13 @@ export async function markReadConversation({ return result; } -type ApplyBadgeClearResultInput = { +type ApplyBadgeClearResultInput = { badgeClear: MarkConversationReadResponse['badgeClear']; - startBadgeFreshnessEpoch: number; - currentBadgeFreshnessEpoch: number; userId: string | null; updateBadgeRows: ( queryKey: readonly ['badges', string], updater: (badges: BadgeCountRow[] | undefined) => BadgeCountRow[] | undefined ) => void; - setBadgeCount: (badgeCount: number) => Promise; }; export function filterClearedBadgeBucket( @@ -45,26 +42,14 @@ export function filterClearedBadgeBucket( return badges?.filter(row => row.badgeBucket !== badgeClear.badgeBucket); } -export function applyBadgeClearResult({ +export function applyBadgeClearResult({ badgeClear, - startBadgeFreshnessEpoch, - currentBadgeFreshnessEpoch, userId, updateBadgeRows, - setBadgeCount, -}: ApplyBadgeClearResultInput): boolean { - if (badgeClear === null) { - return false; - } - - if (userId !== null) { - updateBadgeRows(['badges', userId], badges => filterClearedBadgeBucket(badges, badgeClear)); - } - - if (currentBadgeFreshnessEpoch !== startBadgeFreshnessEpoch) { - return false; +}: ApplyBadgeClearResultInput): void { + if (badgeClear === null || userId === null) { + return; } - void setBadgeCount(badgeClear.badgeCount); - return true; + updateBadgeRows(['badges', userId], badges => filterClearedBadgeBucket(badges, badgeClear)); } diff --git a/apps/mobile/src/components/kilo-chat/hooks/use-mark-read.ts b/apps/mobile/src/components/kilo-chat/hooks/use-mark-read.ts index efb7ae1e2e..7c9a482a72 100644 --- a/apps/mobile/src/components/kilo-chat/hooks/use-mark-read.ts +++ b/apps/mobile/src/components/kilo-chat/hooks/use-mark-read.ts @@ -1,7 +1,6 @@ import { useCallback } from 'react'; import { useMutation, useQueryClient } from '@tanstack/react-query'; import * as Sentry from '@sentry/react-native'; -import * as Notifications from 'expo-notifications'; import { type KiloChatClient, type MarkConversationReadResponse } from '@kilocode/kilo-chat'; import { type BadgeCountRow } from '@kilocode/notifications'; @@ -9,7 +8,6 @@ import { useMarkConversationRead } from '@kilocode/kilo-chat-hooks'; import { useCurrentUserId } from './use-current-user-id'; import { applyBadgeClearResult, markReadConversation } from './mark-read-operation'; -import { advanceBadgeFreshnessEpoch, readBadgeFreshnessEpoch } from '@/lib/badge-freshness'; type MarkReadInput = { sandboxId: string; @@ -48,18 +46,13 @@ export function useMarkRead(client: KiloChatClient) { extra: { hasUser: userId !== null }, }); }, - onMutate: () => ({ startBadgeFreshnessEpoch: advanceBadgeFreshnessEpoch() }), - onSuccess: (result, _variables, context) => { - const currentBadgeFreshnessEpoch = readBadgeFreshnessEpoch(); + onSuccess: result => { applyBadgeClearResult({ badgeClear: result.badgeClear, - startBadgeFreshnessEpoch: context.startBadgeFreshnessEpoch, - currentBadgeFreshnessEpoch, userId, updateBadgeRows: (queryKey, updater) => { queryClient.setQueryData(queryKey, updater); }, - setBadgeCount: Notifications.setBadgeCountAsync, }); }, onSettled: () => { diff --git a/apps/mobile/src/components/kilo-chat/mark-read-state.test.ts b/apps/mobile/src/components/kilo-chat/mark-read-state.test.ts index 9b2805a8da..1afce4055a 100644 --- a/apps/mobile/src/components/kilo-chat/mark-read-state.test.ts +++ b/apps/mobile/src/components/kilo-chat/mark-read-state.test.ts @@ -13,7 +13,6 @@ import { filterClearedBadgeBucket, markReadConversation, } from './hooks/mark-read-operation'; -import { reconcileHydratedBadgeCount, totalBadgeCount } from '@/lib/badge-hydration'; type UpdateBadgeRows = ( queryKey: readonly ['badges', string], @@ -105,116 +104,28 @@ describe('markReadConversation', () => { ).toEqual([{ badgeBucket: 'bucket-1', badgeCount: 2 }]); }); - it('does not update badge cache or OS badge count when badgeClear is null', () => { + it('does not update the badge cache when badgeClear is null', () => { const updateBadgeRows = createUpdateBadgeRowsMock(); - const setBadgeCount = vi.fn<(badgeCount: number) => Promise>(async () => { - const result = await Promise.resolve(true); - return result; - }); - const applied = applyBadgeClearResult({ + applyBadgeClearResult({ badgeClear: null, - startBadgeFreshnessEpoch: 0, - currentBadgeFreshnessEpoch: 0, userId: 'user-1', updateBadgeRows, - setBadgeCount, }); - expect(applied).toBe(false); expect(updateBadgeRows).not.toHaveBeenCalled(); - expect(setBadgeCount).not.toHaveBeenCalled(); }); - it('updates badge cache and OS badge count when badgeClear includes a count with unchanged freshness', () => { + it('updates the badge cache when badgeClear contains a cleared row', () => { const updateBadgeRows = createUpdateBadgeRowsMock(); - const setBadgeCount = vi.fn<(badgeCount: number) => Promise>(async () => { - const result = await Promise.resolve(true); - return result; - }); - const applied = applyBadgeClearResult({ + applyBadgeClearResult({ badgeClear: { badgeBucket: 'server-bucket', badgeCount: 3 }, - startBadgeFreshnessEpoch: 4, - currentBadgeFreshnessEpoch: 4, userId: 'user-1', updateBadgeRows, - setBadgeCount, }); - expect(applied).toBe(true); expect(updateBadgeRows).toHaveBeenCalledOnce(); expect(updateBadgeRows).toHaveBeenCalledWith(['badges', 'user-1'], expect.any(Function)); - expect(setBadgeCount).toHaveBeenCalledWith(3); - }); - - it('keeps badge cache updates but skips stale OS badge counts when freshness advanced', () => { - const updateBadgeRows = createUpdateBadgeRowsMock(); - const setBadgeCount = vi.fn<(badgeCount: number) => Promise>(async () => { - const result = await Promise.resolve(true); - return result; - }); - - const applied = applyBadgeClearResult({ - badgeClear: { badgeBucket: 'server-bucket', badgeCount: 0 }, - startBadgeFreshnessEpoch: 8, - currentBadgeFreshnessEpoch: 9, - userId: 'user-1', - updateBadgeRows, - setBadgeCount, - }); - - expect(applied).toBe(false); - expect(updateBadgeRows).toHaveBeenCalledOnce(); - expect(updateBadgeRows).toHaveBeenCalledWith(['badges', 'user-1'], expect.any(Function)); - expect(setBadgeCount).not.toHaveBeenCalled(); - }); -}); - -describe('badge hydration reconciliation', () => { - it('totals all hydrated badge buckets for the native OS badge', () => { - expect( - totalBadgeCount([ - { badgeBucket: 'kiloclaw:sandbox-1', badgeCount: 2 }, - { badgeBucket: 'kiloclaw:sandbox-1:conversation-1', badgeCount: 3 }, - ]) - ).toBe(5); - }); - - it('updates the native OS badge when hydration is still fresh', () => { - const setBadgeCount = vi.fn<(badgeCount: number) => Promise>(async () => { - const result = await Promise.resolve(true); - return result; - }); - - const applied = reconcileHydratedBadgeCount({ - badgeRows: [ - { badgeBucket: 'kiloclaw:sandbox-1', badgeCount: 2 }, - { badgeBucket: 'kiloclaw:sandbox-1:conversation-1', badgeCount: 3 }, - ], - startBadgeFreshnessEpoch: 10, - currentBadgeFreshnessEpoch: 10, - setBadgeCount, - }); - - expect(applied).toBe(true); - expect(setBadgeCount).toHaveBeenCalledWith(5); - }); - - it('does not overwrite a newer native OS badge update from stale hydration', () => { - const setBadgeCount = vi.fn<(badgeCount: number) => Promise>(async () => { - const result = await Promise.resolve(true); - return result; - }); - - const applied = reconcileHydratedBadgeCount({ - badgeRows: [{ badgeBucket: 'kiloclaw:sandbox-1', badgeCount: 4 }], - startBadgeFreshnessEpoch: 10, - currentBadgeFreshnessEpoch: 11, - setBadgeCount, - }); - - expect(applied).toBe(false); - expect(setBadgeCount).not.toHaveBeenCalled(); }); }); diff --git a/apps/mobile/src/lib/active-sessions-live-sync.ts b/apps/mobile/src/lib/active-sessions-live-sync.ts index 79c86453f9..52e29c67f4 100644 --- a/apps/mobile/src/lib/active-sessions-live-sync.ts +++ b/apps/mobile/src/lib/active-sessions-live-sync.ts @@ -333,6 +333,8 @@ export class ActiveSessionsLiveSync { /** One app-level mount owns the registry and socket lease. */ let attachedSync: ActiveSessionsLiveSync | null = null; +export const refreshActiveSessionsFromPush = (): void => attachedSync?.scheduleRefresh('manual'); + /** Returns false if no current owner handles this exact key. */ export async function refreshActiveSessionsNow( queryKey: QueryKey diff --git a/apps/mobile/src/lib/badge-freshness.ts b/apps/mobile/src/lib/badge-freshness.ts deleted file mode 100644 index 66327bd29b..0000000000 --- a/apps/mobile/src/lib/badge-freshness.ts +++ /dev/null @@ -1,10 +0,0 @@ -let badgeFreshnessEpoch = 0; - -export function readBadgeFreshnessEpoch(): number { - return badgeFreshnessEpoch; -} - -export function advanceBadgeFreshnessEpoch(): number { - badgeFreshnessEpoch += 1; - return badgeFreshnessEpoch; -} diff --git a/apps/mobile/src/lib/badge-hydration.ts b/apps/mobile/src/lib/badge-hydration.ts deleted file mode 100644 index 74cc9db6d5..0000000000 --- a/apps/mobile/src/lib/badge-hydration.ts +++ /dev/null @@ -1,26 +0,0 @@ -import { type BadgeCountRow } from '@kilocode/notifications'; - -type ReconcileHydratedBadgeCountInput = { - badgeRows: BadgeCountRow[]; - startBadgeFreshnessEpoch: number; - currentBadgeFreshnessEpoch: number; - setBadgeCount: (badgeCount: number) => Promise; -}; - -export function totalBadgeCount(badgeRows: BadgeCountRow[]): number { - return badgeRows.reduce((total, row) => total + row.badgeCount, 0); -} - -export function reconcileHydratedBadgeCount({ - badgeRows, - startBadgeFreshnessEpoch, - currentBadgeFreshnessEpoch, - setBadgeCount, -}: ReconcileHydratedBadgeCountInput): boolean { - if (currentBadgeFreshnessEpoch !== startBadgeFreshnessEpoch) { - return false; - } - - void setBadgeCount(totalBadgeCount(badgeRows)); - return true; -} diff --git a/apps/mobile/src/lib/glanceable/publisher.test.ts b/apps/mobile/src/lib/glanceable/publisher.test.ts index 0845573591..34004ba3e7 100644 --- a/apps/mobile/src/lib/glanceable/publisher.test.ts +++ b/apps/mobile/src/lib/glanceable/publisher.test.ts @@ -1,3 +1,4 @@ +/* eslint-disable max-lines -- one cohesive publisher state-machine suite sharing the fake-sink harness */ import { afterEach, describe, expect, it, vi } from 'vitest'; import { @@ -99,7 +100,7 @@ describe('GlanceablePublisher', () => { publisher.dispose(); }); - it('coalesces later happy updates and emits only the latest', () => { + it('coalesces later happy updates but publishes needs-input changes immediately', () => { vi.useFakeTimers(); const { sink, calls } = makeSink(); const publisher = new GlanceablePublisher({ sinks: [sink], now: () => NOW, coalesceMs: 1000 }); @@ -110,6 +111,11 @@ describe('GlanceablePublisher', () => { vi.advanceTimersByTime(1000); expect(count(calls, 'startOrUpdate')).toBe(2); expect(lastSnapshot(calls, 'startOrUpdate').running).toBe(3); + // The badge reads `needsInput`, so a change to it skips the coalesce wait. + publisher.handleSessions([{ status: 'permission' }], PUB_CTX); + expect(lastSnapshot(calls, 'startOrUpdate').needsInput).toBe(1); + publisher.handleSessions([{ status: 'busy' }], PUB_CTX); + expect(lastSnapshot(calls, 'startOrUpdate').needsInput).toBe(0); publisher.dispose(); }); diff --git a/apps/mobile/src/lib/glanceable/publisher.ts b/apps/mobile/src/lib/glanceable/publisher.ts index 189a65f287..9bae4a8188 100644 --- a/apps/mobile/src/lib/glanceable/publisher.ts +++ b/apps/mobile/src/lib/glanceable/publisher.ts @@ -132,6 +132,10 @@ export class GlanceablePublisher { // First eligible emit starts the activity immediately, no coalesce wait. this.emit(snapshot, ctx); this.activityStarted = true; + } else if (snapshot.needsInput !== this.current?.needsInput) { + // Badge changes are actionable and must reach the launcher immediately. + this.cancelCoalesce(); + this.emit(snapshot, ctx); } else { this.scheduleCoalesced(snapshot, ctx); } diff --git a/apps/mobile/src/lib/hooks/use-current-user-id.test.ts b/apps/mobile/src/lib/hooks/use-current-user-id.test.ts new file mode 100644 index 0000000000..b828fae1d7 --- /dev/null +++ b/apps/mobile/src/lib/hooks/use-current-user-id.test.ts @@ -0,0 +1,54 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { useCurrentUserId } from './use-current-user-id'; + +const query = vi.hoisted(() => ({ + data: undefined as { id: string; email: string } | undefined, + isLoading: false, + isError: false, + isFetched: false, + refetch: vi.fn(), +})); + +vi.mock('@tanstack/react-query', () => ({ useQuery: () => query })); +vi.mock('@/lib/trpc', () => ({ + useTRPC: () => ({ user: { getMe: { queryOptions: () => ({}) } } }), +})); + +describe('useCurrentUserId', () => { + beforeEach(() => { + Object.assign(query, { + data: undefined, + isLoading: false, + isError: false, + isFetched: false, + }); + }); + + it('keeps the error state visible while a failed request retries', () => { + Object.assign(query, { isLoading: true, isFetched: true }); + + expect(useCurrentUserId().isError).toBe(true); + }); + + it('does not turn the initial loading state into an error', () => { + Object.assign(query, { isLoading: true }); + + expect(useCurrentUserId().isError).toBe(false); + }); + + it('keeps a settled request error visible', () => { + Object.assign(query, { isError: true, isFetched: true }); + + expect(useCurrentUserId().isError).toBe(true); + }); + + it('clears the error state after a successful retry', () => { + Object.assign(query, { + data: { id: 'user-1', email: 'user@example.com' }, + isFetched: true, + }); + + expect(useCurrentUserId().isError).toBe(false); + }); +}); diff --git a/apps/mobile/src/lib/hooks/use-current-user-id.ts b/apps/mobile/src/lib/hooks/use-current-user-id.ts index 70bcac3341..8e8472704a 100644 --- a/apps/mobile/src/lib/hooks/use-current-user-id.ts +++ b/apps/mobile/src/lib/hooks/use-current-user-id.ts @@ -8,7 +8,7 @@ type UseCurrentUserIdOptions = { export function useCurrentUserId(options: UseCurrentUserIdOptions = {}) { const trpc = useTRPC(); - const { data, isLoading, isError, refetch } = useQuery({ + const { data, isLoading, isError, isFetched, refetch } = useQuery({ ...trpc.user.getMe.queryOptions(), enabled: options.enabled ?? true, }); @@ -17,7 +17,7 @@ export function useCurrentUserId(options: UseCurrentUserIdOptions = {}) { userId: data?.id, email: data?.email, isLoading, - isError, + isError: isError || (isLoading && isFetched), refetch: () => { void refetch(); }, diff --git a/apps/mobile/src/lib/hooks/use-unread-counts-invalidation.ts b/apps/mobile/src/lib/hooks/use-unread-counts-invalidation.ts index 6570edcbb7..449bd86c32 100644 --- a/apps/mobile/src/lib/hooks/use-unread-counts-invalidation.ts +++ b/apps/mobile/src/lib/hooks/use-unread-counts-invalidation.ts @@ -4,7 +4,6 @@ import { useEffect } from 'react'; import { AppState } from 'react-native'; import { useCurrentUserId } from '@/components/kilo-chat/hooks/use-current-user-id'; -import { advanceBadgeFreshnessEpoch } from '@/lib/badge-freshness'; import { parseNotificationData } from '@/lib/notifications'; /** @@ -30,7 +29,6 @@ export function useUnreadCountsInvalidation() { } const invalidate = () => { - advanceBadgeFreshnessEpoch(); void queryClient.invalidateQueries({ queryKey: ['badges', userId], }); diff --git a/apps/mobile/src/lib/hooks/use-unread-counts.ts b/apps/mobile/src/lib/hooks/use-unread-counts.ts index 0a3fc7b092..40cbe83571 100644 --- a/apps/mobile/src/lib/hooks/use-unread-counts.ts +++ b/apps/mobile/src/lib/hooks/use-unread-counts.ts @@ -1,5 +1,4 @@ import { useQuery } from '@tanstack/react-query'; -import * as Notifications from 'expo-notifications'; import { useMemo } from 'react'; import { @@ -10,8 +9,6 @@ import { import { useCurrentUserId } from '@/components/kilo-chat/hooks/use-current-user-id'; import { useKiloChatTokenGetter } from '@/components/kilo-chat/hooks/use-kilo-chat-token'; -import { readBadgeFreshnessEpoch } from '@/lib/badge-freshness'; -import { reconcileHydratedBadgeCount } from '@/lib/badge-hydration'; import { NOTIFICATIONS_URL } from '@/lib/config'; /** @@ -34,7 +31,6 @@ export function useUnreadCounts() { enabled: userId !== null, staleTime: 30_000, queryFn: async () => { - const startBadgeFreshnessEpoch = readBadgeFreshnessEpoch(); const token = await getToken(); const response = await fetch(`${NOTIFICATIONS_URL}/v1/badges`, { headers: { Authorization: `Bearer ${token}` }, @@ -43,12 +39,6 @@ export function useUnreadCounts() { throw new Error(`Failed to fetch badges: ${response.status}`); } const body = listBadgesResponseSchema.parse(await response.json()); - reconcileHydratedBadgeCount({ - badgeRows: body.buckets, - startBadgeFreshnessEpoch, - currentBadgeFreshnessEpoch: readBadgeFreshnessEpoch(), - setBadgeCount: Notifications.setBadgeCountAsync, - }); return body.buckets; }, }); diff --git a/apps/mobile/src/lib/notifications.test.ts b/apps/mobile/src/lib/notifications.test.ts index ed2c83d4ba..3493334f44 100644 --- a/apps/mobile/src/lib/notifications.test.ts +++ b/apps/mobile/src/lib/notifications.test.ts @@ -34,7 +34,9 @@ type ResponseListener = (response: Response) => void; const mocks = vi.hoisted(() => ({ platform: { OS: 'android' as string }, + setBadgeCountAsync: vi.fn(), setNotificationChannelAsync: vi.fn(), + setNotificationHandler: vi.fn(), getPermissionsAsync: vi.fn(), requestPermissionsAsync: vi.fn(), getExpoPushTokenAsync: vi.fn(), @@ -49,6 +51,7 @@ const mocks = vi.hoisted(() => ({ startTokenListeners: new Set<(event: { activityPushToStartToken: string }) => void>(), registerActivityToken: vi.fn(), unregisterActivityToken: vi.fn(), + refreshActiveSessionsFromPush: vi.fn(), defineTask: vi.fn(), registerTaskAsync: vi.fn(), captureEvent: vi.fn(), @@ -59,11 +62,12 @@ vi.mock('react-native', () => ({ })); vi.mock('expo-notifications', () => ({ + setBadgeCountAsync: mocks.setBadgeCountAsync, setNotificationChannelAsync: mocks.setNotificationChannelAsync, getPermissionsAsync: mocks.getPermissionsAsync, requestPermissionsAsync: mocks.requestPermissionsAsync, getExpoPushTokenAsync: mocks.getExpoPushTokenAsync, - setNotificationHandler: vi.fn(), + setNotificationHandler: mocks.setNotificationHandler, addNotificationResponseReceivedListener: (listener: ResponseListener) => { mocks.listeners.add(listener); return { remove: () => mocks.listeners.delete(listener) }; @@ -139,6 +143,9 @@ vi.mock('@/lib/trpc', () => ({ }, })); vi.mock('@/lib/query-client', () => ({ queryClient: {} })); +vi.mock('@/lib/active-sessions-live-sync', () => ({ + refreshActiveSessionsFromPush: mocks.refreshActiveSessionsFromPush, +})); vi.mock('@/lib/persist/read-cache', () => ({ readCachedUserId: () => null })); vi.mock('@kilocode/notifications', async importOriginal => ({ @@ -167,7 +174,12 @@ async function loadNotifications() { deleteItemAsync: vi.fn().mockResolvedValue(undefined), getItemAsync: vi.fn().mockResolvedValue(null), }); - return { ...(await import('./notifications')), pending }; + const [notifications, registry, persist] = await Promise.all([ + import('./notifications'), + import('@/lib/glanceable/sink-registry'), + import('@/lib/glanceable/persist'), + ]); + return { ...notifications, pending, persist, registry }; } function deferred(): { promise: Promise; resolve: () => void } { @@ -192,6 +204,7 @@ async function flushMicrotasks(): Promise { beforeEach(() => { vi.clearAllMocks(); mocks.platform.OS = 'android'; + mocks.setBadgeCountAsync.mockResolvedValue(true); mocks.setNotificationChannelAsync.mockResolvedValue(undefined); mocks.getPermissionsAsync.mockResolvedValue({ status: 'denied' }); mocks.requestPermissionsAsync.mockResolvedValue({ status: 'denied' }); @@ -513,6 +526,252 @@ const secureStoreMock = { }), }; +describe('glanceable app badge sink', () => { + async function loadBadgeSink() { + const loaded = await loadNotifications(); + loaded._setGlanceableSinksLoaderForTests(() => undefined); + loaded.setupNotificationBackgroundHandler(); + const [sink] = loaded.registry.getGlanceableSinks(); + if (!sink) { + throw new Error('The app badge sink was not registered'); + } + return { loaded, sink }; + } + + it('sets the needs-input count for a local happy snapshot', async () => { + const { sink } = await loadBadgeSink(); + + sink.publish(glanceableSnapshot({ needsInput: 3 })); + await flushMicrotasks(); + + expect(mocks.setBadgeCountAsync).toHaveBeenCalledWith(3); + }); + + it('keeps the stale count until a successful retry publishes a replacement', async () => { + const { sink } = await loadBadgeSink(); + + sink.publish(glanceableSnapshot({ needsInput: 3 })); + sink.publish(glanceableSnapshot({ status: 'stale', needsInput: 3 })); + sink.publish(glanceableSnapshot({ revision: 2, needsInput: 1 })); + await flushMicrotasks(); + + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[3], [3], [1]]); + }); + + it('leaves a push-set launcher badge alone across re-publications of the same count', async () => { + const { sink } = await loadBadgeSink(); + + // e7 baseline: a needs-input session published count 1 to the launcher. + sink.publish(glanceableSnapshot({ needsInput: 1 })); + await flushMicrotasks(); + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[1]]); + + // The OS then moves the launcher badge itself: a visible foreground push + // with badge 2 is applied through the shouldSetBadge: true path. Every + // tray refresh re-publishes the unchanged snapshot afterwards; none of + // those re-publications may undo the badge the push carried. + sink.publish(glanceableSnapshot({ revision: 2, needsInput: 1 })); + sink.publish(glanceableSnapshot({ revision: 3, status: 'stale', needsInput: 1 })); + await flushMicrotasks(); + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[1]]); + + // A real count change still reaches the launcher. + sink.publish(glanceableSnapshot({ revision: 4, needsInput: 2 })); + await flushMicrotasks(); + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[1], [2]]); + }); + + it.each([ + ['busy-only', { status: 'happy', running: 2, needsInput: 0 }], + ['empty', { status: 'empty', running: 0, needsInput: 0, needsInputSince: null }], + ['waiting', { status: 'waiting', running: 0, needsInput: 0, needsInputSince: null }], + ['signed-out', { status: 'signed_out', running: 0, needsInput: 0, needsInputSince: null }], + ['privacy', { status: 'privacy', running: 0, needsInput: 0, needsInputSince: null }], + ] as const)('clears the badge for a %s snapshot', async (_label, overrides) => { + const { sink } = await loadBadgeSink(); + + sink.publish(glanceableSnapshot(overrides)); + await flushMicrotasks(); + + expect(mocks.setBadgeCountAsync).toHaveBeenCalledWith(0); + }); + + it('serializes native writes so the newest count finishes last', async () => { + const gate = deferred(); + mocks.setBadgeCountAsync + .mockImplementationOnce(async () => { + await gate.promise; + return true; + }) + .mockResolvedValue(true); + const { sink } = await loadBadgeSink(); + + sink.publish(glanceableSnapshot({ needsInput: 2 })); + sink.publish(glanceableSnapshot({ revision: 2, needsInput: 7 })); + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[2]]); + + gate.resolve(); + await flushMicrotasks(); + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[2], [7]]); + }); + + it('captures a failed write and continues with the next count', async () => { + const error = new Error('badge write failed'); + mocks.setBadgeCountAsync.mockRejectedValueOnce(error).mockResolvedValue(true); + const { sink } = await loadBadgeSink(); + + sink.publish(glanceableSnapshot({ needsInput: 2 })); + sink.publish(glanceableSnapshot({ revision: 2, needsInput: 5 })); + await flushMicrotasks(); + + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[2], [5]]); + expect(mocks.captureException).toHaveBeenCalledWith(error, { + tags: { + 'error.subsystem': 'notifications', + 'error.operation': 'set_glanceable_badge', + }, + }); + }); + + it('applies foreground glanceable counts and allows visible push badges', async () => { + const { loaded } = await loadBadgeSink(); + loaded.persist._setLastGlanceableSnapshotForTests(glanceableSnapshot({ needsInput: 2 })); + mockSecureStoreKeys(); + loaded.setupNotificationHandler(); + const registration = mocks.setNotificationHandler.mock.calls[0]?.[0] as { + handleNotification: (notification: { + request: { content: { data: unknown } }; + }) => Promise<{ shouldSetBadge: boolean }>; + }; + + const ordinary = await registration.handleNotification({ + request: { + content: { + data: { + type: 'chat.message', + sandboxId: 'sandbox-1', + conversationId: 'conversation-1', + messageId: 'message-1', + }, + }, + }, + }); + expect(ordinary.shouldSetBadge).toBe(true); + expect(mocks.setBadgeCountAsync).not.toHaveBeenCalled(); + + const stale = await registration.handleNotification({ + request: { + content: { + data: activeGlanceablePush({ + updatedAt: '2025-12-31T00:00:00.000Z', + needsInput: 9, + }), + }, + }, + }); + expect(stale.shouldSetBadge).toBe(false); + expect(mocks.setBadgeCountAsync).not.toHaveBeenCalled(); + await flushMicrotasks(); + expect(mocks.setBadgeCountAsync).toHaveBeenCalledWith(2); + mocks.setBadgeCountAsync.mockClear(); + + const badgeWrite = deferred(); + mocks.setBadgeCountAsync.mockImplementation(async () => { + await badgeWrite.promise; + return true; + }); + const handling = registration.handleNotification({ + request: { + content: { + data: activeGlanceablePush({ + updatedAt: '2026-01-02T00:00:00.000Z', + needsInput: 4, + }), + }, + }, + }); + await flushMicrotasks(); + const completedBeforeWrite = await Promise.race([handling, Promise.resolve(null)]); + badgeWrite.resolve(); + const glanceable = await handling; + + expect(glanceable.shouldSetBadge).toBe(true); + expect(mocks.setBadgeCountAsync).toHaveBeenCalledWith(4); + expect(mocks.refreshActiveSessionsFromPush).toHaveBeenCalledOnce(); + expect(completedBeforeWrite).toBeNull(); + }); + + it('does not re-assert the local count when a later glanceable push is discarded', async () => { + const { loaded } = await loadBadgeSink(); + loaded.persist._setLastGlanceableSnapshotForTests(glanceableSnapshot({ needsInput: 2 })); + mockSecureStoreKeys(); + loaded.setupNotificationHandler(); + const registration = mocks.setNotificationHandler.mock.calls[0]?.[0] as { + handleNotification: (notification: { + request: { content: { data: unknown } }; + }) => Promise<{ shouldSetBadge: boolean }>; + }; + + // A fresh glanceable push confirms count 2 on the launcher through the sink. + await registration.handleNotification({ + request: { + content: { + data: activeGlanceablePush({ + updatedAt: '2026-01-02T00:00:00.000Z', + needsInput: 2, + }), + }, + }, + }); + await flushMicrotasks(); + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[2]]); + + // A visible push then moves the launcher itself (badge 3, shouldSetBadge). + // A stale glanceable push is discarded: it carries no new count, so the + // fallback must NOT re-write the unchanged local truth — that would undo + // the badge the visible push carried (the e7 clobber). + await registration.handleNotification({ + request: { + content: { + data: { + type: 'chat.message', + sandboxId: 'sandbox-1', + conversationId: 'conversation-1', + messageId: 'message-1', + }, + }, + }, + }); + const stale = await registration.handleNotification({ + request: { + content: { + data: activeGlanceablePush({ + updatedAt: '2025-12-31T00:00:00.000Z', + needsInput: 9, + }), + }, + }, + }); + expect(stale.shouldSetBadge).toBe(false); + await flushMicrotasks(); + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[2]]); + + // A real count change still reaches the launcher after the discard. + await registration.handleNotification({ + request: { + content: { + data: activeGlanceablePush({ + updatedAt: '2026-01-03T00:00:00.000Z', + needsInput: 3, + }), + }, + }, + }); + await flushMicrotasks(); + expect(mocks.setBadgeCountAsync.mock.calls).toEqual([[2], [3]]); + }); +}); + describe('applyGlanceablePushData', () => { beforeEach(() => { _resetGlanceablePersistForTests(); @@ -947,6 +1206,7 @@ describe('setupNotificationBackgroundHandler', () => { scopeKey: SCOPE_KEY, updatedAt: '2026-01-02T00:00:00.000Z', organizationBound: true, + needsInput: 6, }) ), }, @@ -965,10 +1225,67 @@ describe('setupNotificationBackgroundHandler', () => { userId: 'u1', organizationId: 'org-9', }); + expect(mocks.setBadgeCountAsync).toHaveBeenCalledWith(6); unregisterGlanceableSink(sink); }); + it('keeps the background task alive until a zero-count badge write finishes', async () => { + const persisted = glanceableSnapshot({ + scopeKey: SCOPE_KEY, + revision: 1, + updatedAt: '2026-01-01T00:00:00.000Z', + needsInput: 1, + }); + secureStore.set('glanceable-snapshot', JSON.stringify(persisted)); + secureStore.set('glanceable-scope-key', SCOPE_KEY); + _setGlanceableSinksLoaderForTests(() => undefined); + const badgeWrite = deferred(); + mocks.setBadgeCountAsync.mockImplementation(async () => { + await badgeWrite.promise; + return true; + }); + + setupNotificationBackgroundHandler(); + const executor = executorFor(mocks.defineTask); + let completed = false; + const apply = async () => { + const result = await executor({ + data: { + notification: null, + data: { + dataString: JSON.stringify( + activeGlanceablePush({ + scopeKey: SCOPE_KEY, + updatedAt: '2026-01-02T00:00:00.000Z', + status: 'empty', + running: 0, + needsInput: 0, + idle: 0, + needsInputSince: null, + }) + ), + }, + }, + error: null, + executionInfo: { + eventId: 'e-clear', + taskName: 'active-agents-glanceable-background-task', + }, + }); + completed = true; + return result; + }; + const applying = apply(); + await flushMicrotasks(); + + expect(mocks.setBadgeCountAsync).toHaveBeenCalledWith(0); + expect(completed).toBe(false); + + badgeWrite.resolve(); + expect(await applying).toBe(0); + }); + it('ignores a headless payload that is not a glanceable push', async () => { _setGlanceableSinksLoaderForTests(() => undefined); diff --git a/apps/mobile/src/lib/notifications.ts b/apps/mobile/src/lib/notifications.ts index 8abc51d8b6..4f44bb8c98 100644 --- a/apps/mobile/src/lib/notifications.ts +++ b/apps/mobile/src/lib/notifications.ts @@ -25,6 +25,7 @@ import { } from '@kilocode/app-shared/glanceable-agents-snapshot'; import { captureEvent } from '@/lib/analytics/posthog'; +import { refreshActiveSessionsFromPush } from '@/lib/active-sessions-live-sync'; import { currentAuthEpoch } from '@/lib/auth/auth-epoch'; import { getTerminalBlankEpoch } from '@/lib/glanceable/cleanup'; import { @@ -33,7 +34,12 @@ import { persistGlanceableSink, restorePersistedGlanceable, } from '@/lib/glanceable/persist'; -import { getGlanceableSinks, registerGlanceableSink } from '@/lib/glanceable/sink-registry'; +import { + getGlanceableSinks, + type GlanceableSink, + registerGlanceableSink, +} from '@/lib/glanceable/sink-registry'; +import { chainSave } from '@/lib/hooks/save-chain'; import { ACTIVE_USER_ID_KEY, ORGANIZATION_STORAGE_KEY } from '@/lib/storage-keys'; import { i18n } from '@/i18n'; import { setPendingDeepLink } from './deep-link-launch'; @@ -56,6 +62,57 @@ function getProjectId(): string { // is registered once and must always read the latest value without stale closures. let activeChatLocation: { sandboxId: string; conversationId: string } | null = null; +let appBadgeWrite: Promise | null = null; + +// The count the last successful native write put on the launcher badge. The +// badge is shared with iOS itself: a visible push that carries its own badge +// moves the icon through the shouldSetBadge path without the app writing. +let badgeWrittenCount: number | null = null; + +async function setAppBadge(count: number): Promise { + try { + await chainSave('glanceable-app-badge', async () => { + await Notifications.setBadgeCountAsync(count); + }); + badgeWrittenCount = count; + } catch (error) { + Sentry.captureException(error, { + tags: { + 'error.subsystem': 'notifications', + 'error.operation': 'set_glanceable_badge', + }, + }); + } +} + +function syncAppBadge(count: number): void { + // Only a change of the needs-input count owns a badge write. Re-asserting + // an unchanged count would undo a badge iOS applied from a visible push + // payload (e7: foreground push badge 2 was clobbered back to needsInput 1 + // before the app was terminated). + if (count === badgeWrittenCount) { + return; + } + appBadgeWrite = setAppBadge(count); +} + +const appBadgeSink: GlanceableSink = { + publish(snapshot) { + syncAppBadge(snapshot.needsInput); + }, + async waitForNativeTerminal() { + if (appBadgeWrite) { + await appBadgeWrite; + } + }, + endImmediate() { + // A terminal snapshot already published zero. + }, + startOrUpdate() { + // The publish operation owns every badge write. + }, +}; + export function setActiveChatLocation( location: { sandboxId: string; conversationId: string } | null ) { @@ -179,6 +236,9 @@ export async function applyGlanceablePushData( // Keep the existing fallback for other sinks; widgets retain their timeline. scheduleGlanceableTerminalEnd(); } + if (appBadgeWrite) { + await appBadgeWrite; + } // Do not finish a background task before ActivityKit accepts the native end. // All publication happens before this await, so it cannot restore an old scope. if (!eligible) { @@ -237,8 +297,20 @@ export function setupNotificationHandler() { // The aggregate glanceable push is a data carrier for the ongoing // notification/widgets, never a visible banner: the local ongoing owns // the display. Apply it to the sinks regardless of the discard outcome. - await applyGlanceablePushData(data); - return suppressed; + const applied = await applyGlanceablePushData(data); + if (applied) { + refreshActiveSessionsFromPush(); + } + if (!applied) { + setTimeout(() => { + // A discarded push carries no new count, so the local truth is + // unchanged. Re-asserting it here would undo a badge iOS applied + // from a later visible push (the same e7 clobber as an unchanged + // re-publication), so it goes through the guarded write too. + syncAppBadge(getLastGlanceableSnapshot()?.needsInput ?? 0); + }, 0); + } + return { ...suppressed, shouldSetBadge: applied }; } if ( @@ -274,6 +346,7 @@ export function _setGlanceableSinksLoaderForTests(loader: (() => void) | null): * sinks must be registered here before `applyGlanceablePushData` runs. */ function ensureGlanceableSinksLoaded(): void { + registerGlanceableSink(appBadgeSink); if (glanceableSinksLoaderForTests) { glanceableSinksLoaderForTests(); return; diff --git a/services/notifications/src/__tests__/dispatch-push.test.ts b/services/notifications/src/__tests__/dispatch-push.test.ts index b1bcf81842..fcc39c9367 100644 --- a/services/notifications/src/__tests__/dispatch-push.test.ts +++ b/services/notifications/src/__tests__/dispatch-push.test.ts @@ -85,6 +85,12 @@ function getDO(name = 'user-1') { return env.NOTIFICATION_CHANNEL_DO.get(id); } +afterEach(() => { + for (const [messages] of vi.mocked(sendPushNotifications).mock.calls) { + for (const message of messages) expect(message).not.toHaveProperty('badge'); + } +}); + describe('NotificationChannelDO.dispatchPush', () => { beforeEach(() => { vi.clearAllMocks(); @@ -241,8 +247,6 @@ describe('NotificationChannelDO.dispatchPush', () => { expect(result.kind).toBe('delivered'); expect(sendPushNotifications).toHaveBeenCalledOnce(); - const [[messages]] = vi.mocked(sendPushNotifications).mock.calls; - expect(messages[0].badge).toBe(1); // Bucket persisted to DO storage. const stored = await runInDurableObject(stub, async (_inst, state) => ({ @@ -276,8 +280,6 @@ describe('NotificationChannelDO.dispatchPush', () => { expect(result.outcome.kind).toBe('delivered'); expect(result.totalReads).toBe(1); - const [[messages]] = vi.mocked(sendPushNotifications).mock.calls; - expect(messages[0].badge).toBe(1); }); it('returns failed and avoids delivered idempotency for non-stale Expo ticket errors', async () => { @@ -445,7 +447,6 @@ describe('NotificationChannelDO.dispatchPush', () => { expect(sendPushNotifications).toHaveBeenCalledOnce(); const [[messages]] = vi.mocked(sendPushNotifications).mock.calls; expect(messages.map(message => message.to)).toEqual(['tok-accepted', 'tok-rate-limited']); - expect(messages.map(message => message.badge)).toEqual([1, 1]); expect(receiptSpy).toHaveBeenCalledWith( { ticketTokenPairs: [{ ticketId: 'ticket-accepted', token: 'tok-accepted' }] }, { delaySeconds: 900 } @@ -489,10 +490,6 @@ describe('NotificationChannelDO.dispatchPush', () => { }); expect(second).toEqual({ kind: 'delivered', tokenCount: 1 }); expect(sendPushNotifications).toHaveBeenCalledTimes(2); - const [firstMessages] = vi.mocked(sendPushNotifications).mock.calls[0]; - const [secondMessages] = vi.mocked(sendPushNotifications).mock.calls[1]; - expect(firstMessages[0].badge).toBe(1); - expect(secondMessages[0].badge).toBe(1); const stored = await runInDurableObject(stub, async (_inst, state) => state.storage.get<{ stage: string; ts: number }>('idem:k-retry-ticket-error') @@ -500,7 +497,7 @@ describe('NotificationChannelDO.dispatchPush', () => { expect(stored).toMatchObject({ stage: 'delivered' }); }); - it('accumulates bucket counts across deliveries and exposes total via badge', async () => { + it('accumulates bucket counts across deliveries', async () => { installDbMock({ tokens: [{ user_id: 'u', token: 'tok1' }] }); vi.spyOn(env.EVENT_SERVICE, 'isUserInContext').mockResolvedValue(false); const stub = getDO('user-accumulate'); @@ -514,11 +511,6 @@ describe('NotificationChannelDO.dispatchPush', () => { }) ); - const calls = vi.mocked(sendPushNotifications).mock.calls; - expect(calls[0]?.[0][0].badge).toBe(1); - expect(calls[1]?.[0][0].badge).toBe(2); - expect(calls[2]?.[0][0].badge).toBe(3); - const buckets = await runInDurableObject(stub, async (_inst, state) => { const entries = await state.storage.list({ prefix: 'bucket:' }); return Array.from(entries.entries()); @@ -678,8 +670,20 @@ describe('NotificationChannelDO.dispatchPush', () => { ); expect(afterFail).toBe(1); - const second = await stub.dispatchPush(input); - expect(second.kind).toBe('delivered'); + const retry = await runInDurableObject(stub, async (instance, state) => { + const originalGet = state.storage.get.bind(state.storage); + let totalReads = 0; + state.storage.get = ((key: string | string[]) => { + if (Array.isArray(key)) return originalGet(key); + if (key === 'total') totalReads++; + return originalGet(key); + }) as typeof state.storage.get; + + const outcome = await instance.dispatchPush(input); + return { outcome, totalReads }; + }); + expect(retry.outcome.kind).toBe('delivered'); + expect(retry.totalReads).toBe(0); // Bucket must not be incremented twice across the retry — the first // attempt's `pending` marker gates the second increment out. @@ -687,9 +691,6 @@ describe('NotificationChannelDO.dispatchPush', () => { state.storage.get('bucket:conv1') ); expect(afterRetry).toBe(1); - - const [[messages]] = vi.mocked(sendPushNotifications).mock.calls; - expect(messages[0].badge).toBe(1); }); it('schedules cleanup when writing the pending marker (failed send)', async () => { diff --git a/services/notifications/src/dos/NotificationChannelDO.ts b/services/notifications/src/dos/NotificationChannelDO.ts index b1efc094d7..8ac957b1a9 100644 --- a/services/notifications/src/dos/NotificationChannelDO.ts +++ b/services/notifications/src/dos/NotificationChannelDO.ts @@ -159,22 +159,15 @@ export class NotificationChannelDO extends DurableObject { // 4. Badge math. On a retry the badge was already incremented during // the prior attempt; re-applying the delta would double-count. - // The total is recomputed in either case (other writers may have - // advanced it). - let badgeTotal: number | undefined; - if (input.badge) { - if (!isRetry) { - // Mark `pending` BEFORE the increment so any later failure path - // is gated on the marker and a retry skips the increment. - const ts = Date.now(); - await this.ctx.storage.put(idemKey, { stage: 'pending', ts }); - // Also schedule cleanup at this point — if Expo keeps failing and - // no future push ever lands, `pending` would otherwise leak. - await this.ensureCleanupAlarm(ts); - badgeTotal = await this.incrementBucket(input.badge.badgeBucket, input.badge.delta); - } else { - badgeTotal = await this.getTotal(); - } + if (input.badge && !isRetry) { + // Mark `pending` BEFORE the increment so any later failure path + // is gated on the marker and a retry skips the increment. + const ts = Date.now(); + await this.ctx.storage.put(idemKey, { stage: 'pending', ts }); + // Also schedule cleanup at this point — if Expo keeps failing and + // no future push ever lands, `pending` would otherwise leak. + await this.ensureCleanupAlarm(ts); + await this.incrementBucket(input.badge.badgeBucket, input.badge.delta); } // 5. Tokens. Missing Expo tokens only means no OS push can be sent; the @@ -288,7 +281,6 @@ export class NotificationChannelDO extends DurableObject { // version at registration) get a channelId; older clients fall back // to the default channel. iOS ignores channelId either way. ...(app_version != null && { channelId }), - ...(badgeTotal !== undefined && { badge: badgeTotal }), sound: input.push.sound ?? undefined, priority: input.push.priority ?? 'default', } satisfies ExpoPushMessage; @@ -550,7 +542,7 @@ export class NotificationChannelDO extends DurableObject { // Read-modify-write of a bucket counter. The DO is single-threaded, so // this is race-free without explicit locking. - private async incrementBucket(bucket: string, delta: number): Promise { + private async incrementBucket(bucket: string, delta: number): Promise { const key = `${BUCKET_PREFIX}${bucket}`; const total = await this.getTotal(); const current = (await this.ctx.storage.get(key)) ?? 0; @@ -563,7 +555,6 @@ export class NotificationChannelDO extends DurableObject { const nextTotal = Math.max(0, total + delta); await this.ctx.storage.put(TOTAL_KEY, nextTotal); - return nextTotal; } // Aggregate badge count. Existing DOs without the aggregate fall back to one