121 lines
4.6 KiB
TypeScript
121 lines
4.6 KiB
TypeScript
import { protectRoute } from '~~/server/utils/auth';
|
|
import { startGeneration, addClientToGeneration, removeClientFromGeneration, sendToClient, isGenerationStreaming, getGenerationStatus } from '~~/server/utils/generation';
|
|
import { eventHandler, setHeader, setResponseStatus } from 'h3';
|
|
import { eq } from 'drizzle-orm';
|
|
import { useDrizzle } from '~~/server/utils/drizzle';
|
|
import { generations, messages } from '~~/db/schema';
|
|
|
|
export default eventHandler(async (event) => {
|
|
await protectRoute(event);
|
|
|
|
const generationId = getRouterParam(event, 'id');
|
|
|
|
if (!generationId) {
|
|
throw createError({
|
|
statusCode: 400,
|
|
statusMessage: 'Missing generation ID'
|
|
});
|
|
}
|
|
|
|
// Fetch generation from database
|
|
const generation = await getGenerationStatus(generationId);
|
|
if (!generation) {
|
|
throw createError({
|
|
statusCode: 404,
|
|
statusMessage: 'Generation not found'
|
|
});
|
|
}
|
|
|
|
// Verify ownership
|
|
if (generation.userId !== event.context.user.id) {
|
|
throw createError({
|
|
statusCode: 403,
|
|
statusMessage: 'Unauthorized'
|
|
});
|
|
}
|
|
|
|
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);
|
|
|
|
try {
|
|
const stream = new ReadableStream({
|
|
async start(controller) {
|
|
try {
|
|
// If generation is already completed, send the completed message
|
|
if (generation.status === 'completed' && generation.messageId) {
|
|
const db = useDrizzle();
|
|
const [message] = await db
|
|
.select()
|
|
.from(messages)
|
|
.where(eq(messages.id, generation.messageId));
|
|
|
|
if (message) {
|
|
sendToClient(controller, {
|
|
type: 'complete',
|
|
data: message
|
|
});
|
|
}
|
|
controller.close();
|
|
return;
|
|
}
|
|
|
|
// If generation failed, send the error
|
|
if (generation.status === 'failed') {
|
|
sendToClient(controller, {
|
|
type: 'error',
|
|
data: { error: generation.error || 'Generation failed' }
|
|
});
|
|
controller.close();
|
|
return;
|
|
}
|
|
|
|
// If already streaming, just add this client
|
|
if (isGenerationStreaming(generationId)) {
|
|
addClientToGeneration(generationId, controller);
|
|
} else {
|
|
// Start generation if in pending status
|
|
if (generation.status === 'pending') {
|
|
// Fetch the original messages context (stored in topic messages)
|
|
const db = useDrizzle();
|
|
const topicMessages = await db
|
|
.select()
|
|
.from(messages)
|
|
.where(eq(messages.topicId, generation.topicId));
|
|
|
|
const chatMessages = topicMessages.map(m => ({
|
|
type: m.isUser ? 'user' as const : ('agent' as const),
|
|
message: m.content
|
|
}));
|
|
|
|
addClientToGeneration(generationId, controller);
|
|
await startGeneration(generationId, generation.userId, generation.topicId, chatMessages, controller);
|
|
}
|
|
}
|
|
|
|
event.node.req.on('close', () => {
|
|
removeClientFromGeneration(generationId, controller);
|
|
});
|
|
} catch (error) {
|
|
console.error('Stream start error:', error);
|
|
sendToClient(controller, {
|
|
type: 'error',
|
|
data: { error: 'Stream initialization failed' }
|
|
});
|
|
controller.close();
|
|
}
|
|
}
|
|
});
|
|
|
|
return sendStream(event, stream);
|
|
} catch (error) {
|
|
console.error('Stream error:', error);
|
|
throw createError({
|
|
statusCode: 500,
|
|
statusMessage: 'Stream error'
|
|
});
|
|
}
|
|
}); |