upon provider error, reset to easily resubmit last user message

Author: dlants

node/chat/chat.spec.ts

@@ -364,7 +364,8 @@ describe("tea/chat.spec.ts", () => {
     await withDriver(async (driver) => {
       await driver.showSidebar();
 
-      await driver.inputMagentaText("Generate some text for me");
+      const userPrompt = "Generate some text for me";
+      await driver.inputMagentaText(userPrompt);
       await driver.send();
 
       const pendingRequest = await driver.mockAnthropic.awaitPendingRequest();
@@ -383,11 +384,8 @@ describe("tea/chat.spec.ts", () => {
         "Error Connection error during response",
       );
 
-      const displayText = await driver.getDisplayBufferText();
-      expect(displayText).toContain("# user:");
-      expect(displayText).toContain("Generate some text for me");
-
-      expect(displayText).toContain("Error Connection error");
+      // Check that the input buffer is pre-populated with the last user message
+      await driver.assertInputBufferContains(userPrompt);
     });
   });
 });

node/chat/chat.ts

@@ -40,6 +40,7 @@ export type ConversationState =
   | {
       state: "error";
       error: Error;
+      lastAssistantMessage?: Message.Model;
     };
 
 export type Model = {
@@ -98,6 +99,10 @@ export type Msg =
     }
   | {
       type: "show-message-debug-info";
+    }
+  | {
+      type: "sidebar-setup-resubmit";
+      lastUserMessage: string;
     };
 
 export function init({ nvim, lsp }: { nvim: Nvim; lsp: Lsp }) {
@@ -173,28 +178,67 @@ export function init({ nvim, lsp }: { nvim: Nvim; lsp: Lsp }) {
 
       case "conversation-state": {
         model.conversation = msg.conversation;
-        if (msg.conversation.state == "stopped") {
-          const lastMessage = model.messages[model.messages.length - 1];
-          if (lastMessage?.role === "assistant") {
-            lastMessage.parts.push({
-              type: "stop-msg",
-              stopReason: msg.conversation.stopReason,
-              usage: msg.conversation.usage,
-            });
+
+        switch (msg.conversation.state) {
+          case "stopped": {
+            const lastMessage = model.messages[model.messages.length - 1];
+            if (lastMessage?.role === "assistant") {
+              lastMessage.parts.push({
+                type: "stop-msg",
+                stopReason: msg.conversation.stopReason,
+                usage: msg.conversation.usage,
+              });
+            }
+            return [model, maybeAutorespond(model)];
           }
-        }
 
-        if (msg.conversation.state == "error") {
-          const lastMessage = model.messages[model.messages.length - 1];
-          if (lastMessage?.role == "assistant") {
-            // get rid of the latest assistant message so we can re-submit.
-            model.messages.pop();
+          case "error": {
+            const lastAssistantMessage =
+              model.messages[model.messages.length - 1];
+            if (lastAssistantMessage?.role == "assistant") {
+              model.messages.pop();
+
+              // save the last message so we can show a nicer error message.
+              (
+                model.conversation as Extract<
+                  ConversationState,
+                  { state: "error" }
+                >
+              ).lastAssistantMessage = lastAssistantMessage;
+            }
+
+            const lastUserMessage = model.messages[model.messages.length - 1];
+            if (lastUserMessage?.role == "user") {
+              model.messages.pop();
+              const sidebarResubmitSetupThunk: Thunk<Msg> = (dispatch) =>
+                new Promise((resolve) => {
+                  dispatch({
+                    type: "sidebar-setup-resubmit",
+                    lastUserMessage: lastUserMessage.parts
+                      .map((p) => (p.type == "text" ? p.text : ""))
+                      .join(""),
+                  });
+                  resolve();
+                });
+
+              return [model, sidebarResubmitSetupThunk];
+            }
+            return [model];
           }
-        }
 
-        return [model, maybeAutorespond(model)];
+          case "message-in-flight":
+            return [model];
+
+          default:
+            return assertUnreachable(msg.conversation);
+        }
       }
 
+      case "sidebar-setup-resubmit":
+        // this action is really just there so the parent (magenta app) can observe it and manipulate the sidebar
+        // accordingly
+        return [model];
+
       case "send-message": {
         const lastMessage = model.messages[model.messages.length - 1];
         if (lastMessage && lastMessage.role == "user") {
@@ -504,7 +548,8 @@ export function init({ nvim, lsp }: { nvim: Nvim; lsp: Lsp }) {
   }) => {
     if (
       model.messages.length == 0 &&
-      Object.keys(model.contextManager.files).length == 0
+      Object.keys(model.contextManager.files).length == 0 &&
+      model.conversation.state == "stopped"
     ) {
       return d`${LOGO}`;
     }
@@ -532,7 +577,16 @@ export function init({ nvim, lsp }: { nvim: Nvim; lsp: Lsp }) {
           ? withBindings(d`Stopped (${model.conversation.stopReason})`, {
               "<CR>": () => dispatch({ type: "show-message-debug-info" }),
             })
-          : d`Error ${model.conversation.error.message}${model.conversation.error.stack ? "\n" + model.conversation.error.stack : ""}`
+          : d`Error ${model.conversation.error.message}${model.conversation.error.stack ? "\n" + model.conversation.error.stack : ""}${
+              model.conversation.lastAssistantMessage
+                ? "\n\nLast assistant message:\n" +
+                  JSON.stringify(
+                    model.conversation.lastAssistantMessage,
+                    null,
+                    2,
+                  )
+                : ""
+            }`
     }${
       model.conversation.state != "message-in-flight" &&
       !contextManagerModel.isContextEmpty(model.contextManager)

node/magenta.ts

@@ -64,7 +64,29 @@ export class Magenta {
     this.chatApp = TEA.createApp({
       nvim: this.nvim,
       initialModel: this.chatModel.initModel(),
-      update: (msg, model) => this.chatModel.update(msg, model, { nvim }),
+      update: (msg, model) => {
+        if (msg.type == "sidebar-setup-resubmit") {
+          if (
+            this.sidebar &&
+            this.sidebar.state &&
+            this.sidebar.state.inputBuffer
+          ) {
+            this.sidebar.state.inputBuffer
+              .setLines({
+                start: 0,
+                end: -1,
+                lines: msg.lastUserMessage.split("\n") as Line[],
+              })
+              .catch((error) => {
+                this.nvim.logger?.error(
+                  `Error updating sidebar input: ${error}`,
+                );
+              });
+          }
+        }
+
+        return this.chatModel.update(msg, model, { nvim });
+      },
       View: this.chatModel.view,
     });
 

node/providers/anthropic.ts

@@ -458,6 +458,7 @@ export class AnthropicProvider implements Provider {
       }
     };
 
+    let requestActive = true;
     try {
       this.request = this.client.messages
         .stream(
@@ -466,10 +467,16 @@ export class AnthropicProvider implements Provider {
           ) as Anthropic.Messages.MessageStreamParams,
         )
         .on("text", (text: string) => {
+          if (!requestActive) {
+            return;
+          }
           buf.push(text);
           flushBuffer();
         })
         .on("inputJson", (_delta, snapshot) => {
+          if (!requestActive) {
+            return;
+          }
           this.nvim.logger?.debug(
             `anthropic stream inputJson: ${JSON.stringify(snapshot)}`,
           );
@@ -559,12 +566,20 @@ export class AnthropicProvider implements Provider {
         usage.cacheMisses = response.usage.cache_creation_input_tokens;
       }
 
+      if (!requestActive) {
+        throw new Error(`request no longer active`);
+      }
+
       return {
         toolRequests,
         stopReason: response.stop_reason || "end_turn",
         usage,
       };
     } finally {
+      requestActive = false;
+      if (this.request) {
+        this.request.abort();
+      }
       this.request = undefined;
     }
   }