feat: PR 7 — drop badge_counts table; per-user DO storage; JWT-authed badge endpoints

apps/mobile/.env

@@ -5,5 +5,6 @@ CLOUD_AGENT_WS_URL=wss://cloud-agent-next.kilosessions.ai
 SESSION_INGEST_WS_URL=wss://ingest.kilosessions.ai
 APPSFLYER_DEV_KEY=jnoVs6KzXanpbKrqXckPu9
 APPSFLYER_APP_ID=6761193135
-EXPO_PUBLIC_KILO_CHAT_URL=https://kilo-chat.kilosessions.ai
-EXPO_PUBLIC_EVENT_SERVICE_URL=wss://event-service.kilosessions.ai
+KILO_CHAT_URL=https://chat.kiloapps.io
+EVENT_SERVICE_URL=wss://events.kiloapps.io
+NOTIFICATIONS_URL=https://notifications.kiloapps.io

apps/mobile/src/components/kilo-chat/hooks/use-mark-read.ts

@@ -1,28 +1,78 @@
 import { useCallback } from 'react';
-import { useMutation } from '@tanstack/react-query';
+import { useMutation, useQueryClient } from '@tanstack/react-query';
 import * as Notifications from 'expo-notifications';
+import { toast } from 'sonner-native';
 
 import { badgeBucketForConversation } from '@kilocode/notifications';
 
-import { useTRPC } from '@/lib/trpc';
+import { NOTIFICATIONS_URL } from '@/lib/config';
+
+import { useCurrentUserId } from './use-current-user-id';
+import { useKiloChatTokenGetter } from './use-kilo-chat-token';
+
+type BadgeBucket = { badgeBucket: string; badgeCount: number };
+type MarkReadContext = {
+  previousBadges?: BadgeBucket[];
+  queryKey?: readonly ['badges', string];
+};
 
 export function useMarkRead() {
-  const trpc = useTRPC();
-  const mutation = useMutation(
-    trpc.user.markChatRead.mutationOptions({
-      onSuccess: result => {
-        if (typeof result.badgeCount === 'number') {
-          void Notifications.setBadgeCountAsync(result.badgeCount);
-        }
-      },
-    })
-  );
+  const queryClient = useQueryClient();
+  const userId = useCurrentUserId();
+  const getToken = useKiloChatTokenGetter();
+
+  const mutation = useMutation({
+    mutationFn: async (badgeBucket: string): Promise<{ badgeCount: number }> => {
+      const token = await getToken();
+      const response = await fetch(`${NOTIFICATIONS_URL}/v1/badges/mark-read`, {
+        method: 'POST',
+        headers: {
+          Authorization: `Bearer ${token}`,
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify({ badgeBucket }),
+      });
+      if (!response.ok) {
+        throw new Error(`Failed to mark badge read: ${response.status}`);
+      }
+      return (await response.json()) as { badgeCount: number };
+    },
+    onMutate: async (badgeBucket): Promise<MarkReadContext> => {
+      if (userId === null) {
+        return {};
+      }
+
+      const queryKey = ['badges', userId] as const;
+      await queryClient.cancelQueries({ queryKey });
+      const previousBadges = queryClient.getQueryData<BadgeBucket[]>(queryKey);
+
+      queryClient.setQueryData<BadgeBucket[]>(queryKey, badges =>
+        badges?.filter(row => row.badgeBucket !== badgeBucket)
+      );
+
+      return { previousBadges, queryKey };
+    },
+    onError: (error, _badgeBucket, context) => {
+      if (context?.queryKey && context.previousBadges) {
+        queryClient.setQueryData(context.queryKey, context.previousBadges);
+      }
+      toast.error(error.message);
+    },
+    onSuccess: result => {
+      if (typeof result.badgeCount === 'number') {
+        void Notifications.setBadgeCountAsync(result.badgeCount);
+      }
+    },
+    onSettled: () => {
+      if (userId !== null) {
+        void queryClient.invalidateQueries({ queryKey: ['badges', userId] });
+      }
+    },
+  });
 
   return useCallback(
     (sandboxId: string, conversationId: string) => {
-      mutation.mutate({
-        badgeBucket: badgeBucketForConversation(sandboxId, conversationId),
-      });
+      mutation.mutate(badgeBucketForConversation(sandboxId, conversationId));
     },
     [mutation]
   );

apps/mobile/src/lib/config.ts

@@ -21,3 +21,4 @@ export const SESSION_INGEST_WS_URL: string = required('sessionIngestWsUrl');
 
 export const KILO_CHAT_URL: string = required('kiloChatUrl');
 export const EVENT_SERVICE_URL: string = required('eventServiceUrl');
+export const NOTIFICATIONS_URL: string = required('notificationsUrl');

apps/mobile/src/lib/env-keys.js

@@ -7,6 +7,7 @@ export const ENV_KEYS = {
   sessionIngestWsUrl: 'SESSION_INGEST_WS_URL',
   appsFlyerDevKey: 'APPSFLYER_DEV_KEY',
   appsFlyerAppId: 'APPSFLYER_APP_ID',
-  kiloChatUrl: 'EXPO_PUBLIC_KILO_CHAT_URL',
-  eventServiceUrl: 'EXPO_PUBLIC_EVENT_SERVICE_URL',
+  kiloChatUrl: 'KILO_CHAT_URL',
+  eventServiceUrl: 'EVENT_SERVICE_URL',
+  notificationsUrl: 'NOTIFICATIONS_URL',
 };

apps/mobile/src/lib/hooks/use-unread-counts-invalidation.ts

@@ -3,11 +3,11 @@ import * as Notifications from 'expo-notifications';
 import { useEffect } from 'react';
 import { AppState } from 'react-native';
 
+import { useCurrentUserId } from '@/components/kilo-chat/hooks/use-current-user-id';
 import { parseNotificationData } from '@/lib/notifications';
-import { useTRPC } from '@/lib/trpc';
 
 /**
- * Keeps the `user.getUnreadCounts` cache in sync with real-time notification
+ * Keeps the `['badges', userId]` cache in sync with real-time notification
  * traffic so per-instance badges on the dashboard reflect pushes received while
  * the app is open or resumed from background.
  *
@@ -20,14 +20,16 @@ import { useTRPC } from '@/lib/trpc';
  */
 export function useUnreadCountsInvalidation() {
   const queryClient = useQueryClient();
-  const trpc = useTRPC();
+  const userId = useCurrentUserId();
 
   useEffect(() => {
-    // `trpc` is stable (memoized inside TRPCProvider) but `queryKey()` returns
-    // a fresh array on each call, so we resolve it inside each invalidation.
+    if (userId === null) {
+      return undefined;
+    }
+
     const invalidate = () => {
       void queryClient.invalidateQueries({
-        queryKey: trpc.user.getUnreadCounts.queryKey(),
+        queryKey: ['badges', userId],
       });
     };
 
@@ -48,5 +50,5 @@ export function useUnreadCountsInvalidation() {
       received.remove();
       appStateSubscription.remove();
     };
-  }, [queryClient, trpc]);
+  }, [queryClient, userId]);
 }

