feat: efficient concurrent sessions of same thread support

Author: NamesMT

apps/backend-convex/convex/_counters.ts

@@ -0,0 +1,4 @@
+import { ShardedCounter } from '@convex-dev/sharded-counter'
+import { components } from './_generated/api'
+
+export const messagesInThreadCounter = new ShardedCounter(components.shardedCounter, { defaultShards: 1 })

apps/backend-convex/convex/_generated/api.d.ts

@@ -8,6 +8,7 @@
  * @module
  */
 
+import type * as _counters from "../_counters.js";
 import type * as authInfo from "../authInfo.js";
 import type * as crons from "../crons.js";
 import type * as http_ai from "../http/ai.js";
@@ -31,6 +32,7 @@ import type {
  * ```
  */
 declare const fullApi: ApiFromModules<{
+  _counters: typeof _counters;
   authInfo: typeof authInfo;
   crons: typeof crons;
   "http/ai": typeof http_ai;
@@ -195,4 +197,28 @@ export declare const components: {
       >;
     };
   };
+  shardedCounter: {
+    public: {
+      add: FunctionReference<
+        "mutation",
+        "internal",
+        { count: number; name: string; shard?: number; shards?: number },
+        number
+      >;
+      count: FunctionReference<"query", "internal", { name: string }, number>;
+      estimateCount: FunctionReference<
+        "query",
+        "internal",
+        { name: string; readFromShards?: number; shards?: number },
+        any
+      >;
+      rebalance: FunctionReference<
+        "mutation",
+        "internal",
+        { name: string; shards?: number },
+        any
+      >;
+      reset: FunctionReference<"mutation", "internal", { name: string }, any>;
+    };
+  };
 };

apps/backend-convex/convex/convex.config.ts

@@ -1,7 +1,9 @@
 import rateLimiter from '@convex-dev/rate-limiter/convex.config'
+import shardedCounter from '@convex-dev/sharded-counter/convex.config'
 import { defineApp } from 'convex/server'
 
 const app: ReturnType<typeof defineApp> = defineApp()
 app.use(rateLimiter)
+app.use(shardedCounter)
 
 export default app

apps/backend-convex/convex/http/ai.ts

@@ -163,7 +163,7 @@ aiApp
       }
 
       // Get conversation history
-      const messages = await c.env.runQuery(api.messages.list, { threadId, lockerKey })
+      const messages = await c.env.runQuery(api.messages.listByThread, { threadId, lockerKey })
 
       // Prepare messages for AI API (exclude the streaming messages)
       const messagesContext = messages

apps/backend-convex/convex/messages.ts

@@ -1,9 +1,10 @@
 import { objectPick } from '@namesmt/utils'
 import { ConvexError, v } from 'convex/values'
+import { messagesInThreadCounter } from './_counters'
 import { internalMutation, internalQuery, mutation, query } from './_generated/server'
 import { assertThreadAccess } from './threads'
 
-export const list = query({
+export const listByThread = query({
   args: {
     threadId: v.id('threads'),
     lockerKey: v.optional(v.string()),
@@ -24,6 +25,22 @@ export const list = query({
   },
 })
 
+export const countByThread = query({
+  args: {
+    threadId: v.id('threads'),
+    lockerKey: v.optional(v.string()),
+  },
+  handler: async (ctx, args) => {
+    const thread = await ctx.db.get(args.threadId)
+    if (!thread)
+      throw new ConvexError('Thread not found')
+
+    await assertThreadAccess(ctx, { thread, lockerKey: args.lockerKey })
+
+    return await messagesInThreadCounter.count(ctx, args.threadId)
+  },
+})
+
 export const get = query({
   args: {
     messageId: v.id('messages'),
@@ -62,6 +79,8 @@ export const add = mutation({
 
     await assertThreadAccess(ctx, { thread, lockerKey: args.lockerKey })
 
+    await messagesInThreadCounter.inc(ctx, args.threadId)
+
     return await ctx.db.insert('messages', {
       ...objectPick(args, ['threadId', 'role', 'content', 'isStreaming', 'streamId', 'provider', 'model']),
       timestamp: Date.now(),

apps/backend-convex/convex/threads.ts

@@ -6,6 +6,7 @@ import { clearUndefined } from '@namesmt/utils'
 import { openrouter } from '@openrouter/ai-sdk-provider'
 import { generateText } from 'ai'
 import { ConvexError, v } from 'convex/values'
+import { messagesInThreadCounter } from './_counters'
 import { api, internal } from './_generated/api'
 import { action, internalMutation, mutation, query } from './_generated/server'
 
@@ -138,6 +139,8 @@ export const branchThreadFromMessage = mutation({
       })
     }))
 
+    await messagesInThreadCounter.inc(ctx, thread._id)
+
     return newThreadId
   },
 })
@@ -188,7 +191,7 @@ export const generateThreadTitle = action({
   handler: async (ctx, args) => {
     const thread = await ctx.runQuery(api.threads.get, { threadId: args.threadId, lockerKey: args.lockerKey })
 
-    const messages = await ctx.runQuery(api.messages.list, { threadId: args.threadId, lockerKey: args.lockerKey })
+    const messages = await ctx.runQuery(api.messages.listByThread, { threadId: args.threadId, lockerKey: args.lockerKey })
 
     const { text } = await generateText({
       model: openrouter('qwen/qwen3-8b:free'),

apps/backend-convex/package.json

@@ -19,6 +19,7 @@
   "dependencies": {
     "@ai-sdk/openai": "^1.3.22",
     "@convex-dev/rate-limiter": "^0.2.7",
+    "@convex-dev/sharded-counter": "^0.1.8",
     "@hono/zod-validator": "^0.7.0",
     "@openrouter/ai-sdk-provider": "^0.7.2",
     "ai": "^4.3.16",

apps/frontend/app/components/chat/ChatInterface.vue

@@ -42,6 +42,7 @@ const nearTopBottom = computed(() => {
 
 const threadIdRef = useThreadIdRef()
 const isThreadFrozen = computed(() => chatContext.activeThread.value?.frozen)
+const fetchKey = ref(0)
 
 const cachedThreadsMessages: {
   [threadId: string]: Array<CustomMessage>
@@ -52,10 +53,10 @@ const streamingMessages = ref(0)
 const isFetching = ref(false)
 const chatInput = ref('')
 
-// Fetch messages and resume streams
+// Fetch messages as needed and resume streams
 const { ignoreUpdates: ignorePathUpdate } = watchIgnorable(
-  threadIdRef,
-  async (threadId, oldThreadId) => {
+  [threadIdRef, fetchKey],
+  async ([threadId], [oldThreadId]) => {
     if (oldThreadId)
       cachedThreadsMessages[oldThreadId] = messages.value
 
@@ -65,7 +66,7 @@ const { ignoreUpdates: ignorePathUpdate } = watchIgnorable(
 
     if (threadId) {
       isFetching.value = true
-      await convex.query(api.messages.list, { threadId: threadId as Doc<'threads'>['_id'], lockerKey: getLockerKey(threadId) })
+      await convex.query(api.messages.listByThread, { threadId: threadId as Doc<'threads'>['_id'], lockerKey: getLockerKey(threadId) })
         .then((existingMessages) => {
           if (threadIdRef.value === threadId) {
             messages.value = existingMessages.map(customMessageTransform)
@@ -77,7 +78,6 @@ const { ignoreUpdates: ignorePathUpdate } = watchIgnorable(
 
           // If the owner have deleted the thread, remove it locally
           // (or the demo crons cleaned it)
-          console.log({ a: getConvexErrorMessage(e), e })
           if (getConvexErrorMessage(e) === 'Thread not found') {
             toast({ variant: 'destructive', description: t('chat.toast.threadRemovedExternal') })
 
@@ -109,6 +109,23 @@ const { ignoreUpdates: ignorePathUpdate } = watchIgnorable(
   { immediate: true },
 )
 
+// Subscribe to a counter to check for messages from other concurrent sessions
+watchImmediate(threadIdRef, (threadId) => {
+  console.log(`Subscribing to messages count of: ${threadId}`)
+  const { unsubscribe } = convex.onUpdate(
+    api.messages.countByThread,
+    { threadId: threadId as Id<'threads'>, lockerKey: getLockerKey(threadId) },
+    (count) => {
+      if (count > messages.value.length)
+        ++fetchKey.value
+    },
+  )
+  watchOnce(threadIdRef, () => {
+    unsubscribe()
+    console.log(`Unsubscribed from ${threadId}`)
+  })
+})
+
 interface HandleSubmitArgs {
   input: string
   confirmMultiStream?: boolean

apps/frontend/app/pages/chat/[...all].vue

@@ -37,7 +37,7 @@ provideSidebarContext({
     </div> -->
 
     <SidebarProvider>
-      <ChatSidebar />
+      <ChatSidebar class="z-5" />
       <ChatInterface class="h-full w-full" />
       <ChatFloatingMenu class="absolute left-2 top-2 z-10" />
     </SidebarProvider>