diff --git a/server/bundles/io.cloudbeaver.service.ai/schema/service.ai.graphqls b/server/bundles/io.cloudbeaver.service.ai/schema/service.ai.graphqls index 83d0cff23f2..049d711f2e9 100644 --- a/server/bundles/io.cloudbeaver.service.ai/schema/service.ai.graphqls +++ b/server/bundles/io.cloudbeaver.service.ai/schema/service.ai.graphqls @@ -272,6 +272,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") } diff --git a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/WebAIUtils.java b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/WebAIUtils.java index 4a22f6c4902..cff82e72173 100644 --- a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/WebAIUtils.java +++ b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/WebAIUtils.java @@ -24,6 +24,8 @@ 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.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 static CompletableFuture scheduleConversationSubmissi @Nullable AIConfirmation confirmation, @NotNull String jobName ) { - webSession.setAttribute(getWaitingAttr(conversation), true); CompletableFuture 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 static CompletableFuture scheduleConversationSubmissi } }); } 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 static WebAISendChatMessageInfo submitPrompt( 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()); - } - 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 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"); } - 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), diff --git a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/DBWServiceAI.java b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/DBWServiceAI.java index 975fc9b6e85..1c30e530fd4 100644 --- a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/DBWServiceAI.java +++ b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/DBWServiceAI.java @@ -139,6 +139,12 @@ boolean setLastChatMessage( @NotNull String messageId ) throws DBWebException; + @WebAction + boolean cancelChatMessage( + @NotNull WebSession webSession, + @NotNull String conversationId + ) throws DBWebException; + @NotNull @WebAction WebAIDataSourceSettings getDataSourceAiSettings( diff --git a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceAI.java b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceAI.java index c76bfbb00a7..9f51f264611 100644 --- a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceAI.java +++ b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceAI.java @@ -25,6 +25,7 @@ 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.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 boolean setLastChatMessage( 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( diff --git a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceBindingAI.java b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceBindingAI.java index 0615a4731ad..69487e3bfa1 100644 --- a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceBindingAI.java +++ b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/gql/WebServiceBindingAI.java @@ -163,6 +163,12 @@ public void bindWiring(DBWBindingContext model) { 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), diff --git a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/model/WebAiChatResponseConsumer.java b/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/model/WebAiChatResponseConsumer.java index 77296210d41..bb61a90fd84 100644 --- a/server/bundles/io.cloudbeaver.service.ai/src/io/cloudbeaver/service/ai/model/WebAiChatResponseConsumer.java +++ b/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.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 void warning(@NotNull String message) { @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 void error(@NotNull Throwable throwable) { } @Override - public void complete(@NotNull List meta, boolean finishConversation) { + public void complete(@NotNull List 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); + } } } diff --git a/webapp/packages/core-sdk/src/queries/ai/cancelConversation.gql b/webapp/packages/core-sdk/src/queries/ai/cancelConversation.gql new file mode 100644 index 00000000000..91a36576d7f --- /dev/null +++ b/webapp/packages/core-sdk/src/queries/ai/cancelConversation.gql @@ -0,0 +1,3 @@ +mutation cancelConversation($conversationId: ID!) { + result: aiCancelChatMessage(conversationId: $conversationId) +} diff --git a/webapp/packages/plugin-ai-chat/src/AIChat/AIChatConversation/AIChatConversationsResource.ts b/webapp/packages/plugin-ai-chat/src/AIChat/AIChatConversation/AIChatConversationsResource.ts index d5926e7e09b..2d89ce9f452 100644 --- a/webapp/packages/plugin-ai-chat/src/AIChat/AIChatConversation/AIChatConversationsResource.ts +++ b/webapp/packages/plugin-ai-chat/src/AIChat/AIChatConversation/AIChatConversationsResource.ts @@ -92,6 +92,11 @@ export class AIChatConversationsResource extends CachedMapResource { + const { result } = await this.graphQLService.sdk.cancelConversation({ conversationId }); + return result; + } + protected async loader(originalKey: ResourceKey): Promise> { const conversationList: AIChatConversationInfo[] = []; diff --git a/webapp/packages/plugin-ai-chat/src/AIChat/AIChatMessage/AIChatMessageForm.tsx b/webapp/packages/plugin-ai-chat/src/AIChat/AIChatMessage/AIChatMessageForm.tsx index b8e89c71efe..42a2ce57320 100644 --- a/webapp/packages/plugin-ai-chat/src/AIChat/AIChatMessage/AIChatMessageForm.tsx +++ b/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>(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>(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>(function AIC return (
-
+
>(function AIC autoFocus onChange={v => setValue(v)} /> - + {!aiChatConversationsService.processing ? ( + + ) : ( + +
+ + )}
{children} diff --git a/webapp/packages/plugin-ai-chat/src/locales/en.ts b/webapp/packages/plugin-ai-chat/src/locales/en.ts index 5b148391a25..4dc15459e33 100644 --- a/webapp/packages/plugin-ai-chat/src/locales/en.ts +++ b/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'], diff --git a/webapp/packages/plugin-ai-chat/src/locales/it.ts b/webapp/packages/plugin-ai-chat/src/locales/it.ts index 609e6e89072..546bca837b2 100644 --- a/webapp/packages/plugin-ai-chat/src/locales/it.ts +++ b/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'], diff --git a/webapp/packages/plugin-ai-chat/src/locales/ru.ts b/webapp/packages/plugin-ai-chat/src/locales/ru.ts index 4cbbbc3f435..e3cbb528842 100644 --- a/webapp/packages/plugin-ai-chat/src/locales/ru.ts +++ b/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', 'Активная конфигурация'], diff --git a/webapp/packages/plugin-ai-chat/src/locales/zh.ts b/webapp/packages/plugin-ai-chat/src/locales/zh.ts index a7d0b2c7cc8..6fa2ea9c49d 100644 --- a/webapp/packages/plugin-ai-chat/src/locales/zh.ts +++ b/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'],