dbeaver/pro#6455 add cancel button (#4533)

* dbeaver/pro#6455 add cancel button

* dbeaver/pro#6455 add cancel button

* dbeaver/pro#6455 add cancel button

* dbeaver/pro#6455 fix cancel before job start

* dbeaver/pro#6455 chat message now in common model

* dbeaver/pro#6455 add synchronization for cancel

* dbeaver/pro#6455 add message cancel

---------

Co-authored-by: naumov <iamemptyhuh@gmail.com>
Co-authored-by: Daria Marutkina <125263541+dariamarutkina@users.noreply.github.com>

server/bundles/io.cloudbeaver.service.ai/schema/service.ai.graphqls

@@ -274,6 +274,9 @@ extend type Mutation @since(version: "23.2.2") {
     "Clears the chat messages in the AI chat conversation. The messages will be removed from the conversation."
     aiClearLastChatMessages(conversationId: ID!, messageId: ID!): Boolean! @since(version: "25.1.1")
 
+    "Cancels the in-progress AI response generation in the specified conversation."
+    aiCancelChatMessage(conversationId: ID!): Boolean! @since(version: "26.1.5")
+
     "Saves AI settings (e.g. supporting confirming metadata transfer) for the specified connection."
     aiSaveDataSourceSettings(dataSourceId: DataSourceIdInput!, settings: AIDataSourceSettingsInput!): AIDataSourceSettingsInfo! @since(version: "25.3.3")
 }

server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/WebAIUtils.java

@@ -24,6 +24,8 @@ import io.cloudbeaver.service.ai.model.WebAISendChatMessageInfo;
 import io.cloudbeaver.service.ai.model.WebAiChatResponseConsumer;
 import io.cloudbeaver.service.ai.model.events.WSAiChatMessageEvent;
 import io.cloudbeaver.utils.ServletAppUtils;
+import org.eclipse.core.runtime.IStatus;
+import org.eclipse.core.runtime.Status;
 import org.jkiss.code.NotNull;
 import org.jkiss.code.Nullable;
 import org.jkiss.dbeaver.DBException;
@@ -39,9 +41,9 @@ import org.jkiss.dbeaver.model.app.DBPProject;
 import org.jkiss.dbeaver.model.navigator.DBNDatabaseNode;
 import org.jkiss.dbeaver.model.navigator.DBNNode;
 import org.jkiss.dbeaver.model.navigator.DBNUtils;
+import org.jkiss.dbeaver.model.runtime.AbstractJob;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
 import org.jkiss.dbeaver.model.struct.DBSObject;
-import org.jkiss.dbeaver.utils.RuntimeUtils;
 import org.jkiss.utils.CommonUtils;
 
 import java.time.Clock;
@@ -139,10 +141,11 @@ public class WebAIUtils {
         @Nullable AIConfirmation confirmation,
         @NotNull String jobName
     ) {
-        webSession.setAttribute(getWaitingAttr(conversation), true);
         CompletableFuture<AIChatConversation> result = new CompletableFuture<>();
-        RuntimeUtils.scheduleJob(
-            jobName, monitor -> {
+        AbstractJob job = new AbstractJob(jobName) {
+            @NotNull
+            @Override
+            protected IStatus run(@NotNull DBRProgressMonitor monitor) {
                 try {
                     AIChatResponseConsumer subscriber = new WebAiChatResponseConsumer(conversation, webSession, aiChatSession);
                     aiChatSession.processAICompletion(
@@ -159,15 +162,25 @@ public class WebAIUtils {
                         }
                     });
                 } catch (DBException e) {
-                    log.error("Error processing AI completion", e);
-                    var errorMessage = conversation.addMessage(AIMessage.errorMessage(e));
-                    webSession.addSessionEvent(new WSAiChatMessageEvent(new WebAIMessage(errorMessage, conversation)));
-                    aiChatSession.notifyMessageAdd(conversation, errorMessage);
+                    if (monitor.isCanceled()) {
+                        log.debug("AI completion cancelled", e);
+                    } else {
+                        log.error("Error processing AI completion", e);
+                        var errorMessage = conversation.addMessage(AIMessage.errorMessage(e));
+                        webSession.addSessionEvent(new WSAiChatMessageEvent(new WebAIMessage(errorMessage, conversation)));
+                        aiChatSession.notifyMessageAdd(conversation, errorMessage);
+                    }
                 } finally {
-                    webSession.removeAttribute(getWaitingAttr(conversation));
+                    // Only clear the flag if it still points to this job
+                    if (webSession.getAttribute(getWaitingAttr(conversation)) == this) {
+                        webSession.removeAttribute(getWaitingAttr(conversation));
+                    }
                 }
+                return Status.OK_STATUS;
             }
-        );
+        };
+        webSession.setAttribute(getWaitingAttr(conversation), job);
+        job.schedule();
         return result;
     }
 
