RunStreamPresenter.server.ts

import { PrismaClient, prisma } from "~/db.server";
import { logger } from "~/services/logger.server";
import { singleton } from "~/utils/singleton";
import { createSSELoader } from "~/utils/sse";
import { throttle } from "~/utils/throttle";
import { tracePubSub } from "~/v3/services/tracePubSub.server";

const PING_INTERVAL = 1000;
const STREAM_TIMEOUT = 30 * 1000; // 30 seconds

export class RunStreamPresenter {
  #prismaClient: PrismaClient;

  constructor(prismaClient: PrismaClient = prisma) {
    this.#prismaClient = prismaClient;
  }

  public createLoader() {
    const prismaClient = this.#prismaClient;

    return createSSELoader({
      timeout: STREAM_TIMEOUT,
      interval: PING_INTERVAL,
      handler: async (context) => {
        const runFriendlyId = context.params.runParam;

        if (!runFriendlyId) {
          throw new Response("Missing runParam", { status: 400 });
        }

        const run = await prismaClient.taskRun.findFirst({
          where: {
            friendlyId: runFriendlyId,
          },
          select: {
            traceId: true,
          },
        });

        if (!run) {
          throw new Response("Not found", { status: 404 });
        }

        logger.info("RunStreamPresenter.start", {
          runFriendlyId,
          traceId: run.traceId,
        });

        // Subscribe to trace updates
        const { unsubscribe, eventEmitter } = await tracePubSub.subscribeToTrace(run.traceId);

        // Store throttled send function and message listener for cleanup
        let throttledSend: ReturnType<typeof throttle> | undefined;
        let messageListener: ((event: string) => void) | undefined;

        return {
          initStream: ({ send }) => {
            // Create throttled send function
            throttledSend = throttle((args: { event?: string; data: string }) => {
              try {
                send(args);
              } catch (error) {
                if (error instanceof Error) {
                  if (error.name !== "TypeError") {
                    logger.debug("Error sending SSE in RunStreamPresenter", {
                      error: {
                        name: error.name,
                        message: error.message,
                        stack: error.stack,
                      },
                    });
                  }
                }
                // Abort the stream on send error
                context.controller.abort("Send error");
              }
            }, 1000);

            // Set up message listener for pub/sub events
            messageListener = (event: string) => {
              throttledSend?.({ data: event });
            };
            eventEmitter.addListener("message", messageListener);

            context.debug("Subscribed to trace pub/sub");
          },

          iterator: ({ send }) => {
            // Send ping to keep connection alive
            try {
              send({ event: "ping", data: new Date().toISOString() });
            } catch (error) {
              // If we can't send a ping, the connection is likely dead
              return false;
            }
          },

          cleanup: () => {
            logger.info("RunStreamPresenter.cleanup", {
              runFriendlyId,
              traceId: run.traceId,
            });

            // Remove message listener
            if (messageListener) {
              eventEmitter.removeListener("message", messageListener);
            }
            eventEmitter.removeAllListeners();

            // Unsubscribe from Redis pub/sub
            unsubscribe()
              .then(() => {
                logger.info("RunStreamPresenter.cleanup.unsubscribe succeeded", {
                  runFriendlyId,
                  traceId: run.traceId,
                });
              })
              .catch((error) => {
                logger.error("RunStreamPresenter.cleanup.unsubscribe failed", {
                  runFriendlyId,
                  traceId: run.traceId,
                  error: {
                    name: error.name,
                    message: error.message,
                    stack: error.stack,
                  },
                });
              });
          },
        };
      },
    });
  }
}

// Export a singleton loader for the route to use
export const runStreamLoader = singleton("runStreamLoader", () => {
  const presenter = new RunStreamPresenter();
  return presenter.createLoader();
});