🛠️ fix: Custom Endpoint issues, Improve SSE Response Handling (#1510)

* fix(custom): prevent presets using removed custom endpoints from causing frontend errors

* refactor(abortMiddleware): send 204 status when abortController is not found/active, set expected header `application/json` when not set

* fix(useSSE): general improvements:
- Add endpointType to fetch URL in useSSE hook
- use EndpointURLs enum
- handle 204 response by setting `data` to initiated response
- add better error handling UX, make clear when there is an explicit error

api/server/middleware/abortMiddleware.js

@@ -14,7 +14,7 @@ async function abortMessage(req, res) {
   }
 
   if (!abortControllers.has(abortKey) && !res.headersSent) {
-    return res.status(404).send({ message: 'Request not found' });
+    return res.status(204).send({ message: 'Request not found' });
   }
 
   const { abortController } = abortControllers.get(abortKey);
@@ -26,6 +26,8 @@ async function abortMessage(req, res) {
     return sendMessage(res, finalEvent);
   }
 
+  res.setHeader('Content-Type', 'application/json');
+
   res.send(JSON.stringify(finalEvent));
 }
 

client/src/components/Chat/Landing.tsx

@@ -23,6 +23,7 @@ export default function Landing({ Header }: { Header?: ReactNode }) {
   const endpointType = getEndpointField(endpointsConfig, endpoint, 'type');
   const iconURL = getEndpointField(endpointsConfig, endpoint, 'iconURL');
   const iconKey = endpointType ? 'unknown' : endpoint ?? 'unknown';
+  const Icon = icons[iconKey];
 
   return (
     <div className="relative h-full">
@@ -31,7 +32,8 @@ export default function Landing({ Header }: { Header?: ReactNode }) {
         <div className="mb-3 h-[72px] w-[72px]">
           <div className="gizmo-shadow-stroke relative flex h-full items-center justify-center rounded-full bg-white text-black">
             {endpoint &&
-              icons[iconKey]({
+              Icon &&
+              Icon({
                 size: 41,
                 context: 'landing',
                 className: 'h-2/3 w-2/3',

client/src/components/Chat/Menus/Endpoints/MenuItem.tsx

@@ -82,7 +82,7 @@ const MenuItem: FC<MenuItemProps> = ({
         <div className="flex grow items-center justify-between gap-2">
           <div>
             <div className="flex items-center gap-2">
-              {
+              {Icon && (
                 <Icon
                   size={18}
                   endpoint={endpoint}
@@ -90,7 +90,7 @@ const MenuItem: FC<MenuItemProps> = ({
                   className="icon-md shrink-0 dark:text-white"
                   iconURL={getEndpointField(endpointsConfig, endpoint, 'iconURL')}
                 />
-              }
+              )}
               <div>
                 {title}
                 <div className="text-token-text-tertiary">{description}</div>

client/src/components/Chat/Menus/Presets/PresetItems.tsx

@@ -97,7 +97,8 @@ const PresetItems: FC<{
 
             const iconKey = getEndpointField(endpointsConfig, preset.endpoint, 'type')
               ? 'unknown'
-              : preset.endpoint ?? 'unknown';
+              : preset.endpointType ?? preset.endpoint ?? 'unknown';
+            const Icon = icons[iconKey];
 
             return (
               <Close asChild key={`preset-${preset.presetId}`}>
@@ -109,12 +110,15 @@ const PresetItems: FC<{
                       title={getPresetTitle(preset)}
                       disableHover={true}
                       onClick={() => onSelectPreset(preset)}
-                      icon={icons[iconKey]({
+                      icon={
+                        Icon &&
+                        Icon({
                           context: 'menu-item',
                           iconURL: getEndpointField(endpointsConfig, preset.endpoint, 'iconURL'),
                           className: 'icon-md mr-1 dark:text-white',
                           endpoint: preset.endpoint,
-                      })}
+                        })
+                      }
                       selected={false}
                       data-testid={`preset-item-${preset}`}
                     >

client/src/hooks/useSSE.ts

@@ -4,6 +4,7 @@ import { useParams } from 'react-router-dom';
 import {
   /* @ts-ignore */
   SSE,
+  EndpointURLs,
   createPayload,
   tMessageSchema,
   tConvoUpdateSchema,
@@ -268,10 +269,11 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
 
   const abortConversation = (conversationId = '', submission: TSubmission) => {
     console.log(submission);
-    const { endpoint } = submission?.conversation || {};
+    const { endpoint: _endpoint, endpointType } = submission?.conversation || {};
+    const endpoint = endpointType ?? _endpoint;
     let res: Response;
 
-    fetch(`/api/ask/${endpoint}/abort`, {
+    fetch(`${EndpointURLs[endpoint ?? '']}/abort`, {
       method: 'POST',
       headers: {
         'Content-Type': 'application/json',
@@ -283,7 +285,29 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
     })
       .then((response) => {
         res = response;
+        // Check if the response is JSON
+        const contentType = response.headers.get('content-type');
+        if (contentType && contentType.includes('application/json')) {
           return response.json();
+        } else if (response.status === 204) {
+          const responseMessage = {
+            ...submission.initialResponse,
+            text: submission.initialResponse.text.replace(
+              '<span className="result-streaming">█</span>',
+              '',
+            ),
+          };
+
+          return {
+            requestMessage: submission.message,
+            responseMessage: responseMessage,
+            conversation: submission.conversation,
+          };
+        } else {
+          throw new Error(
+            'Unexpected response from server; Status: ' + res.status + ' ' + res.statusText,
+          );
+        }
       })
       .then((data) => {
         console.log('aborted', data);
@@ -295,6 +319,22 @@ export default function useSSE(submission: TSubmission | null, index = 0) {
       .catch((error) => {
         console.error('Error aborting request');
         console.error(error);
+        const convoId = conversationId ?? v4();
+
+        const text =
+          submission.initialResponse?.text?.length > 45 ? submission.initialResponse?.text : '';
+
+        const errorMessage = {
+          ...submission,
+          ...submission.initialResponse,
+          text: text ?? error.message ?? 'Error cancelling request',
+          unfinished: !!text.length,
+          error: true,
+        };
+
+        const errorResponse = tMessageSchema.parse(errorMessage);
+        setMessages([...submission.messages, submission.message, errorResponse]);
+        newConversation({ template: { conversationId: convoId } });
         setIsSubmitting(false);
       });
     return;