361 lines
9.6 KiB
TypeScript
361 lines
9.6 KiB
TypeScript
import ReconnectingEventSource from "reconnecting-eventsource";
|
|
import { z } from "zod";
|
|
import { getRedisPublisher, getRedisSubscriber } from "@/lib/redis";
|
|
|
|
const CHANNEL_PREFIX = "erika:sse:";
|
|
const DEFAULT_RETRY_MS = 3_000;
|
|
const DEFAULT_PUBLISH_TIMEOUT_MS = 2_000;
|
|
const DEFAULT_HEARTBEAT_MS = 90_000;
|
|
const DEFAULT_CONNECTION_MAX_MS = 30 * 60_000;
|
|
|
|
type EventMap = Record<string, z.ZodType>;
|
|
type EventHandlers<T extends EventMap> = Partial<{
|
|
[K in keyof T]: (data: z.infer<T[K]>) => void;
|
|
}>;
|
|
|
|
type SubscriptionOptions = {
|
|
endpoint?: string | URL;
|
|
onerror?: (event: Event) => void;
|
|
onopen?: (event: Event) => void;
|
|
};
|
|
|
|
type StreamMessage<T extends EventMap> = {
|
|
[K in Extract<keyof T, string>]: {
|
|
event: K;
|
|
data: z.input<T[K]>;
|
|
};
|
|
}[Extract<keyof T, string>];
|
|
|
|
type StreamOptions<T extends EventMap> = {
|
|
signal?: AbortSignal;
|
|
motd?: StreamMessage<T> | StreamMessage<T>[];
|
|
};
|
|
|
|
const redisEnvelopeSchema = z.object({
|
|
event: z.string(),
|
|
data: z.unknown(),
|
|
}).strict();
|
|
|
|
function durationFromEnvironment(name: string, fallback: number) {
|
|
const value = Number(process.env[name]);
|
|
return Number.isFinite(value) && value > 0 ? value : fallback;
|
|
}
|
|
|
|
function channelName(topic: string) {
|
|
return `${CHANNEL_PREFIX}${topic}`;
|
|
}
|
|
|
|
function encodeEvent(event: string, data: unknown) {
|
|
return `event: ${event}\ndata: ${JSON.stringify(data)}\n\n`;
|
|
}
|
|
|
|
async function publish(topic: string, event: string, data: unknown) {
|
|
let timeout: ReturnType<typeof setTimeout> | undefined;
|
|
const operation = getRedisPublisher().then((publisher) =>
|
|
publisher.publish(
|
|
channelName(topic),
|
|
JSON.stringify({ event, data }),
|
|
),
|
|
);
|
|
const deadline = new Promise<never>((_, reject) => {
|
|
timeout = setTimeout(
|
|
() => reject(new Error("Redis SSE publish timed out")),
|
|
DEFAULT_PUBLISH_TIMEOUT_MS,
|
|
);
|
|
});
|
|
|
|
try {
|
|
await Promise.race([operation, deadline]);
|
|
} finally {
|
|
if (timeout) clearTimeout(timeout);
|
|
}
|
|
}
|
|
|
|
export class SseEndpoint<T extends EventMap> {
|
|
constructor(
|
|
readonly topic: string,
|
|
private readonly events: T,
|
|
readonly adminOnly = false,
|
|
) {}
|
|
|
|
pub<K extends Extract<keyof T, string>>(event: K, data: z.input<T[K]>) {
|
|
const payload = this.events[event].parse(data);
|
|
|
|
return publish(this.topic, event, payload).catch((error: unknown) => {
|
|
console.error(`Failed to publish SSE event ${this.topic}.${event}`, error);
|
|
});
|
|
}
|
|
|
|
subMany(
|
|
handlers: EventHandlers<T>,
|
|
{
|
|
endpoint = `/sse/${encodeURIComponent(this.topic)}`,
|
|
onerror,
|
|
onopen,
|
|
}: SubscriptionOptions = {},
|
|
) {
|
|
const eventSource = new ReconnectingEventSource(endpoint);
|
|
const listeners: Array<[string, (event: MessageEvent<string>) => void]> = [];
|
|
|
|
for (const event of Object.keys(handlers) as Array<Extract<keyof T, string>>) {
|
|
const callback = handlers[event];
|
|
if (!callback) continue;
|
|
|
|
const listener = (message: MessageEvent<string>) => {
|
|
try {
|
|
const payload = this.events[event].safeParse(JSON.parse(message.data));
|
|
if (payload.success) {
|
|
callback(payload.data);
|
|
} else {
|
|
console.error(
|
|
`Ignored invalid SSE event ${this.topic}.${event}`,
|
|
payload.error,
|
|
);
|
|
}
|
|
} catch (error) {
|
|
console.error(`Ignored malformed SSE event ${this.topic}.${event}`, error);
|
|
}
|
|
};
|
|
|
|
listeners.push([event, listener]);
|
|
eventSource.addEventListener(event, listener);
|
|
}
|
|
|
|
if (onerror) eventSource.onerror = onerror;
|
|
if (onopen) eventSource.onopen = onopen;
|
|
|
|
const clean = () => {
|
|
for (const [event, listener] of listeners) {
|
|
eventSource.removeEventListener(event, listener);
|
|
}
|
|
eventSource.close();
|
|
};
|
|
|
|
return { clean, eventSource };
|
|
}
|
|
|
|
sub<K extends Extract<keyof T, string>>(
|
|
event: K,
|
|
handler: (data: z.infer<T[K]>) => void,
|
|
options?: SubscriptionOptions,
|
|
) {
|
|
return this.subMany(
|
|
{ [event]: handler } as unknown as EventHandlers<T>,
|
|
options,
|
|
);
|
|
}
|
|
|
|
parseRedisMessage(message: string) {
|
|
try {
|
|
const envelope = redisEnvelopeSchema.safeParse(JSON.parse(message));
|
|
if (!envelope.success || !Object.hasOwn(this.events, envelope.data.event)) {
|
|
return null;
|
|
}
|
|
|
|
const event = envelope.data.event as Extract<keyof T, string>;
|
|
const payload = this.events[event].safeParse(envelope.data.data);
|
|
if (!payload.success) return null;
|
|
|
|
return { event, data: payload.data };
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
async stream({ signal, motd }: StreamOptions<T> = {}) {
|
|
const initialMessages = (Array.isArray(motd) ? motd : motd ? [motd] : [])
|
|
.map(({ event, data }) =>
|
|
encodeEvent(event, this.events[event].parse(data)),
|
|
);
|
|
const subscriber = await getRedisSubscriber();
|
|
const topic = this.topic;
|
|
const channel = channelName(this.topic);
|
|
const pendingMessages: string[] = [];
|
|
let receive: ((message: string) => void) | undefined;
|
|
|
|
const listener = (message: string) => {
|
|
if (receive) {
|
|
receive(message);
|
|
} else {
|
|
pendingMessages.push(message);
|
|
}
|
|
};
|
|
|
|
await subscriber.subscribe(channel, listener);
|
|
|
|
let heartbeat: ReturnType<typeof setInterval> | undefined;
|
|
let renewal: ReturnType<typeof setTimeout> | undefined;
|
|
let controller: ReadableStreamDefaultController<Uint8Array> | undefined;
|
|
let closed = false;
|
|
const encoder = new TextEncoder();
|
|
|
|
const unsubscribe = () => {
|
|
void subscriber.unsubscribe(channel, listener).catch((error: unknown) => {
|
|
console.error(`Failed to unsubscribe SSE topic ${this.topic}`, error);
|
|
});
|
|
};
|
|
|
|
const clean = () => {
|
|
if (closed) return;
|
|
closed = true;
|
|
if (heartbeat) clearInterval(heartbeat);
|
|
if (renewal) clearTimeout(renewal);
|
|
signal?.removeEventListener("abort", close);
|
|
unsubscribe();
|
|
};
|
|
|
|
const close = () => {
|
|
clean();
|
|
try {
|
|
controller?.close();
|
|
} catch {
|
|
// The request or reader may have already closed the stream.
|
|
}
|
|
};
|
|
|
|
const write = (value: string) => {
|
|
if (closed) return;
|
|
try {
|
|
controller?.enqueue(encoder.encode(value));
|
|
} catch {
|
|
close();
|
|
}
|
|
};
|
|
|
|
const handleMessage = (message: string) => {
|
|
const parsed = this.parseRedisMessage(message);
|
|
if (!parsed) {
|
|
console.error(`Ignored invalid Redis SSE message for ${this.topic}`);
|
|
return;
|
|
}
|
|
write(encodeEvent(parsed.event, parsed.data));
|
|
};
|
|
|
|
const stream = new ReadableStream<Uint8Array>({
|
|
start(streamController) {
|
|
controller = streamController;
|
|
write(`retry: ${DEFAULT_RETRY_MS}\n\n: connected ${topic}\n\n`);
|
|
for (const message of initialMessages) write(message);
|
|
|
|
receive = handleMessage;
|
|
for (const message of pendingMessages.splice(0)) {
|
|
handleMessage(message);
|
|
}
|
|
|
|
heartbeat = setInterval(
|
|
() => write(": heartbeat\n\n"),
|
|
durationFromEnvironment("SSE_HEARTBEAT_MS", DEFAULT_HEARTBEAT_MS),
|
|
);
|
|
renewal = setTimeout(
|
|
close,
|
|
durationFromEnvironment(
|
|
"SSE_CONNECTION_MAX_MS",
|
|
DEFAULT_CONNECTION_MAX_MS,
|
|
),
|
|
);
|
|
signal?.addEventListener("abort", close, { once: true });
|
|
if (signal?.aborted) close();
|
|
},
|
|
cancel() {
|
|
clean();
|
|
},
|
|
});
|
|
|
|
return new Response(stream, {
|
|
headers: {
|
|
"Content-Type": "text/event-stream; charset=utf-8",
|
|
"Cache-Control": "no-cache, no-transform",
|
|
"Connection": "keep-alive",
|
|
"Content-Encoding": "identity",
|
|
"X-Accel-Buffering": "no",
|
|
},
|
|
});
|
|
}
|
|
}
|
|
|
|
type EndpointDefinition<T extends EventMap> = {
|
|
events: T;
|
|
adminOnly?: boolean;
|
|
};
|
|
|
|
export function createSseEndpoints<
|
|
T extends Record<string, EndpointDefinition<EventMap>>,
|
|
>(definitions: T) {
|
|
const endpoints = {} as {
|
|
[K in keyof T]: SseEndpoint<T[K]["events"]>;
|
|
};
|
|
|
|
for (const topic of Object.keys(definitions) as Array<keyof T>) {
|
|
const definition = definitions[topic];
|
|
endpoints[topic] = new SseEndpoint(
|
|
String(topic),
|
|
definition.events,
|
|
definition.adminOnly,
|
|
);
|
|
}
|
|
|
|
return endpoints;
|
|
}
|
|
|
|
const updateActionSchema = z.enum(["created", "updated", "deleted"]);
|
|
|
|
export const sse = createSseEndpoints({
|
|
active: {
|
|
events: {
|
|
version: z.string().min(1),
|
|
refresh: z.string().min(1).nullable(),
|
|
},
|
|
},
|
|
forms: {
|
|
events: {
|
|
update: z.object({
|
|
formId: z.string().min(1),
|
|
entity: z.enum(["form", "question"]),
|
|
action: updateActionSchema,
|
|
}).strict(),
|
|
},
|
|
},
|
|
submissions: {
|
|
adminOnly: true,
|
|
events: {
|
|
update: z.object({
|
|
formId: z.string().min(1),
|
|
submissionId: z.string().min(1).optional(),
|
|
action: updateActionSchema,
|
|
}).strict(),
|
|
},
|
|
},
|
|
uploads: {
|
|
adminOnly: true,
|
|
events: {
|
|
update: z.object({
|
|
uploadId: z.string().min(1),
|
|
jobId: z.string().min(1).optional(),
|
|
platform: z.enum(["youtube", "tiktok"]).optional(),
|
|
status: z.enum([
|
|
"scheduled",
|
|
"processing",
|
|
"pending",
|
|
"completed",
|
|
"partial",
|
|
"failed",
|
|
]),
|
|
}).strict(),
|
|
},
|
|
},
|
|
leaderboards: {
|
|
events: {
|
|
xp: z.object({}).strict(),
|
|
vc: z.object({
|
|
channelId: z.string().regex(/^\d+$/),
|
|
}).strict(),
|
|
},
|
|
},
|
|
});
|
|
|
|
export type SseTopic = keyof typeof sse;
|
|
|
|
export function isSseTopic(topic: string): topic is SseTopic {
|
|
return Object.hasOwn(sse, topic);
|
|
}
|