Merge pull request #5347 from continuedev/dallin/apply-state

Edit tool fixes

Author: RomneyDa

core/index.d.ts

@@ -1216,6 +1216,7 @@ export type EditStatus =
   | "done";
 
 export type ApplyStateStatus =
+  | "not-started" // Apply state created but not necessarily streaming
   | "streaming" // Changes are being applied to the file
   | "done" // All changes have been applied, awaiting user to accept/reject
   | "closed"; // All changes have been applied. Note that for new files, we immediately set the status to "closed"

gui/src/components/StyledMarkdownPreview/StepContainerPreToolbar/ApplyActions.tsx

@@ -48,6 +48,14 @@ export function ApplyActions(props: ApplyActionsProps) {
   );
 
   switch (props.applyState ? props.applyState.status : null) {
+    case "not-started":
+      return (
+        <div className="flex select-none items-center rounded bg-zinc-700 pl-2 pr-1">
+          <span className="text-lightgray inline-flex w-min items-center gap-2 text-center text-xs">
+            Pending
+          </span>
+        </div>
+      );
     case "streaming":
       return (
         <div className="flex select-none items-center rounded bg-zinc-700 pl-2 pr-1">
@@ -84,18 +92,15 @@ export function ApplyActions(props: ApplyActionsProps) {
         </div>
       );
     case "closed":
-      if (!hasRejected && props.applyState?.numDiffs === 0) {
-        if (showApplied) {
+      if (isSuccessful) {
+        if (showApplied || props.disableManualApply) {
           return (
             <span className="flex select-none items-center rounded bg-zinc-700 text-slate-400 max-sm:px-0.5 sm:pl-2">
               <span className="max-sm:hidden">Applied</span>
               <CheckIcon className="h-3.5 w-3.5 hover:brightness-125 sm:px-1" />
             </span>
           );
         }
-        if (props.disableManualApply) {
-          return null;
-        }
         return applyButton("Reapply");
       }
     default:

gui/src/components/StyledMarkdownPreview/StepContainerPreToolbar/GeneratingCodeLoader.tsx

@@ -3,22 +3,24 @@ import Spinner from "../../gui/Spinner";
 export interface GeneratingCodeLoaderProps {
   showLineCount: boolean;
   codeBlockContent: string;
+  isPending: boolean;
 }
 
 export function GeneratingCodeLoader({
   showLineCount,
   codeBlockContent,
+  isPending,
 }: GeneratingCodeLoaderProps) {
   const numLinesCodeBlock = codeBlockContent.split("\n").length;
   const linesGeneratedText =
     numLinesCodeBlock === 1
       ? `1 line generated`
-      : `${numLinesCodeBlock} lines generated`;
+      : `${numLinesCodeBlock} lines ${isPending ? "pending" : "generated"}`;
 
   return (
     <span className="text-lightgray inline-flex items-center gap-2">
       {showLineCount ? linesGeneratedText : "Generating"}
-      <Spinner />
+      {!isPending && <Spinner />}
     </span>
   );
 }

gui/src/components/StyledMarkdownPreview/StepContainerPreToolbar/index.tsx

@@ -52,7 +52,7 @@ export interface StepContainerPreToolbarProps {
   codeBlockContent: string;
   language: string | null;
   relativeFilepath?: string;
-  isGeneratingCodeBlock: boolean;
+  isFinalCodeblock: boolean;
   codeBlockIndex: number; // To track which codeblock we are applying
   codeBlockStreamId: string;
   range?: string;
@@ -65,7 +65,7 @@ export function StepContainerPreToolbar({
   codeBlockContent,
   language,
   relativeFilepath,
-  isGeneratingCodeBlock,
+  isFinalCodeblock,
   codeBlockIndex,
   codeBlockStreamId,
   range,
@@ -115,6 +115,9 @@ export function StepContainerPreToolbar({
   const hasFileExtension =
     relativeFilepath && /\.[0-9a-z]+$/i.test(relativeFilepath);
 
+  const isStreaming = useAppSelector((store) => store.session.isStreaming);
+  const isGeneratingCodeBlock = isFinalCodeblock && isStreaming;
+
   // If we are creating a file, we already render that in the button
   // so we don't want to dispaly it twice here
   const displayFilepath = relativeFilepath ?? appliedFileUri;
@@ -272,20 +275,23 @@ export function StepContainerPreToolbar({
         </div>
 
         <div className="flex items-center gap-2.5">
-          {isGeneratingCodeBlock ? (
+          {!isGeneratingCodeBlock && (
+            <div className="xs:flex hidden items-center gap-2.5">
+              <InsertButton onInsert={onClickInsertAtCursor} />
+              <CopyButton text={codeBlockContent} />
+            </div>
+          )}
+
+          {isGeneratingCodeBlock || applyState?.status === "not-started" ? (
             <GeneratingCodeLoader
               showLineCount={!isExpanded}
               codeBlockContent={codeBlockContent}
+              isPending={
+                applyState?.status === "not-started" && !isGeneratingCodeBlock
+              }
             />
           ) : (
-            <>
-              <div className="xs:flex hidden items-center gap-2.5">
-                <InsertButton onInsert={onClickInsertAtCursor} />
-                <CopyButton text={codeBlockContent} />
-              </div>
-
-              {renderActionButtons()}
-            </>
+            renderActionButtons()
           )}
         </div>
       </ToolbarDiv>

gui/src/components/StyledMarkdownPreview/index.tsx

@@ -126,7 +126,7 @@ interface StyledMarkdownPreviewProps {
   itemIndex?: number;
   useParentBackgroundColor?: boolean;
   disableManualApply?: boolean;
-  singleCodeblockStreamId?: string;
+  forceStreamId?: string;
   expandCodeblocks?: boolean;
 }
 
@@ -195,27 +195,16 @@ const StyledMarkdownPreview = memo(function StyledMarkdownPreview(
   const pastFileInfoRef = useUpdatingRef(pastFileInfo);
 
   const isLastItem = useMemo(() => {
-    return props.itemIndex && props.itemIndex === history.length - 1;
+    return props.itemIndex === history.length - 1;
   }, [history.length, props.itemIndex]);
   const isLastItemRef = useUpdatingRef(isLastItem);
 
-  const isStreaming = useAppSelector((store) => store.session.isStreaming);
-  const isStreamingRef = useUpdatingRef(isStreaming);
-
-  const codeblockState = useRef<{ streamId: string; isGenerating: boolean }[]>(
-    [],
-  );
-
+  const codeblockStreamIds = useRef<string[]>([]);
   useEffect(() => {
-    if (props.singleCodeblockStreamId) {
-      codeblockState.current = [
-        {
-          streamId: props.singleCodeblockStreamId,
-          isGenerating: false,
-        },
-      ];
+    if (props.forceStreamId) {
+      codeblockStreamIds.current = [props.forceStreamId];
     }
-  }, [props.singleCodeblockStreamId]);
+  }, [props.forceStreamId, codeblockStreamIds]);
 
   const [reactContent, setMarkdownSource] = useRemark({
     remarkPlugins: [
@@ -310,19 +299,12 @@ const StyledMarkdownPreview = memo(function StyledMarkdownPreview(
 
           const language = getLanguageFromClassName(className);
 
-          const isGeneratingCodeBlock =
-            preChildProps["data-islastcodeblock"] &&
-            isLastItemRef.current &&
-            isStreamingRef.current;
-
-          if (codeblockState.current[codeBlockIndex] === undefined) {
-            codeblockState.current[codeBlockIndex] = {
-              streamId: uuidv4(),
-              isGenerating: isGeneratingCodeBlock,
-            };
-          } else {
-            codeblockState.current[codeBlockIndex].isGenerating =
-              isGeneratingCodeBlock;
+          const isFinalCodeblock =
+            preChildProps["data-islastcodeblock"] && isLastItemRef.current;
+
+          if (codeblockStreamIds.current[codeBlockIndex] === undefined) {
+            codeblockStreamIds.current[codeBlockIndex] =
+              props.forceStreamId ?? uuidv4();
           }
 
           return (
@@ -331,11 +313,9 @@ const StyledMarkdownPreview = memo(function StyledMarkdownPreview(
               codeBlockIndex={codeBlockIndex}
               language={language}
               relativeFilepath={relativeFilePath}
-              isGeneratingCodeBlock={isGeneratingCodeBlock}
+              isFinalCodeblock={isFinalCodeblock}
               range={range}
-              codeBlockStreamId={
-                codeblockState.current[codeBlockIndex].streamId
-              }
+              codeBlockStreamId={codeblockStreamIds.current[codeBlockIndex]}
               expanded={props.expandCodeblocks}
               disableManualApply={props.disableManualApply}
             >
@@ -381,7 +361,7 @@ const StyledMarkdownPreview = memo(function StyledMarkdownPreview(
   const codeWrapState = uiConfig?.codeWrap ? "pre-wrap" : "pre";
   return (
     <StyledMarkdown
-      contentEditable='false'
+      contentEditable="false"
       fontSize={getFontSize()}
       whiteSpace={codeWrapState}
       bgColor={props.useParentBackgroundColor ? "" : vscBackground}

gui/src/components/mainInput/InputToolbar.tsx

@@ -4,7 +4,10 @@ import { modelSupportsImages, modelSupportsTools } from "core/llm/autodetect";
 import { useRef } from "react";
 import { useAppDispatch, useAppSelector } from "../../redux/hooks";
 import { selectUseActiveFile } from "../../redux/selectors";
-import { selectCurrentToolCall } from "../../redux/selectors/selectCurrentToolCall";
+import {
+  selectCurrentToolCall,
+  selectCurrentToolCallApplyState,
+} from "../../redux/selectors/selectCurrentToolCall";
 import { selectSelectedChatModel } from "../../redux/slices/configSlice";
 import {
   selectHasCodeToEdit,
@@ -53,15 +56,16 @@ function InputToolbar(props: InputToolbarProps) {
   const hasCodeToEdit = useAppSelector(selectHasCodeToEdit);
   const toolCallState = useAppSelector(selectCurrentToolCall);
   const isEditModeAndNoCodeToEdit = isInEditMode && !hasCodeToEdit;
-  const activeToolCallStreamId = useAppSelector(
-    (store) => store.session.activeToolStreamId,
+  const currentToolCallApplyState = useAppSelector(
+    selectCurrentToolCallApplyState,
   );
 
   const isEnterDisabled =
     props.disabled ||
     isEditModeAndNoCodeToEdit ||
     toolCallState?.status === "generated" ||
-    !!activeToolCallStreamId;
+    (currentToolCallApplyState &&
+      currentToolCallApplyState.status !== "closed");
 
   const toolsSupported = defaultModel && modelSupportsTools(defaultModel);
 

gui/src/components/mainInput/Lump/LumpToolbar.tsx

@@ -109,7 +109,9 @@ export function LumpToolbar() {
   if (toolCallState?.status === "generated") {
     return (
       <Container>
-        <GeneratingIndicator />
+        <div className="flex flex-row items-center pb-0.5 pr-1 text-xs text-gray-400">
+          <span className="hidden sm:flex">Pending tool call</span>
+        </div>
 
         <div className="flex gap-2 pb-0.5">
           <StopButton

gui/src/hooks/useSetup.ts

@@ -10,6 +10,7 @@ import {
   setSelectedProfile,
 } from "../redux";
 import { useAppDispatch, useAppSelector } from "../redux/hooks";
+import { selectCurrentToolCallApplyState } from "../redux/selectors/selectCurrentToolCall";
 import {
   selectSelectedChatModel,
   setConfigResult,
@@ -252,41 +253,39 @@ function useSetup() {
     dispatch(updateIndexingStatus(data));
   });
 
-  const activeToolStreamId = useAppSelector(
-    (store) => store.session.activeToolStreamId,
+  const currentToolCallApplyState = useAppSelector(
+    selectCurrentToolCallApplyState,
   );
   useWebviewListener(
     "updateApplyState",
     async (state) => {
       dispatch(updateApplyState(state));
-      const lastHistoryMsg = history.at(-1);
-      const [streamId, toolCallId] = activeToolStreamId ?? [];
+
+      // Handle apply status updates that are associated with current tool call
       if (
-        toolCallId &&
         state.status === "closed" &&
-        lastHistoryMsg?.toolCallState?.toolCallId === toolCallId
+        currentToolCallApplyState &&
+        currentToolCallApplyState.streamId === state.streamId
       ) {
-        if (state.streamId === streamId) {
-          // const output: ContextItem = {
-          //   name: "Edit tool output",
-          //   content: "Completed edit",
-          //   description: "",
-          // };
-          dispatch(
-            acceptToolCall({
-              toolCallId,
-            }),
-          );
-          // dispatch(setToolCallOutput([]));
-          dispatch(
-            streamResponseAfterToolCall({
-              toolCallId,
-            }),
-          );
-        }
+        // const output: ContextItem = {
+        //   name: "Edit tool output",
+        //   content: "Completed edit",
+        //   description: "",
+        // };
+        dispatch(
+          acceptToolCall({
+            toolCallId: currentToolCallApplyState.toolCallId!,
+          }),
+        );
+        // dispatch(setToolCallOutput([]));
+        dispatch(
+          streamResponseAfterToolCall({
+            toolCallId: currentToolCallApplyState.toolCallId!,
+          }),
+        );
       }
     },
-    [activeToolStreamId, history],
+    [currentToolCallApplyState, history],
   );
 }
 

gui/src/pages/gui/Chat.tsx

@@ -30,7 +30,10 @@ import { IdeMessengerContext } from "../../context/IdeMessenger";
 import { useWebviewListener } from "../../hooks/useWebviewListener";
 import { useAppDispatch, useAppSelector } from "../../redux/hooks";
 import { selectUseHub } from "../../redux/selectors";
-import { selectCurrentToolCall } from "../../redux/selectors/selectCurrentToolCall";
+import {
+  selectCurrentToolCall,
+  selectCurrentToolCallApplyState,
+} from "../../redux/selectors/selectCurrentToolCall";
 import { selectSelectedChatModel } from "../../redux/slices/configSlice";
 import { submitEdit } from "../../redux/slices/editModeState";
 import {
@@ -152,6 +155,10 @@ export function Chat() {
 
   const { widget, highlights } = useFindWidget(stepsDivRef);
 
+  const currentToolCallApplyState = useAppSelector(
+    selectCurrentToolCallApplyState,
+  );
+
   const sendInput = useCallback(
     (
       editorState: JSONContent,
@@ -164,6 +171,14 @@ export function Chat() {
           "Cannot submit message while awaiting tool confirmation",
         );
       }
+      if (
+        currentToolCallApplyState &&
+        currentToolCallApplyState.status !== "closed"
+      ) {
+        return console.error(
+          "Cannot submit message while awaiting tool call apply",
+        );
+      }
       if (selectedChatModel?.provider === "free-trial") {
         const newCount = incrementFreeTrialCount();
 
@@ -352,6 +367,7 @@ export function Chat() {
                           toolCallState={item.toolCallState!}
                           toolCall={toolCall}
                           output={history[index + 1]?.contextItems}
+                          historyIndex={index}
                         />
                       </div>
                     );