apps/mobile/src/lib/hooks/use-unread-counts.ts

@@ -1,26 +1,44 @@
 import { useQuery } from '@tanstack/react-query';
 import { useMemo } from 'react';
 
+import { useCurrentUserId } from '@/components/kilo-chat/hooks/use-current-user-id';
+import { useKiloChatTokenGetter } from '@/components/kilo-chat/hooks/use-kilo-chat-token';
 import { badgeBucketForInstance } from '@/lib/badge-buckets';
-import { useTRPC } from '@/lib/trpc';
+import { NOTIFICATIONS_URL } from '@/lib/config';
+
+type Bucket = { badgeBucket: string; badgeCount: number };
 
 /**
- * Fetches unread message counts for the current user and returns a Map keyed by
- * instance badge bucket for O(1) lookup from dashboard cards. Conversation
- * buckets are summed into their parent instance bucket.
+ * Fetches unread message counts for the current user from the notifications
+ * worker and returns a Map keyed by instance badge bucket for O(1) lookup from
+ * dashboard cards. Conversation buckets are summed into their parent instance
+ * bucket.
  *
  * Freshness is driven by invalidations, not polling:
  *   - Foreground chat push → invalidate (see `use-unread-counts-invalidation`).
  *   - App returns to active → invalidate.
- *   - `markChatRead` optimistically clears the relevant row.
+ *   - `useMarkRead` optimistically clears the relevant row.
  */
 export function useUnreadCounts() {
-  const trpc = useTRPC();
-  const query = useQuery(
-    trpc.user.getUnreadCounts.queryOptions(undefined, {
-      staleTime: 30_000,
-    })
-  );
+  const userId = useCurrentUserId();
+  const getToken = useKiloChatTokenGetter();
+
+  const query = useQuery<Bucket[]>({
+    queryKey: ['badges', userId],
+    enabled: userId !== null,
+    staleTime: 30_000,
+    queryFn: async () => {
+      const token = await getToken();
+      const response = await fetch(`${NOTIFICATIONS_URL}/v1/badges`, {
+        headers: { Authorization: `Bearer ${token}` },
+      });
+      if (!response.ok) {
+        throw new Error(`Failed to fetch badges: ${response.status}`);
+      }
+      const body = (await response.json()) as { buckets: Bucket[] };
+      return body.buckets;
+    },
+  });
 
   const byBadgeBucket = useMemo(() => {
     const map = new Map<string, number>();

apps/web/src/routers/user-router.test.ts

@@ -1,7 +1,7 @@
 import { createCallerForUser } from '@/routers/test-utils';
 import { db } from '@/lib/drizzle';
-import { badge_counts, kilocode_users } from '@kilocode/db/schema';
-import { eq, inArray } from 'drizzle-orm';
+import { kilocode_users } from '@kilocode/db/schema';
+import { eq } from 'drizzle-orm';
 import { insertTestUser } from '@/tests/helpers/user.helper';
 import type { User } from '@kilocode/db/schema';
 
@@ -416,25 +416,3 @@ describe('session and API token reset mutations', () => {
     expect(updated.api_token_pepper).toBe('api-pepper-before');
   });
 });
-
-describe('user router - getUnreadCounts', () => {
-  it('does not return counts from other users', async () => {
-    const user = await insertTestUser({
-      google_user_email: `unread-counts-me-${crypto.randomUUID()}@example.com`,
-    });
-    const other = await insertTestUser({
-      google_user_email: `unread-counts-other-${crypto.randomUUID()}@example.com`,
-    });
-    await db.insert(badge_counts).values([
-      { user_id: user.id, badge_bucket: 'sandbox-mine', badge_count: 4 },
-      { user_id: other.id, badge_bucket: 'sandbox-theirs', badge_count: 9 },
-    ]);
-
-    const caller = await createCallerForUser(user.id);
-    const result = await caller.user.getUnreadCounts();
-
-    expect(result).toEqual([{ badgeBucket: 'sandbox-mine', badgeCount: 4 }]);
-
-    await db.delete(badge_counts).where(inArray(badge_counts.user_id, [user.id, other.id]));
-  });
-});

apps/web/src/routers/user-router.ts

@@ -20,9 +20,8 @@ import {
   kiloclaw_instances,
   kiloclaw_subscriptions,
   user_push_tokens,
-  badge_counts,
 } from '@kilocode/db/schema';
-import { eq, and, isNull, inArray, sql, gt, gte, sum } from 'drizzle-orm';
+import { eq, and, isNull, inArray, sql, gte } from 'drizzle-orm';
 import crypto from 'crypto';
 import { checkDiscordGuildMembership } from '@/lib/integrations/discord-guild-membership';
 import { AuthProviderIdSchema } from '@/lib/auth/provider-metadata';
@@ -722,45 +721,4 @@ export const userRouter = createTRPCRouter({
       .from(user_push_tokens)
       .where(eq(user_push_tokens.user_id, ctx.user.id));
   }),
-
-  // ─── Badge Counts ──────────────────────────────────────────────────
-
-  // Called by the mobile app when the user opens a chat. Resets the badge
-  // count for that channel to 0 and returns the new total across all
-  // channels, which the app applies as the OS badge count.
-  markChatRead: baseProcedure
-    .input(z.object({ badgeBucket: z.string().min(1) }))
-    .mutation(async ({ ctx, input }) => {
-      await db
-        .update(badge_counts)
-        .set({ badge_count: 0 })
-        .where(
-          and(
-            eq(badge_counts.user_id, ctx.user.id),
-            eq(badge_counts.badge_bucket, input.badgeBucket)
-          )
-        );
-
-      const [totals] = await db
-        .select({ total: sum(badge_counts.badge_count) })
-        .from(badge_counts)
-        .where(eq(badge_counts.user_id, ctx.user.id));
-
-      return { badgeCount: Number(totals?.total ?? 0) };
-    }),
-
-  // Per-channel unread counts for showing badges on the mobile dashboard. For
-  // kiloclaw chats, `channelId` equals `sandbox_id` (see NotificationChannelDO).
-  // Destroyed instances are filtered implicitly on the client — the dashboard only
-  // renders cards for instances returned by `listAllInstances`, which already
-  // excludes destroyed ones.
-  getUnreadCounts: baseProcedure.query(async ({ ctx }) => {
-    return readDb
-      .select({
-        badgeBucket: badge_counts.badge_bucket,
-        badgeCount: badge_counts.badge_count,
-      })
-      .from(badge_counts)
-      .where(and(eq(badge_counts.user_id, ctx.user.id), gt(badge_counts.badge_count, 0)));
-  }),
 });

packages/db/src/migrations/0108_drop_badge_counts.sql

@@ -0,0 +1 @@
+DROP TABLE "badge_counts" CASCADE;