@@ -257,21 +270,25 @@ public class WebAIUtils {
             throw new DBWebException("AI services restricted for '%s'. Please contact your administrator if you need it.".formatted(
                 conversation.getDataSource()));
         }
-        String caption = conversation.getCaption();
-        AIChatMessage promptMessage = conversation.addMessage(message);
-        webSession.addSessionEvent(new WSAiChatMessageEvent(new WebAIMessage(promptMessage, conversation)));
-        aiChatSession.notifyMessageAdd(conversation, promptMessage);
-        if (!CommonUtils.equalObjects(caption, conversation.getCaption())) {
-            aiChatSession.notifyConversationRenamed(conversation, conversation.getCaption());
+        AIChatMessage promptMessage;
+        AIChatMessage result;
+        synchronized (conversation) {
+            String caption = conversation.getCaption();
+            promptMessage = conversation.addMessage(message);
+            webSession.addSessionEvent(new WSAiChatMessageEvent(new WebAIMessage(promptMessage, conversation)));
+            aiChatSession.notifyMessageAdd(conversation, promptMessage);
+            if (!CommonUtils.equalObjects(caption, conversation.getCaption())) {
+                aiChatSession.notifyConversationRenamed(conversation, conversation.getCaption());
+            }
+            if (!AIUtils.hasValidConfiguration()) {
+                throw new DBWebException("Invalid AI configuration");
+            }
+            if (webSession.getAttribute(WebAIUtils.getWaitingAttr(conversation)) != null) {
+                throw new DBWebException("Conversation is already waiting for response");
+            }
+            result = new AIChatMessage(conversation.getNextMessageId(), AIMessage.assistantMessage("", null));
+            WebAIUtils.scheduleConversationSubmission(webSession, aiChatSession, conversation, null, "AI completion");
         }
-        if (!AIUtils.hasValidConfiguration()) {
-            throw new DBWebException("Invalid AI configuration");
-        }
-        if (webSession.getAttribute(WebAIUtils.getWaitingAttr(conversation)) != null) {
-            throw new DBWebException("Conversation is already waiting for response");
-        }
-        AIChatMessage result = new AIChatMessage(conversation.getNextMessageId(), AIMessage.assistantMessage("", null));
-        WebAIUtils.scheduleConversationSubmission(webSession, aiChatSession, conversation, null, "AI completion");
         return new WebAISendChatMessageInfo(
             new WebAIChatConversation(webSession, conversation),
             new WebAIMessage(promptMessage, conversation),

server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/DBWServiceAI.java

@@ -139,6 +139,12 @@ public interface DBWServiceAI extends DBWService {
         @NotNull String messageId
     ) throws DBWebException;
 
+    @WebAction
+    boolean cancelChatMessage(
+        @NotNull WebSession webSession,
+        @NotNull String conversationId
+    ) throws DBWebException;
+
     @NotNull
     @WebAction
     WebAIDataSourceSettings getDataSourceAiSettings(

server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceAI.java

@@ -25,6 +25,7 @@ import io.cloudbeaver.model.session.WebSession;
 import io.cloudbeaver.server.CBApplication;
 import io.cloudbeaver.service.ai.WebAIUtils;
 import io.cloudbeaver.service.ai.model.*;
+import io.cloudbeaver.service.ai.model.events.WSAiChatMessageEvent;
 import io.cloudbeaver.service.ai.model.inputs.DataSourceId;
 import io.cloudbeaver.service.ai.model.inputs.WebAIChatConversationInput;
 import io.cloudbeaver.service.ai.model.inputs.WebAIConfigurationProfileInput;
@@ -46,12 +47,14 @@ import org.jkiss.dbeaver.model.ai.engine.AIDatabaseContext;
 import org.jkiss.dbeaver.model.ai.engine.AIEngine;
 import org.jkiss.dbeaver.model.ai.engine.AIEngineProperties;
 import org.jkiss.dbeaver.model.ai.engine.AIModel;
+import org.jkiss.dbeaver.model.ai.internal.AIChatMessages;
 import org.jkiss.dbeaver.model.ai.prompt.AIPromptGenerateSql;
 import org.jkiss.dbeaver.model.ai.registry.*;
 import org.jkiss.dbeaver.model.app.DBPProject;
 import org.jkiss.dbeaver.model.data.json.JSONUtils;
 import org.jkiss.dbeaver.model.logical.DBSLogicalDataSource;
 import org.jkiss.dbeaver.model.preferences.DBPPropertyDescriptor;
+import org.jkiss.dbeaver.model.runtime.AbstractJob;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
 import org.jkiss.dbeaver.model.websocket.event.WSWorkspaceConfigurationChangedEvent;
 import org.jkiss.dbeaver.runtime.DBWorkbench;
@@ -397,6 +400,33 @@ public class WebServiceAI implements DBWServiceAI {
         return true;
     }
 
+    @Override
+    public boolean cancelChatMessage(
+        @NotNull WebSession webSession,
+        @NotNull String conversationId
+    ) throws DBWebException {
+        WebAIUtils.validateAiPluginEnabled();
+        AIChatConversation conversation = WebAIUtils.getAiChatConversation(webSession, conversationId);
+        synchronized (conversation) {
+            boolean completionStarted = conversation.isActive();
+            conversation.cancelConversation();
+            boolean hadPendingJob = false;
+            if (webSession.getAttribute(WebAIUtils.getWaitingAttr(conversation)) instanceof AbstractJob job) {
+                hadPendingJob = true;
+                job.cancel();
+            }
+            webSession.removeAttribute(WebAIUtils.getWaitingAttr(conversation));
+            if (hadPendingJob && !completionStarted) {
+                // The completion job was still queued, so its response consumer will not be called.
+                // We need to add a cancellation message to the conversation and notify the client.
+                AIChatMessage cancelMessage = conversation.addMessage(
+                    AIMessage.warningMessage(AIChatMessages.ai_chat_conversation_cancelled));
+                webSession.addSessionEvent(new WSAiChatMessageEvent(new WebAIMessage(cancelMessage, conversation)));
+            }
+        }
+        return true;
+    }
+
     @NotNull
     @Override
     public WebAIDataSourceSettings getDataSourceAiSettings(

server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceBindingAI.java

@@ -163,6 +163,12 @@ public class WebServiceBindingAI extends WebServiceBindingBase<DBWServiceAI> imp
                     getArgumentVal(env, "conversationId"),
                     getArgumentVal(env, "messageId")
                 )
+            ).dataFetcher(
+                "aiCancelChatMessage",
+                env -> getService(env).cancelChatMessage(
+                    getWebSession(env),
+                    getArgumentVal(env, "conversationId")
+                )
             ).dataFetcher(
                 "aiCreateProfile", env -> getService(env).createProfile(
                     getWebSession(env),

server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/model/WebAiChatResponseConsumer.java

@@ -1,18 +1,18 @@
 /*
  * DBeaver - Universal Database Manager
- * Copyright (C) 2010-2026 DBeaver Corp
+ * Copyright (C) 2010-2026 DBeaver Corp and others
  *
- * All Rights Reserved.
+ * 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
  *
- * NOTICE:  All information contained herein is, and remains
- * the property of DBeaver Corp and its suppliers, if any.
- * The intellectual and technical concepts contained
- * herein are proprietary to DBeaver Corp and its suppliers
- * and may be covered by U.S. and Foreign Patents,
- * patents in process, and are protected by trade secret or copyright law.
- * Dissemination of this information or reproduction of this material
- * is strictly forbidden unless prior written permission is obtained
- * from DBeaver Corp.
+ *     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.
  */
 package io.cloudbeaver.service.ai.model;
 
@@ -22,9 +22,11 @@ import io.cloudbeaver.service.ai.model.events.WSAiChatMessageErrorEvent;
 import io.cloudbeaver.service.ai.model.events.WSAiChatMessageEvent;
 import org.jkiss.code.NotNull;
 import org.jkiss.dbeaver.model.ai.*;
+import org.jkiss.dbeaver.model.ai.internal.AIChatMessages;
 import org.jkiss.utils.CommonUtils;
 
 import java.util.List;
+import java.util.concurrent.CancellationException;
 
 public class WebAiChatResponseConsumer implements AIChatResponseConsumer {
     private final StringBuilder responseBuilder;
@@ -78,7 +80,12 @@ public class WebAiChatResponseConsumer implements AIChatResponseConsumer {
 
     @Override
     public void error(@NotNull Throwable throwable) {
-        var errorMessage = conversation.addMessage(AIMessage.errorMessage(throwable));
+
+        AIMessage aiMessage = throwable instanceof CancellationException cancellationException
+            ? AIMessage.warningMessage(cancellationException.getMessage())
+            : AIMessage.errorMessage(throwable);
+
+        AIChatMessage errorMessage = conversation.addMessage(aiMessage);
         if (responseBuilder.isEmpty()) {
             webSession.addSessionEvent(
                 new WSAiChatMessageEvent(new WebAIMessage(errorMessage, conversation)));
@@ -94,12 +101,19 @@ public class WebAiChatResponseConsumer implements AIChatResponseConsumer {
     }
 
     @Override
-    public void complete(@NotNull List<AIMessageMeta> meta, boolean finishConversation) {
+    public void complete(@NotNull List<AIMessageMeta> meta, boolean finishConversation, boolean isCanceled) {
         if (responseBuilder.isEmpty()) {
+            if (isCanceled) {
+                warning(AIChatMessages.ai_chat_conversation_cancelled);
+            }
             return;
         }
         AIChatMessage responseMessage = conversation.addMessage(AIMessage.assistantMessage(responseBuilder.toString(), meta));
         chatSession.notifyMessageAdd(conversation, responseMessage);
         webSession.addSessionEvent(new WSAiChatMessageChunkEvent(conversation.getId(), responseMessage.id(), null, true));
+
+        if (isCanceled) {
+            warning(AIChatMessages.ai_chat_conversation_cancelled);
+        }
     }
 }

webapp/packages/core-sdk/src/queries/ai/cancelConversation.gql

@@ -0,0 +1,3 @@
+mutation cancelConversation($conversationId: ID!) {
+  result: aiCancelChatMessage(conversationId: $conversationId)
+}

webapp/packages/plugin-ai-chat/src/AIChat/AIChatConversation/AIChatConversationsResource.ts

@@ -92,6 +92,11 @@ export class AIChatConversationsResource extends CachedMapResource<string, AICha
     return this.get(conversation.id)!;
   }
 
+  async cancelConversation(conversationId: string): Promise<boolean> {
+    const { result } = await this.graphQLService.sdk.cancelConversation({ conversationId });
+    return result;
+  }
+
   protected async loader(originalKey: ResourceKey<string>): Promise<Map<string, AIChatConversationInfo>> {
     const conversationList: AIChatConversationInfo[] = [];
 

webapp/packages/plugin-ai-chat/src/AIChat/AIChatMessage/AIChatMessageForm.tsx

@@ -13,10 +13,12 @@ import { ActionIconButton, AutoResizeTextarea, Form, s, useS, useTranslate } fro
 import { useService } from '@cloudbeaver/core-di';
 import { getOS, OperatingSystem } from '@cloudbeaver/core-utils';
 import { NotificationService } from '@cloudbeaver/core-events';
+import { Command } from '@dbeaver/ui-kit';
 
 import { AIChatMessageService } from './AIChatMessageService.js';
 import { AIChatConversationsService } from '../AIChatConversation/AIChatConversationsService.js';
 import { AIChatContext } from '../AIChatContext.js';
+import { AIChatConversationsResource } from '../AIChatConversation/AIChatConversationsResource.js';
 import classes from './AIChatMessageForm.module.css';
 
 interface Props {
@@ -30,6 +32,7 @@ export const AIChatMessageForm = observer<PropsWithChildren<Props>>(function AIC
   const notificationService = useService(NotificationService);
   const aiChatMessageService = useService(AIChatMessageService);
   const aiChatConversationsService = useService(AIChatConversationsService);
+  const aiChatConversationsResource = useService(AIChatConversationsResource);
 
   const [value, setValue] = useState('');
 
@@ -54,6 +57,16 @@ export const AIChatMessageForm = observer<PropsWithChildren<Props>>(function AIC
     }
   }
 
+  async function cancel() {
+    if (currentConversationId) {
+      try {
+        await aiChatConversationsResource.cancelConversation(currentConversationId);
+      } catch (exception: any) {
+        notificationService.logException(exception, 'plugin_ai_chat_conversation_cancel_failed');
+      }
+    }
+  }
+
   function getPlaceholder() {
     const OS = getOS();
     const symbol = OS === OperatingSystem.macOS ? '⌘' : 'Ctrl';
@@ -65,7 +78,12 @@ export const AIChatMessageForm = observer<PropsWithChildren<Props>>(function AIC
 
   return (
     <div className={s(styles, { container: true })}>
-      <Form className="tw:flex tw:items-end tw:gap-2" disableEnterSubmit={disabled} contents onSubmit={sendMessage}>
+      <Form
+        className="tw:flex tw:items-end tw:gap-2"
+        disableEnterSubmit={disabled}
+        contents
+        onSubmit={aiChatConversationsService.processing ? cancel : sendMessage}
+      >
         <div className={s(styles, { textareaContainer: true })}>
           <AutoResizeTextarea
             className={s(styles, { textarea: true })}
@@ -74,7 +92,16 @@ export const AIChatMessageForm = observer<PropsWithChildren<Props>>(function AIC
             autoFocus
             onChange={v => setValue(v)}
           />
-          <ActionIconButton name="/icons/send.svg" disabled={disabled} img onClick={sendMessage} />
+          {!aiChatConversationsService.processing ? (
+            <ActionIconButton name="/icons/send.svg" disabled={disabled} img onClick={sendMessage} />
+          ) : (
+            <Command
+              className="tw:cursor-pointer tw:w-10 tw:h-10 tw:bg-[var(--theme-primary)] tw:rounded-md tw:flex tw:items-center tw:justify-center tw:focus:opacity-80 tw:hover:opacity-80 tw:transition-opacity"
+              onClick={cancel}
+            >
+              <div className="tw:w-3 tw:h-3 tw:bg-[var(--theme-surface)] tw:rounded-xs" />
+            </Command>
+          )}
         </div>
       </Form>
       {children}

webapp/packages/plugin-ai-chat/src/locales/en.ts

@@ -44,6 +44,8 @@ export default [
   ['plugin_ai_chat_conversation_history_date_group_last_7_days', 'Last 7 days'],
   ['plugin_ai_chat_conversation_history_date_group_over_week_ago', 'Over a week ago'],
 
+  ['plugin_ai_chat_conversation_cancel_failed', 'Failed to cancel the conversation'],
+
   ['plugin_ai_chat_scope_change', 'Configure AI context'],
   ['plugin_ai_chat_scope_change_fail', 'Failed to change context'],
   ['plugin_ai_chat_profile_group', 'Active configuration'],

webapp/packages/plugin-ai-chat/src/locales/it.ts

@@ -44,6 +44,8 @@ export default [
   ['plugin_ai_chat_conversation_history_date_group_last_7_days', 'Last 7 days'],
   ['plugin_ai_chat_conversation_history_date_group_over_week_ago', 'Over a week ago'],
 
+  ['plugin_ai_chat_conversation_cancel_failed', 'Impossibile annullare la conversazione'],
+
   ['plugin_ai_chat_scope_change', 'Configure AI context'],
   ['plugin_ai_chat_scope_change_fail', 'Failed to change context'],
   ['plugin_ai_chat_profile_group', 'Active configuration'],

webapp/packages/plugin-ai-chat/src/locales/ru.ts

@@ -44,6 +44,8 @@ export default [
   ['plugin_ai_chat_conversation_history_date_group_last_7_days', 'Последние 7 дней'],
   ['plugin_ai_chat_conversation_history_date_group_over_week_ago', 'Более недели назад'],
 
+  ['plugin_ai_chat_conversation_cancel_failed', 'Не удалось отменить разговор'],
+
   ['plugin_ai_chat_scope_change', 'Настроить AI контекст'],
   ['plugin_ai_chat_scope_change_fail', 'Не удалось изменить контекст'],
   ['plugin_ai_chat_profile_group', 'Активная конфигурация'],

webapp/packages/plugin-ai-chat/src/locales/zh.ts

@@ -44,6 +44,8 @@ export default [
   ['plugin_ai_chat_conversation_history_date_group_last_7_days', 'Last 7 days'],
   ['plugin_ai_chat_conversation_history_date_group_over_week_ago', 'Over a week ago'],
 
+  ['plugin_ai_chat_conversation_cancel_failed', '无法取消对话'],
+
   ['plugin_ai_chat_scope_change', 'Configure AI context'],
   ['plugin_ai_chat_scope_change_fail', 'Failed to change context'],
   ['plugin_ai_chat_profile_group', 'Active configuration'],