fix(chat): cancel run (#168)

Signed-off-by: Petr Kadlec <[email protected]>
i-am-bee · Jan 7, 2025 · f8b1201 · f8b1201
1 parent d794c1f
commit f8b1201
Show file tree

Hide file tree

Showing 8 changed files with 231 additions and 100 deletions.
diff --git a/src/middleware.ts b/src/middleware.ts
@@ -101,7 +101,7 @@ const getCspHeader = (nonce: string) => {
     base-uri 'self';
     form-action 'self';
     frame-ancestors 'none';
-    upgrade-insecure-requests;
+    ${process.env.NODE_ENV === 'production' ? 'upgrade-insecure-requests;' : ''}
   `;
 
   return cspHeader.replace(/\s{2,}/g, ' ').trim();

diff --git a/src/modules/chat/api/useDeleteMessage.ts b/src/modules/chat/api/useDeleteMessage.ts
@@ -0,0 +1,48 @@
+/**
+ * Copyright 2024 IBM Corp.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { useMutation, useQueryClient } from '@tanstack/react-query';
+import { deleteMessage } from '@/app/api/threads-messages';
+import { useProjectContext } from '@/layout/providers/ProjectProvider';
+import { messagesWithFilesQuery } from '../queries';
+
+export function useDeleteMessage() {
+  const { project, organization } = useProjectContext();
+  const queryClient = useQueryClient();
+
+  const mutation = useMutation({
+    mutationFn: ({ threadId, messageId }: DeleteMutationParams) =>
+      deleteMessage(organization.id, project.id, threadId, messageId),
+    onSuccess: (_, { threadId }) => {
+      queryClient.invalidateQueries({
+        queryKey: [
+          messagesWithFilesQuery(
+            organization.id,
+            project.id,
+            threadId,
+          ).queryKey.at(0),
+        ],
+      });
+    },
+  });
+
+  return mutation;
+}
+
+interface DeleteMutationParams {
+  threadId: string;
+  messageId: string;
+}
diff --git a/src/modules/chat/assistant-plan/PlanStep.tsx b/src/modules/chat/assistant-plan/PlanStep.tsx
@@ -57,6 +57,7 @@ import { ToolApprovalValue } from '../types';
 import classes from './PlanStep.module.scss';
 import { useUserSetting } from '@/layout/hooks/useUserSetting';
 import { useProjectContext } from '@/layout/providers/ProjectProvider';
+import { getToolApproval, getToolReferenceFromToolCall } from './utils';
 
 interface Props {
   step: AssistantPlanStep;
@@ -91,18 +92,7 @@ export function PlanStep({ step, toolCall, allStepsDone }: Props) {
   const status = getStepStatus(step, run);
 
   const toolKey = toolCall.type;
-  const tool =
-    toolKey === 'system'
-      ? {
-          type: toolKey,
-          id: toolCall.toolId,
-        }
-      : toolKey === 'user'
-        ? {
-            type: toolKey,
-            id: toolCall.toolId,
-          }
-        : { type: toolKey, id: toolKey };
+  const tool = getToolReferenceFromToolCall(toolCall);
 
   const { toolName, toolIcon } = useToolInfo({
     organization,
@@ -115,18 +105,10 @@ export function PlanStep({ step, toolCall, allStepsDone }: Props) {
   const { setExpandedStep } = useExpandedStepActions();
   const expanded = expandedState?.stepId === step.id;
 
-  const toolApproval = (
-    run?.status === 'requires_action' &&
-    run.required_action?.type === 'submit_tool_approvals'
-      ? run.required_action.submit_tool_approvals.tool_calls
-      : []
-  )
-    .map((tool) => ({
-      id: tool.id,
-      toolId: getToolApprovalId(tool),
-      type: tool.type,
-    }))
-    .find((toolApproval) => toolApproval.toolId === tool.id);
+  const toolApproval = useMemo(
+    () => getToolApproval(toolCall, run),
+    [run, toolCall],
+  );
 
   const handleToolApprovalSubmit = (value: ToolApprovalValue) => {
     if (value === 'always' && thread && toolApproval?.toolId) {

diff --git a/src/modules/chat/assistant-plan/PlanWithSources.tsx b/src/modules/chat/assistant-plan/PlanWithSources.tsx
@@ -34,7 +34,11 @@ import { runStepsQuery } from '../queries';
 import { BotChatMessage } from '../types';
 import { PlanView } from './PlanView';
 import classes from './PlanWithSources.module.scss';
-import { updatePlanWithRunStep } from './utils';
+import {
+  getToolApproval,
+  getToolReferenceFromToolCall,
+  updatePlanWithRunStep,
+} from './utils';
 import { TraceData } from '../trace/types';
 import { useBuildTraceData } from '../trace/useBuildTraceData';
 import { TraceInfoView } from '../trace/TraceInfoView';
@@ -108,6 +112,25 @@ function PlanWithSourcesComponent({ message, inView }: Props) {
     threadId: thread?.id,
   });
 
+  useEffect(() => {
+    if (isOpen) return;
+
+    const steps = plan.steps || [];
+    const hasToolApprovalRequest = steps.some((step) => {
+      if (step.status === 'in_progress') {
+        const lastToolCall = step.toolCalls.at(-1);
+        if (lastToolCall) {
+          const toolApproval = getToolApproval(lastToolCall, message.run);
+          return Boolean(toolApproval);
+        }
+      }
+      return false;
+    });
+    if (hasToolApprovalRequest) {
+      setIsOpen(true);
+    }
+  }, [isOpen, message.run, plan.steps]);
+
   useEffect(() => {
     setShowButton(!debugMode && !plan.pending);
   }, [debugMode, plan.pending]);

diff --git a/src/modules/chat/assistant-plan/utils.ts b/src/modules/chat/assistant-plan/utils.ts
@@ -21,13 +21,15 @@ import {
   RunStepDeltaEventResponse,
   StepToolCall,
   SystemToolResult,
+  ThreadRun,
 } from '@/app/api/threads-runs/types';
 import {
   isArXivToolResult,
   isRunStepDeltaEventDetailsThought,
   isWebSearchToolResult,
   isWikipediaToolResult,
 } from '@/app/api/threads-runs/utils';
+import { getToolApprovalId } from '@/modules/tools/utils';
 import { isNotNull } from '@/utils/helpers';
 import { v4 as uuid } from 'uuid';
 
@@ -219,3 +221,37 @@ export function updatePlanWithRunStepDelta(
   }
   return plan;
 }
+
+export function getToolReferenceFromToolCall(toolCall: StepToolCall) {
+  const toolKey = toolCall.type;
+  return toolKey === 'system'
+    ? {
+        type: toolKey,
+        id: toolCall.toolId,
+      }
+    : toolKey === 'user'
+      ? {
+          type: toolKey,
+          id: toolCall.toolId,
+        }
+      : { type: toolKey, id: toolKey };
+}
+
+export function getToolApproval(toolCall: StepToolCall, run?: ThreadRun) {
+  const tool = getToolReferenceFromToolCall(toolCall);
+
+  if (
+    run?.status === 'requires_action' &&
+    run.required_action?.type === 'submit_tool_approvals'
+  ) {
+    return run.required_action.submit_tool_approvals.tool_calls
+      .map((tool) => ({
+        id: tool.id,
+        toolId: getToolApprovalId(tool),
+        type: tool.type,
+      }))
+      .find((toolApproval) => toolApproval.toolId === tool.id);
+  }
+
+  return null;
+}
diff --git a/src/modules/chat/layout/InputBar.tsx b/src/modules/chat/layout/InputBar.tsx
@@ -40,7 +40,7 @@ import { FilesMenu } from './FilesMenu';
 import classes from './InputBar.module.scss';
 import { PromptSuggestions } from './PromptSuggestions';
 import { ThreadSettings } from './ThreadSettings';
-import { useMessages } from '../providers/useMessages';
+import { UserChatMessage } from '../types';
 
 interface Props {
   showSuggestions?: boolean;
@@ -100,6 +100,13 @@ export const InputBar = memo(function InputBar({
     inputRef.current?.focus();
   };
 
+  const handleAfterRemoveSentMessage = useCallback(
+    (message: UserChatMessage) => {
+      setValue('input', message.content, { shouldValidate: true });
+    },
+    [setValue],
+  );
+
   const isPending = status !== 'ready';
   const inputValue = watch('input');
   const isFileUploadEnabled = isFeatureEnabled(FeatureName.Files);
@@ -130,7 +137,10 @@ export const InputBar = memo(function InputBar({
         handleSubmit(({ input }) => {
           onMessageSubmit?.();
           resetForm();
-          sendMessage(input).then((result) => {
+
+          sendMessage(input, {
+            onAfterRemoveSentMessage: handleAfterRemoveSentMessage,
+          }).then((result) => {
             onMessageSent?.(result);
           });
         })();
@@ -212,10 +222,8 @@ export const InputBar = memo(function InputBar({
                   size="sm"
                   hasIconOnly
                   iconDescription="Cancel"
-                  disabled={status === 'waiting'}
-                  onClick={() => {
-                    cancel();
-                  }}
+                  disabled={status === 'waiting' || status === 'aborting'}
+                  onClick={() => cancel()}
                 />
               )}
             </div>