import { TRPCError, tracked } from "@trpc/server"; import { z } from "zod/v4"; import { adminProcedure, createTRPCRouter, protectedProcedure, } from "~/server/api/trpc"; import { getCurrentlyTyping, removeStaleTyping, removeTyping, } from "~/server/controllers/chat.controller"; import { db } from "~/server/db"; import { type ChatMessage, editMessage, setReadMessage, } from "~/server/services/messages.service"; import { ee, type MessageUpdate } from "~/server/sse"; import { zAsyncIterable } from "~/server/utils/sse_zod"; import { zChatId, zMessageId, zUserId } from "~/server/utils/zod_types"; export const messagesSSERouter = createTRPCRouter({ add: protectedProcedure .input( z.object({ message: z.string(), chatId: zChatId, sender: zUserId, }), ) .mutation(async ({ input }) => { const newMessage = await db .insertInto("messages") .values({ message: input.message, sender: input.sender, chatid: input.chatId, time: new Date(), }) .returningAll() .executeTakeFirstOrThrow(); const userInfos = await db .selectFrom("users") .select([ "users.nome as sender_nome", "users.isAdmin as sender_isAdmin", ]) .where("id", "=", input.sender) .executeTakeFirst(); if (!userInfos) { throw new TRPCError({ code: "NOT_FOUND", message: "User not found or no name", }); } await removeTyping({ chatId: input.chatId, userId: input.sender, username: userInfos.sender_nome, }); await removeStaleTyping(input.chatId); ee.emit("chatUpdates", input.chatId, { eventType: "typing", data: await getCurrentlyTyping(input.chatId), }); ee.emit("add", input.chatId, { ...newMessage, ...userInfos, }); return { ...newMessage, ...userInfos }; }), edit: adminProcedure .input( z.object({ chatId: zChatId, messageId: zMessageId, message: z.string() }), ) .mutation(async ({ input }) => { await editMessage({ db, messageId: input.messageId, message: input.message, }); ee.emit("messageUpdates", input.chatId, { eventType: "edit", data: { messageId: input.messageId, message: input.message, }, }); }), delete: adminProcedure .input(z.object({ chatId: zChatId, messageId: zMessageId })) .mutation(async ({ input }) => { await db .deleteFrom("messages") .where("messageid", "=", input.messageId) .execute(); ee.emit("messageUpdates", input.chatId, { eventType: "delete", data: input.messageId, }); }), read: protectedProcedure .input(z.object({ chatId: zChatId })) .mutation(async ({ input, ctx }) => { await setReadMessage({ db, chatId: input.chatId, reader: ctx.session.id, }); return true; }), infinite: protectedProcedure .input( z.object({ chatId: zChatId, cursor: z.date().nullish(), take: z.number().min(1).max(50).nullish(), }), ) .query(async ({ input, ctx }) => { try { await setReadMessage({ db, chatId: input.chatId, reader: ctx.session.id, }); const take = input.take ?? 20; const cursor = input.cursor; let qry = db .selectFrom("messages") .selectAll() .innerJoin("users", "users.id", "messages.sender") .select([ "users.nome as sender_nome", "users.isAdmin as sender_isAdmin", ]) .where("users.nome", "is not", null) .where("messages.chatid", "=", input.chatId) .orderBy("messages.time", "desc") .limit(take + 1); if (cursor) { qry = qry.where("messages.time", "<=", cursor); } const page = await qry.execute(); const items = page.reverse(); let nextCursor: typeof cursor | null = null; if (items.length > take) { const prev = items.shift(); // biome-ignore lint/style/noNonNullAssertion: nextCursor = prev!.time; } return { items, nextCursor, }; } catch (e) { throw new TRPCError({ code: "INTERNAL_SERVER_ERROR", message: `Error while fetching messages: ${(e as Error).message}`, }); } }), onAdd: protectedProcedure .input( z.object({ chatId: zChatId, // lastEventId is the last event id that the client has received // On the first call, it will be whatever was passed in the initial setup // If the client reconnects, it will be the last event id that the client received lastEventId: zMessageId.nullish(), }), ) .subscription(async function* ({ input, signal, ctx }) { try { // We start by subscribing to the event emitter so that we don't miss any new events while fetching const iterable = ee.toIterable("add", { signal, }); // Fetch the last message createdAt based on the last event id let lastMessageCreatedAt = await (async () => { const lastEventId = input.lastEventId; console.log("lastEventId", lastEventId); if (!lastEventId) return null; const messageById = await db .selectFrom("messages") .selectAll() .where("messageid", "=", lastEventId) .executeTakeFirstOrThrow(); return messageById?.time ?? null; })(); let qry = db .selectFrom("messages") .selectAll() .innerJoin("users", "users.id", "messages.sender") .select([ "users.nome as sender_nome", "users.isAdmin as sender_isAdmin", ]) .where("messages.chatid", "=", input.chatId) .orderBy("messages.time", "asc"); if (lastMessageCreatedAt) { qry = qry.where("messages.time", ">", lastMessageCreatedAt); } const newMessagesSinceLastMessage = await qry.execute(); await setReadMessage({ db, chatId: input.chatId, reader: ctx.session.id, }); function* maybeYield(msg: ChatMessage) { if (msg.chatid !== input.chatId) { // ignore messages from other channels - the event emitter can emit from other channels return; } if (lastMessageCreatedAt && msg.time <= lastMessageCreatedAt) { // ignore messages that we've already sent - happens if there is a race condition between the query and the event emitter return; } yield tracked(msg.messageid, msg); // update the cursor so that we don't send this post again lastMessageCreatedAt = msg.time; } // yield the messages we fetched from the db for (const post of newMessagesSinceLastMessage) { yield* maybeYield(post); } // yield any new messages from the event emitter for await (const [, post] of iterable) { yield* maybeYield(post); } } catch (e) { if (e instanceof Error && e.name === "AbortError") return; console.error("Error in onAdd subscription:", e); } }), onMessageEvent: protectedProcedure .input(z.object({ chatId: zChatId })) .output( zAsyncIterable({ yield: z.custom(), tracked: true, }), ) .subscription(async function* ({ input, signal }) { try { const iterable = ee.toIterable("messageUpdates", { signal, }); for await (const [chatId, data] of iterable) { if (chatId === input.chatId) { yield tracked(chatId, data); } } } catch (e) { if (e instanceof Error && e.name === "AbortError") return; console.error("Error in onMessageEvent subscription:", e); } }), });