infoalloggi-monorepo/apps/infoalloggi/src/server/utils/sse_zod.ts

/** biome-ignore-all lint/suspicious/noExplicitAny: <intended> */
import type { TrackedEnvelope } from "@trpc/server";
import { isTrackedEnvelope, tracked } from "@trpc/server";
import { z } from "zod";

function isAsyncIterable<TValue, TReturn = unknown>(
	value: unknown,
): value is AsyncIterable<TValue, TReturn> {
	return !!value && typeof value === "object" && Symbol.asyncIterator in value;
}
const trackedEnvelopeSchema =
	z.custom<TrackedEnvelope<unknown>>(isTrackedEnvelope);
/**
 * A Zod schema helper designed specifically for validating async iterables. This schema ensures that:
 * 1. The value being validated is an async iterable.
 * 2. Each item yielded by the async iterable conforms to a specified type.
 * 3. The return value of the async iterable, if any, also conforms to a specified type.
 */
export function zAsyncIterable<
	TYield,
	TReturn = void,
	Tracked extends boolean = false,
>(opts: {
	yield: z.ZodType<TYield, any, any>;
	return?: z.ZodType<TReturn, any, any>;
	tracked?: Tracked;
}) {
	const baseSchema = z.custom<
		AsyncIterable<
			Tracked extends true
				? TrackedEnvelope<z.input<typeof opts.yield>>
				: z.input<typeof opts.yield>,
			z.input<typeof opts.return>
		>
	>((val: unknown): val is AsyncIterable<any, any> => isAsyncIterable(val), {
		message: "Expected AsyncIterable",
	});

	return baseSchema.transform(async function* (iter) {
		const iterator = iter[Symbol.asyncIterator]();
		try {
			// biome-ignore lint/suspicious/noImplicitAnyLet: <intended>
			let next;
			// biome-ignore lint/suspicious/noAssignInExpressions: <intended>
			while ((next = await iterator.next()) && !next.done) {
				if (opts.tracked) {
					const [id, data] = await trackedEnvelopeSchema.parseAsync(next.value);
					yield tracked(id, await opts.yield.parseAsync(data));
					continue;
				}
				yield await opts.yield.parseAsync(next.value);
			}
			if (opts.return) {
				return (await opts.return.parseAsync(next.value)) as TReturn;
			}
			return;
		} finally {
			await iterator.return?.();
		}
	});
}