fix(chat): stop newly provisioned thread runs

Author: FranciscoMoretti

apps/chat/components/multimodal-input.tsx

@@ -622,10 +622,10 @@ function PureMultimodalInput({
 
   const handleStop = useCallback(() => {
     if (session?.user && lastMessageId) {
-      stopStreamMutation.mutate({ messageId: lastMessageId });
+      stopStreamMutation.mutate({ chatId, messageId: lastMessageId });
     }
     stopHelper?.();
-  }, [lastMessageId, session?.user, stopHelper, stopStreamMutation]);
+  }, [chatId, lastMessageId, session?.user, stopHelper, stopStreamMutation]);
 
   return (
     <div className="relative">

apps/chat/lib/db/queries.ts

@@ -1074,10 +1074,13 @@ export async function updateMessageCanceledAt({
   canceledAt: Date | null;
 }) {
   try {
-    return await db
+    const updatedMessages = await db
       .update(message)
       .set({ canceledAt })
-      .where(eq(message.id, messageId));
+      .where(eq(message.id, messageId))
+      .returning({ id: message.id });
+
+    return updatedMessages.length > 0;
   } catch (error) {
     logger.error({ error, messageId }, "updateMessageCanceledAt failed");
     throw error;

apps/chat/trpc/routers/chat.router.ts

@@ -188,30 +188,48 @@ export const chatRouter = createTRPCRouter({
   stopStream: protectedProcedure
     .input(
       z.object({
+        chatId: z.string().uuid(),
         messageId: z.string().uuid(),
       })
     )
     .mutation(async ({ ctx, input }) => {
-      const [msg] = await getMessageById({ id: input.messageId });
-      if (!msg) {
-        throw new TRPCError({
-          code: "NOT_FOUND",
-          message: "Message not found",
-        });
+      // A new chat can still be generating its title when the client stops.
+      // Wait for its reserved assistant row so request ordering cannot lose cancellation.
+      const deadline = Date.now() + 30_000;
+      let retryDelay = 25;
+      let chat = await getChatById({ id: input.chatId });
+
+      while (!chat && Date.now() < deadline) {
+        await new Promise((resolve) => setTimeout(resolve, retryDelay));
+        retryDelay = Math.min(retryDelay * 2, 250);
+        chat = await getChatById({ id: input.chatId });
       }
 
-      const chat = await getChatById({ id: msg.chatId });
       if (!chat || chat.userId !== ctx.user.id) {
         throw new TRPCError({
           code: "NOT_FOUND",
           message: "Chat not found or access denied",
         });
       }
 
-      await updateMessageCanceledAt({
-        messageId: input.messageId,
-        canceledAt: new Date(),
-      });
+      const canceledAt = new Date();
+
+      while (
+        !(await updateMessageCanceledAt({
+          messageId: input.messageId,
+          canceledAt,
+        }))
+      ) {
+        if (Date.now() >= deadline) {
+          throw new TRPCError({
+            code: "NOT_FOUND",
+            message: "Message not found",
+          });
+        }
+
+        await new Promise((resolve) => setTimeout(resolve, retryDelay));
+        retryDelay = Math.min(retryDelay * 2, 250);
+      }
 
       return { success: true };
     }),