fix: cf worker waitUntil timeout

Author: monkeyWie

src/app/routes/chat/-components/chat/ChatMessageItem.tsx

@@ -40,6 +40,7 @@ export function ChatMessageItem({
 	const chatService = useChatService();
 	const intervalRef = useRef<NodeJS.Timeout | null>(null);
 	const skipPoll = useRef<boolean>(false);
+	const pollingGenerationIdRef = useRef<string | null>(null); // Track which generation is being polled
 	const [isLightboxOpen, setIsLightboxOpen] = useState(false);
 	const [currentImageIndex, setCurrentImageIndex] = useState(0);
 	const [isHovered, setIsHovered] = useState(false);
@@ -122,6 +123,20 @@ export function ChatMessageItem({
 		const generationId = message.generationId || message.generation?.id;
 
 		if (isMessageGenerating && generationId && onMessageUpdate) {
+			// If we're already polling this exact generation, don't restart
+			if (intervalRef.current && pollingGenerationIdRef.current === generationId) {
+				return;
+			}
+
+			// Clear any existing interval for a different generation
+			if (intervalRef.current) {
+				clearInterval(intervalRef.current);
+				intervalRef.current = null;
+			}
+
+			// Mark this generation as being polled
+			pollingGenerationIdRef.current = generationId;
+
 			const pollStatus = async () => {
 				try {
 					if (skipPoll.current) {
@@ -143,6 +158,7 @@ export function ChatMessageItem({
 							if (intervalRef.current) {
 								clearInterval(intervalRef.current);
 								intervalRef.current = null;
+								pollingGenerationIdRef.current = null;
 							}
 						}
 					}
@@ -151,21 +167,34 @@ export function ChatMessageItem({
 				}
 			};
 
-			// Start polling every 3 seconds
-			intervalRef.current = setInterval(pollStatus, 3000);
-
-			// Also poll immediately (unless skipFirstPoll is set)
-			pollStatus();
+			// Start polling every 3 seconds after initial 3 second delay
+			const timeoutId = setTimeout(() => {
+				if (isMessageGenerating && pollingGenerationIdRef.current === generationId) {
+					pollStatus();
+					// Start interval polling after first poll completes
+					intervalRef.current = setInterval(pollStatus, 3000);
+				}
+			}, 3000);
 
 			// Cleanup on unmount or when generation is no longer pending
 			return () => {
+				clearTimeout(timeoutId);
 				if (intervalRef.current) {
 					clearInterval(intervalRef.current);
 					intervalRef.current = null;
+					pollingGenerationIdRef.current = null;
 				}
 			};
 		}
-	}, [isMessageGenerating, message.generationId, message.generation?.id, message.id, onMessageUpdate, chatService]);
+
+		// If message is no longer generating, clear the interval
+		if (!isMessageGenerating && intervalRef.current) {
+			clearInterval(intervalRef.current);
+			intervalRef.current = null;
+			pollingGenerationIdRef.current = null;
+		}
+	}, [isMessageGenerating, message.generationId, message.id, onMessageUpdate, chatService]);
+	// Note: Removed message.generation?.id from dependencies to avoid unnecessary re-runs
 
 	// Cleanup interval on unmount
 	useEffect(() => {

src/app/routes/chat/-hooks/useChat.ts

@@ -309,6 +309,20 @@ export const useChat = (initialChatId?: string, selectedProvider?: string, selec
 					setCurrentChatId(result.id);
 					setIsChatIdValidated(true);
 
+					// Trigger image generation if messages were created
+					if (result.messages) {
+						const assistantMessage = result.messages.find((msg: any) => msg.role === "assistant");
+						if (assistantMessage?.generation?.id) {
+							// Mark as generating before triggering
+							assistantMessage.generation.status = "generating";
+
+							// Fire and forget - don't await, let it run in background
+							chatService.createMessageGenerate({ generationId: assistantMessage.generation.id }).catch((error) => {
+								console.error("Error triggering image generation:", error);
+							});
+						}
+					}
+
 					// Return the chat ID - no need to send another message since createChat already handled it
 					return result.id;
 				} catch (error) {
@@ -418,6 +432,11 @@ export const useChat = (initialChatId?: string, selectedProvider?: string, selec
 
 				// Use returned messages to update the chat data instead of revalidating
 				if (result?.messages) {
+					// Trigger image generation in browser first (not blocked by CF Worker timeout)
+					const assistantMessage = result.messages.find((msg) => msg.role === "assistant");
+					const shouldTriggerGeneration = assistantMessage?.generation?.id;
+
+					// Update UI - mark generation as "generating" to avoid immediate polling
 					currentChatMutate((currentData) => {
 						if (!currentData) return currentData;
 
@@ -426,13 +445,35 @@ export const useChat = (initialChatId?: string, selectedProvider?: string, selec
 						const existingMessages =
 							currentData.messages?.filter((msg: any) => msg.id !== optimisticUserMessage.id) || [];
 
-						// Append new messages from server
+						// Append new messages from server, but mark as "generating" if we're about to trigger
+						const newMessages = shouldTriggerGeneration
+							? result.messages.map((msg: any) => {
+									if (msg.role === "assistant" && msg.generation) {
+										return {
+											...msg,
+											generation: {
+												...msg.generation,
+												status: "generating" as const,
+											},
+										};
+									}
+									return msg;
+								})
+							: result.messages;
+
 						return {
 							...currentData,
-							messages: [...existingMessages, ...result.messages],
+							messages: [...existingMessages, ...newMessages],
 							updatedAt: new Date().toISOString(),
 						};
 					}, false);
+
+					if (shouldTriggerGeneration) {
+						// Fire and forget - don't await, let it run in background
+						chatService.createMessageGenerate({ generationId: assistantMessage!.generation!.id }).catch((error) => {
+							console.error("Error triggering image generation:", error);
+						});
+					}
 				} else {
 					// Fallback: revalidate if no messages returned
 					await currentChatMutate();
@@ -546,7 +587,7 @@ export const useChat = (initialChatId?: string, selectedProvider?: string, selec
 									...msg,
 									generation: {
 										...msg.generation,
-										status: "pending" as const,
+										status: "generating" as const,
 										fileIds: null,
 										errorReason: null,
 									},
@@ -558,7 +599,15 @@ export const useChat = (initialChatId?: string, selectedProvider?: string, selec
 				}, false);
 
 				// Call the API
-				await chatService.regenerateMessage({ messageId });
+				const result = await chatService.regenerateMessage({ messageId });
+
+				// Trigger image generation in browser (not blocked by CF Worker timeout)
+				if (result?.generationId) {
+					// Fire and forget - don't await, let it run in background
+					chatService.createMessageGenerate({ generationId: result.generationId }).catch((error) => {
+						console.error("Error triggering image generation:", error);
+					});
+				}
 
 				// The polling mechanism in ChatMessageItem will handle updating the UI
 			} catch (error) {

src/server/api/routes/chat.ts

@@ -1,5 +1,6 @@
 import {
 	CreateChatSchema,
+	CreateMessageGenerateSchema,
 	CreateMessageSchema,
 	DeleteMessageSchema,
 	GetChatByIdSchema,
@@ -75,6 +76,17 @@ const app = new Hono<Env>()
 
 		return c.json(ok(await chatService.getGenerationStatus(req, { userId: user.id })));
 	})
+	.post("/createMessageGenerate", zValidator("json", CreateMessageGenerateSchema), async (c) => {
+		const user = c.var.user!;
+		const req = c.req.valid("json");
+
+		try {
+			const result = await chatService.createMessageGenerate(req, { userId: user.id, executionCtx: c.executionCtx });
+			return c.json(ok(result));
+		} catch (err: any) {
+			return c.json(error(err.code || "error", err.message || "Failed to generate image"));
+		}
+	})
 	.post("/regenerateMessage", zValidator("json", RegenerateMessageSchema), async (c) => {
 		const user = c.var.user!;
 		const req = c.req.valid("json");