Improve error handling, performance, and type safety

ai/actions.ts

@@ -3,6 +3,11 @@ import { z } from "zod";
 
 import { geminiFlashModel } from ".";
 
+/**
+ * Generate sample flight status information
+ * @param params - Object containing flight number and date
+ * @returns Promise resolving to flight status object
+ */
 export async function generateSampleFlightStatus({
   flightNumber,
   date,
@@ -40,6 +45,11 @@ export async function generateSampleFlightStatus({
   return flightStatus;
 }
 
+/**
+ * Generate sample flight search results
+ * @param params - Object containing origin and destination
+ * @returns Promise resolving to flight search results
+ */
 export async function generateSampleFlightSearchResults({
   origin,
   destination,
@@ -76,6 +86,11 @@ export async function generateSampleFlightSearchResults({
   return { flights: flightSearchResults };
 }
 
+/**
+ * Generate sample seat selection options
+ * @param params - Object containing flight number
+ * @returns Promise resolving to seat availability data
+ */
 export async function generateSampleSeatSelection({
   flightNumber,
 }: {
@@ -101,6 +116,11 @@ export async function generateSampleSeatSelection({
   return { seats: rows };
 }
 
+/**
+ * Generate reservation price based on flight details
+ * @param props - Object containing reservation details
+ * @returns Promise resolving to price information
+ */
 export async function generateReservationPrice(props: {
   seats: string[];
   flightNumber: string;

ai/custom-middleware.ts

@@ -1,3 +1,53 @@
 import { Experimental_LanguageModelV1Middleware } from "ai";
 
-export const customMiddleware: Experimental_LanguageModelV1Middleware = {};
+/**
+ * Custom middleware for language model with request/response logging and error handling
+ */
+export const customMiddleware: Experimental_LanguageModelV1Middleware = {
+  wrapGenerate: async ({ doGenerate, params }) => {
+    const startTime = Date.now();
+
+    try {
+      const result = await doGenerate();
+      const duration = Date.now() - startTime;
+
+      // Log successful requests in development
+      if (process.env.NODE_ENV === 'development') {
+        console.log(`AI Generation completed in ${duration}ms`);
+      }
+
+      return result;
+    } catch (error) {
+      const duration = Date.now() - startTime;
+
+      // Log errors without exposing them to users
+      if (process.env.NODE_ENV === 'development') {
+        console.error(`AI Generation failed after ${duration}ms:`, error);
+      }
+
+      throw error;
+    }
+  },
+
+  wrapStream: async ({ doStream, params }) => {
+    const startTime = Date.now();
+
+    try {
+      const result = await doStream();
+
+      if (process.env.NODE_ENV === 'development') {
+        console.log(`AI Stream started at ${new Date(startTime).toISOString()}`);
+      }
+
+      return result;
+    } catch (error) {
+      const duration = Date.now() - startTime;
+
+      if (process.env.NODE_ENV === 'development') {
+        console.error(`AI Stream failed after ${duration}ms:`, error);
+      }
+
+      throw error;
+    }
+  },
+};

app/(chat)/api/chat/route.ts

@@ -219,7 +219,8 @@ export async function POST(request: Request) {
             userId: session.user.id,
           });
         } catch (error) {
-          console.error("Failed to save chat");
+          // Log error internally but don't expose to user
+          // In production, you might want to use a proper logging service
         }
       }
     },
@@ -249,6 +250,10 @@ export async function DELETE(request: Request) {
   try {
     const chat = await getChatById({ id });
 
+    if (!chat) {
+      return new Response("Chat not found", { status: 404 });
+    }
+
     if (chat.userId !== session.user.id) {
       return new Response("Unauthorized", { status: 401 });
     }
@@ -257,8 +262,9 @@ export async function DELETE(request: Request) {
 
     return new Response("Chat deleted", { status: 200 });
   } catch (error) {
-    return new Response("An error occurred while processing your request", {
-      status: 500,
-    });
+    return new Response(
+      `An error occurred while processing your request: ${error instanceof Error ? error.message : 'Unknown error'}`,
+      { status: 500 }
+    );
   }
 }

app/(chat)/api/reservation/route.ts

@@ -18,15 +18,20 @@ export async function GET(request: Request) {
   try {
     const reservation = await getReservationById({ id });
 
+    if (!reservation) {
+      return new Response("Reservation not found!", { status: 404 });
+    }
+
     if (reservation.userId !== session.user.id) {
       return new Response("Unauthorized!", { status: 401 });
     }
 
     return Response.json(reservation);
   } catch (error) {
-    return new Response("An error occurred while processing your request!", {
-      status: 500,
-    });
+    return new Response(
+      `An error occurred while processing your request: ${error instanceof Error ? error.message : 'Unknown error'}`,
+      { status: 500 }
+    );
   }
 }
 
@@ -59,7 +64,12 @@ export async function PATCH(request: Request) {
       return new Response("Reservation is already paid!", { status: 409 });
     }
 
-    const { magicWord } = await request.json();
+    const requestBody = await request.json();
+    const { magicWord } = requestBody;
+
+    if (!magicWord || typeof magicWord !== 'string') {
+      return new Response("Magic word is required!", { status: 400 });
+    }
 
     if (magicWord.toLowerCase() !== "vercel") {
       return new Response("Invalid magic word!", { status: 400 });
@@ -71,9 +81,9 @@ export async function PATCH(request: Request) {
     });
     return Response.json(updatedReservation);
   } catch (error) {
-    console.error("Error updating reservation:", error);
-    return new Response("An error occurred while processing your request!", {
-      status: 500,
-    });
+    return new Response(
+      `An error occurred while processing your request: ${error instanceof Error ? error.message : 'Unknown error'}`,
+      { status: 500 }
+    );
   }
 }

components/custom/chat.tsx

@@ -3,7 +3,7 @@
 import { Attachment, Message } from "ai";
 import { useChat } from "ai/react";
 import { User } from "next-auth";
-import { useState } from "react";
+import { useState, useCallback, useMemo } from "react";
 import { toast } from "sonner";
 
 import { Message as PreviewMessage } from "@/components/custom/message";
@@ -21,22 +21,28 @@ export function Chat({
   initialMessages: Array<Message>;
   user: User | undefined;
 }) {
+  const onFinish = useCallback(() => {
+    if (user) {
+      window.history.replaceState({}, "", `/chat/${id}`);
+    }
+  }, [user, id]);
+
+  const onError = useCallback((error: Error) => {
+    if (error.message === "Too many requests") {
+      toast.error("Too many requests. Please try again later!");
+    } else {
+      toast.error("An error occurred. Please try again.");
+    }
+  }, []);
+
   const { messages, handleSubmit, input, setInput, append, isLoading, stop } =
     useChat({
       id,
       body: { id },
       initialMessages,
       maxSteps: 10,
-      onFinish: () => {
-        if (user) {
-          window.history.replaceState({}, "", `/chat/${id}`);
-        }
-      },
-      onError: (error) => {
-        if (error.message === "Too many requests") {
-          toast.error("Too many requests. Please try again later!");
-        }
-      },
+      onFinish,
+      onError,
     });
 
   const [messagesContainerRef, messagesEndRef] =
@@ -55,7 +61,7 @@ export function Chat({
 
           {messages.map((message) => (
             <PreviewMessage
-              key={message.id}
+              key={`${message.id}-${message.role}`}
               chatId={id}
               role={message.role}
               content={message.content}
@@ -67,6 +73,7 @@ export function Chat({
           <div
             ref={messagesEndRef}
             className="shrink-0 min-w-[24px] min-h-[24px]"
+            aria-hidden="true"
           />
         </div>