From ae2181460b0c384e8377673e0a039089cf54f418 Mon Sep 17 00:00:00 2001
From: Maciej Jastrzebski <mdjastrzebski@gmail.com>
Date: Mon, 1 Apr 2024 23:32:46 +0200
Subject: [PATCH] fix: handle inference error gracefully

---
 src/commands/chat/state/actions.ts | 13 +++++++++++++
 src/commands/chat/ui/ChatUi.tsx    | 22 +++++++++++-----------
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/commands/chat/state/actions.ts b/src/commands/chat/state/actions.ts
index 3e83b4e..a5f46a7 100644
--- a/src/commands/chat/state/actions.ts
+++ b/src/commands/chat/state/actions.ts
@@ -38,6 +38,19 @@ export function addProgramMessage(text: string, level: MessageLevel = 'info') {
   });
 }
 
+export function removeUnansweredUserContextMessage() {
+  useChatState.setState((state) => {
+    const lastMessage = state.contextMessages.at(-1);
+    if (lastMessage?.role !== 'user') {
+      return state;
+    }
+
+    return {
+      contextMessages: state.contextMessages.slice(0, state.contextMessages.length - 1),
+    };
+  });
+}
+
 export function forgetContextMessages() {
   useChatState.setState((state) => {
     return {
diff --git a/src/commands/chat/ui/ChatUi.tsx b/src/commands/chat/ui/ChatUi.tsx
index dfdb38a..40e696d 100644
--- a/src/commands/chat/ui/ChatUi.tsx
+++ b/src/commands/chat/ui/ChatUi.tsx
@@ -9,6 +9,7 @@ import {
   addProgramMessage,
   addUserMessage,
   hideActiveView,
+  removeUnansweredUserContextMessage,
 } from '../state/actions.js';
 import { useChatState } from '../state/state.js';
 import { texts } from '../texts.js';
@@ -38,19 +39,15 @@ export function ChatUi() {
   }, []);
 
   const handleSubmit = (message: string) => {
-    try {
-      hideActiveView();
-
-      const isCommand = processChatCommand(message);
-      if (isCommand) {
-        return;
-      }
+    hideActiveView();
 
-      addUserMessage(message);
-      void fetchAiResponse();
-    } catch (error) {
-      addProgramMessage(`Error: ${extractErrorMessage(error)}`, 'error');
+    const isCommand = processChatCommand(message);
+    if (isCommand) {
+      return;
     }
+
+    addUserMessage(message);
+    void fetchAiResponse();
   };
 
   const showInput = !isLoading && !shouldExit;
@@ -97,6 +94,9 @@ function useAiResponse() {
         addAiResponse(response);
       }
     } catch (error) {
+      // We cannot leave unanswered user message in context, as there is no AI response for it.
+      // Inference APIs require that user and assistant messages should be alternating.
+      removeUnansweredUserContextMessage();
       addProgramMessage(`Error: ${extractErrorMessage(error)}`, 'error');
     } finally {
       setLoading(false);