67 lines
2.5 KiB
TypeScript
67 lines
2.5 KiB
TypeScript
import { protectRoute } from '~~/server/utils/auth';
|
|
import { getPendingGeneration, startGeneration, addClientToGeneration, removeClientFromGeneration, sendToClient } from '~~/server/utils/generation';
|
|
import { eventHandler, setHeader, setResponseStatus } from 'h3';
|
|
import { generations, messages } from '~~/db/schema';
|
|
import { eq } from 'drizzle-orm';
|
|
|
|
export default eventHandler(async (event) => {
|
|
await protectRoute(event);
|
|
|
|
const generationId = getRouterParam(event, 'id');
|
|
|
|
if (!generationId) {
|
|
throw createError({
|
|
statusCode: 400,
|
|
statusMessage: 'Missing generation ID'
|
|
});
|
|
}
|
|
|
|
const pendingGeneration = getPendingGeneration(generationId);
|
|
|
|
if (pendingGeneration && pendingGeneration.expired) {
|
|
throw createError({
|
|
statusCode: 410,
|
|
statusMessage: 'Generation expired - no client connected within 60 seconds'
|
|
});
|
|
}
|
|
|
|
setHeader(event, 'Content-Type', 'text/event-stream');
|
|
setHeader(event, 'Cache-Control', 'no-cache');
|
|
setHeader(event, 'Connection', 'keep-alive');
|
|
setHeader(event, 'X-Accel-Buffering', 'no');
|
|
|
|
setResponseStatus(event, 200);
|
|
|
|
const shouldStartGeneration = pendingGeneration;
|
|
|
|
try {
|
|
const stream = new ReadableStream({
|
|
async start(controller) {
|
|
addClientToGeneration(generationId, controller);
|
|
|
|
if (shouldStartGeneration) {
|
|
startGeneration(generationId, controller);
|
|
} else {
|
|
const generation = await useDrizzle().select().from(generations).where(eq(generations.id, getRouterParam(event, 'id')!))
|
|
if (!generation) throw createError({ statusCode: 404, statusMessage: 'Generation not found' });
|
|
const message = await useDrizzle().select().from(messages).where(eq(messages.id, generation[0].messageId!))
|
|
if (!message) throw createError({ statusCode: 404, statusMessage: 'Message not found' });
|
|
|
|
sendToClient(controller, {
|
|
type: 'complete',
|
|
data: message[0].content
|
|
});
|
|
controller.close();
|
|
}
|
|
|
|
event.node.req.on('close', () => {
|
|
removeClientFromGeneration(generationId, controller);
|
|
});
|
|
}
|
|
});
|
|
|
|
return sendStream(event, stream);
|
|
} catch (error) {
|
|
// assume it failed because the generation is complete so try to send the message
|
|
}
|
|
}); |