Add stop-generation button

main.ts

@@ -34,6 +34,7 @@ export default class LLMAgentPlugin extends Plugin {
 	settings: LLMAgentSettings;
 	eventBus: EventBus;
 	webSocketService: WebSocketService;
+	activeStreamControllers: Map<string, AbortController>;
 
 	async onload() {
 		await this.loadSettings();
@@ -74,6 +75,11 @@ export default class LLMAgentPlugin extends Plugin {
 
 		console.log('LLM Agent plugin загружен');
 
+		this.activeStreamControllers = new Map<string, AbortController>();
+
+		this.eventBus.on('stream-started', this.handleStreamStarted.bind(this));
+		this.eventBus.on('stream-ended', this.handleStreamEnded.bind(this));
+
 		// This adds a settings tab so the user can configure various aspects of the plugin
 		this.addSettingTab(new SampleSettingTab(this.app, this));
 	}
@@ -83,6 +89,11 @@ export default class LLMAgentPlugin extends Plugin {
 			this.webSocketService.disconnect();
 		}
 
+		this.activeStreamControllers.clear();
+
+		this.eventBus.off('stream-started', this.handleStreamStarted);
+		this.eventBus.off('stream-ended', this.handleStreamEnded);
+
 		console.log('LLM Agent plugin выгружен');
 	}
 
@@ -130,6 +141,14 @@ export default class LLMAgentPlugin extends Plugin {
 	handleNewChat() {
 		this.eventBus.emit('new-chat');
 	}
+
+	handleStreamStarted(data: { nodeId: string; controller: AbortController }): void {
+		this.activeStreamControllers.set(data.nodeId, data.controller);
+	}
+
+	handleStreamEnded(data: { nodeId: string }): void {
+		this.activeStreamControllers.delete(data.nodeId);
+	}
 }
 
 // ----------------------------------------------- Sample Setting Tab ---------------------------------------------------------------

src/components/ChatPanel.ts

@@ -15,16 +15,18 @@ import LLMAgentPlugin from 'main';
 import { MarkdownRenderer, setIcon, Notice } from 'obsidian'; // Добавляем импорт MarkdownRenderer
 import { Dialog } from 'src/utils/Dialog';
 
-export class ChatPanel {
+interface ChatPanelProps {
-	container: HTMLElement;
-	props: {
 	chatHistory: ChatHistoryItem[];
 	onSendMessage: (message: string, selectedModel?: string) => void;
 	availableModels?: string[];
 	selectedModel?: string;
 	plugin: LLMAgentPlugin;
 	currentGraphCustomSystemPrompt: string | null;
-	};
+}
+
+export class ChatPanel {
+	container: HTMLElement;
+	props: ChatPanelProps;
 	chatHistory: ChatHistoryItem[];
 	messageInput: HTMLDivElement;
 	sendButton: HTMLButtonElement;
@@ -33,6 +35,8 @@ export class ChatPanel {
 	modelDropdown: HTMLDivElement;
 	selectedModel: string;
 	customPromptIndicator: HTMLSpanElement | null = null;
+	currentStreamingNodeId: string | null;
+	stopButton: HTMLButtonElement;
 
 	// ----------------------------------------------- Reference System Fields ---------------------------------------------------------------
 	private autocompleteManager: AutocompleteManager;
@@ -42,17 +46,7 @@ export class ChatPanel {
 	private currentReferences: FileReference[] = [];
 	private isProcessingReference = false;
 
-	constructor(
+	constructor(container: HTMLElement, props: ChatPanelProps) {
-		container: HTMLElement,
-		props: {
-			chatHistory: ChatHistoryItem[];
-			onSendMessage: (message: string, selectedModel?: string) => void;
-			availableModels?: string[];
-			selectedModel?: string;
-			plugin: LLMAgentPlugin;
-			currentGraphCustomSystemPrompt: string | null;
-		}
-	) {
 		this.container = container;
 		this.props = props;
 		this.chatHistory = props.chatHistory || [];
@@ -331,7 +325,10 @@ export class ChatPanel {
                             </div>
                         </div>
                         <div class="control-buttons-right">
-                            <button class="submit-button" id="send-button" data-tooltip="Enter">
+							<button class="stop-button" id="stop-button" style="display: none;">
+                                <span class="stop-icon">■</span>
+                            </button>
+                            <button class="submit-button" id="send-button">
                                 <span class="submit-text">⏎ Enter</span>
                             </button>
                         </div>
@@ -347,6 +344,7 @@ export class ChatPanel {
 		this.modelSelectButton = this.container.querySelector('#model-select-button') as HTMLButtonElement;
 		this.modelDropdown = this.container.querySelector('#model-dropdown') as HTMLDivElement;
 		this.customPromptIndicator = this.container.querySelector('#custom-prompt-indicator') as HTMLSpanElement;
+		this.stopButton = this.container.querySelector('#stop-button') as HTMLButtonElement;
 
 		this.setupEventListeners();
 		this.setupReferenceSystem();
@@ -418,6 +416,8 @@ export class ChatPanel {
 			this.toggleModelDropdown();
 		});
 
+		this.stopButton.addEventListener('click', () => this.handleStopGeneration());
+
 		// Close dropdown when clicking outside
 		document.addEventListener('click', (e) => {
 			if (!this.modelSelectButton.contains(e.target as Node) && !this.modelDropdown.contains(e.target as Node)) {
@@ -896,6 +896,24 @@ export class ChatPanel {
 		});
 	}
 
+	private handleStopGeneration(): void {
+		if (this.currentStreamingNodeId) {
+			this.props.plugin.eventBus.emit('stop-stream-request', {
+				nodeId: this.currentStreamingNodeId
+			});
+		}
+	}
+
+	updateStopButtonVisibility(isStreaming: boolean, nodeId: string | null): void {
+		this.currentStreamingNodeId = nodeId;
+
+		if (isStreaming && nodeId) {
+			this.stopButton.style.display = 'inline-flex';
+		} else {
+			this.stopButton.style.display = 'none';
+		}
+	}
+
 	private copyMessageContent(content: string): Promise<void> {
 		return navigator.clipboard.writeText(content);
 	}

src/css/styles.css

@@ -227,34 +227,6 @@
 	scroll-margin-bottom: 30px; /* Добавляет отступ в 20px снизу при прокрутке к этому элементу */
 }
 
-.submit-button {
-	padding: 4px 8px; /* Восстанавливаем паддинг */
-	border: none; /* Убираем рамку */
-	border-radius: 8px;
-	cursor: pointer;
-	font-size: 11px;
-	transition: all 0.2s;
-	/* Явно переопределяем стили, которые могли быть затронуты */
-	width: auto; /* Убедимся, что ширина не ограничена */
-	height: auto; /* Убедимся, что высота не ограничена */
-	background-color: var(
-		--button-primary-background
-	) !important; /* Используем переменную для основного цвета кнопки */
-	color: var(
-		--button-primary-text-color
-	) !important; /* Используем переменную для цвета текста кнопки */
-	opacity: 1; /* Если .clickable-icon устанавливает opacity */
-	display: inline-flex; /* Или flex, если нужно выравнивание элементов внутри */
-	align-items: center; /* Для вертикального выравнивания контента, если display: flex */
-	justify-content: center; /* Для горизонтального выравнивания контента, если display: flex */
-}
-
-#send-button:hover {
-	background-color: var(
-		--button-primary-hover-background
-	) !important; /* Переменная для наведения */
-}
-
 .chat-input-container {
 	border-top: 1px solid var(--background-modifier-border);
 	padding: 8px 10px 4px 10px;
@@ -426,20 +398,39 @@
 	border-radius: 0 0 8px 8px;
 }
 
-.submit-button {
+.stop-button, .submit-button {
-	padding: 4px 8px;
+	padding: 4px 8px; /* Восстанавливаем паддинг */
-	border: none;
+	border: none; /* Убираем рамку */
 	border-radius: 8px;
 	cursor: pointer;
-	font-size: 11px;
+	font-size: 14px;
 	transition: all 0.2s;
+	/* Явно переопределяем стили, которые могли быть затронуты */
+	width: auto; /* Убедимся, что ширина не ограничена */
+	height: auto; /* Убедимся, что высота не ограничена */
+	background-color: var(
+		--button-primary-background
+	) !important; /* Используем переменную для основного цвета кнопки */
+	color: var(
+		--button-primary-text-color
+	) !important; /* Используем переменную для цвета текста кнопки */
+	opacity: 1; /* Если .clickable-icon устанавливает opacity */
+	display: inline-flex; /* Или flex, если нужно выравнивание элементов внутри */
+	align-items: center; /* Для вертикального выравнивания контента, если display: flex */
+	justify-content: center; /* Для горизонтального выравнивания контента, если display: flex */
 }
 
-.submit-button:hover:enabled {
+.stop-button:hover, .submit-button:hover {
+	background-color: var(
+		--button-primary-hover-background
+	) !important; /* Переменная для наведения */
+}
+
+.stop-button:hover:enabled, .submit-button:hover:enabled {
 	filter: brightness(1.25);
 }
 
-.submit-button:disabled {
+.stop-button:disabled, .submit-button:disabled {
 	opacity: 0.5;
 	cursor: not-allowed;
 }

src/views/ChatView.ts

@@ -20,6 +20,7 @@ export class ChatView extends ItemView {
 	chatPanel: ChatPanel | null;
 	availableModels: string[];
 	currentGraphCustomSystemPrompt: string | null;
+	currentStreamingNodeId: string | null = null;
 
 	constructor(leaf: WorkspaceLeaf, plugin: LLMAgentPlugin) {
 		super(leaf);
@@ -72,6 +73,9 @@ export class ChatView extends ItemView {
 		this.plugin.eventBus.on('ws-node-type-changed', this.handleNodeTypeChanged.bind(this));
 		this.plugin.eventBus.on('start-new-independent-chat', this.handleStartNewIndependentChat.bind(this));
 		this.plugin.eventBus.on('graph-custom-prompt-updated', this.handleGraphCustomPromptUpdated.bind(this));
+		this.plugin.eventBus.on('stop-stream-request', this.handleStopStreamRequest.bind(this)); // НОВОЕ
+		this.plugin.eventBus.on('stream-started', this.handleStreamStarted.bind(this));
+		this.plugin.eventBus.on('stream-ended', this.handleStreamEnded.bind(this));
 	}
 
 	async onClose(): Promise<void> {
@@ -86,6 +90,9 @@ export class ChatView extends ItemView {
 		this.plugin.eventBus.off('ws-node-type-changed', this.handleNodeTypeChanged);
 		this.plugin.eventBus.off('start-new-independent-chat', this.handleStartNewIndependentChat);
 		this.plugin.eventBus.off('graph-custom-prompt-updated', this.handleGraphCustomPromptUpdated);
+		this.plugin.eventBus.off('stop-stream-request', this.handleStopStreamRequest); // НОВОЕ
+		this.plugin.eventBus.off('stream-started', this.handleStreamStarted);
+		this.plugin.eventBus.off('stream-ended', this.handleStreamEnded);
 	}
 
 	// ----------------------------------------------- API Methods ---------------------------------------------------------------
@@ -226,7 +233,20 @@ export class ChatView extends ItemView {
     handleNodeSelected(data: { nodeId: string }): void {
         this.currentNode = data.nodeId;
         if (this.selectedGraphId) {
-			this.fetchMessagesFromRootToNode(data.nodeId);
+            this.fetchMessagesFromRootToNode(data.nodeId).then(() => {
+                // После загрузки новой ветки проверяем, есть ли стриминг в ней
+                if (this.currentStreamingNodeId) {
+                    const isInNewBranch = this.chatHistory.some(
+                        msg => msg.node_id === this.currentStreamingNodeId
+                    );
+                    if (this.chatPanel) {
+                        this.chatPanel.updateStopButtonVisibility(
+                            isInNewBranch,
+                            isInNewBranch ? this.currentStreamingNodeId : null
+                        );
+                    }
+                }
+            });
         }
     }
 
@@ -280,27 +300,36 @@ export class ChatView extends ItemView {
 			Object.assign(body, { existing_assistant_node_id: existingAssistantNodeId });
 		}
 
+		const abortController = new AbortController();
+		let assistantNodeId: string | null = null;
+
+		let reader: ReadableStreamDefaultReader<Uint8Array> | null = null;
+
+		try {
 			const response = await fetch(`${this.plugin.settings.apiBaseUrl}/chat/stream`, {
 				method: 'POST',
 				headers: { 'Content-Type': 'application/json' },
-			body: JSON.stringify(body)
+				body: JSON.stringify(body),
+				signal: abortController.signal
 			});
 
 			if (!response.ok) {
 				throw new Error(`Ошибка стриминга: ${response.status}`);
 			}
 
-		const reader = response.body?.getReader();
+			if (!response.body) {
-		if (!reader) {
 				throw new Error('Response body is null');
 			}
+			reader = response.body.getReader();
 
 			const decoder = new TextDecoder();
-		let assistantNodeId: string | null = null;
 			let accumulatedContent = '';
 
-		try {
 			while (true) {
+				if (abortController.signal.aborted) {
+					break;
+				}
+
 				const { done, value } = await reader.read();
 				if (done) break;
 
@@ -313,6 +342,12 @@ export class ChatView extends ItemView {
 
 						if (data.type === 'node_created') {
 							assistantNodeId = data.node_id;
+
+							this.plugin.eventBus.emit('stream-started', {
+								nodeId: assistantNodeId,
+								controller: abortController
+							});
+
 							if (data.reused) {
 								this.clearStreamingMessage(assistantNodeId);
 							}
@@ -331,6 +366,8 @@ export class ChatView extends ItemView {
 									graphId: graphId,
 									currentNode: assistantNodeId
 								});
+
+								this.plugin.eventBus.emit('stream-ended', { nodeId: assistantNodeId });
 							}
 							break;
 						} else if (data.type === 'done') {
@@ -339,13 +376,29 @@ export class ChatView extends ItemView {
 								graphId: graphId,
 								currentNode: assistantNodeId
 							});
+
+							if (assistantNodeId) {
+								this.plugin.eventBus.emit('stream-ended', { nodeId: assistantNodeId });
+							}
 							break;
 						}
 					}
 				}
 			}
+		} catch (error: any) {
+			if (error.name === 'AbortError') {
+				console.log('Стрим был отменен пользователем');
+				if (assistantNodeId) {
+					this.plugin.eventBus.emit('stream-ended', { nodeId: assistantNodeId });
+				}
+			} else {
+				throw error; // Пробрасываем другие ошибки
+			}
 		} finally {
-			reader.releaseLock();
+			reader?.releaseLock();
+			if (assistantNodeId) {
+				this.plugin.eventBus.emit('stream-ended', { nodeId: assistantNodeId });
+			}
 		}
 	}
 
@@ -466,6 +519,34 @@ export class ChatView extends ItemView {
 		new Notice('Начат новый чат. Текущий граф не изменен.', 2000);
 	}
 
+	handleStopStreamRequest(data: { nodeId: string }): void {
+		const controller = (this.plugin as any).activeStreamControllers?.get(data.nodeId);
+		if (controller) {
+			controller.abort();
+			new Notice('Генерация остановлена пользователем.', 2000);
+		}
+	}
+
+	handleStreamStarted(data: { nodeId: string; controller: AbortController }): void {
+		this.currentStreamingNodeId = data.nodeId;
+		// TODO: Временно не проверяем, что блок сообщения уже отрисован в панели, 
+		// т.к. у нас пока к моменту вызова 'stream-started' сообщение ещё не отрисовано
+		// const isInCurrentBranch = this.chatHistory.some((msg) => msg.node_id === data.nodeId);
+		// if (this.chatPanel && isInCurrentBranch) {
+		if (this.chatPanel) {
+			this.chatPanel.updateStopButtonVisibility(true, data.nodeId);
+		}
+	}
+
+	handleStreamEnded(data: { nodeId: string }): void {
+		if (this.currentStreamingNodeId === data.nodeId) {
+			this.currentStreamingNodeId = null;
+			if (this.chatPanel) {
+				this.chatPanel.updateStopButtonVisibility(false, null);
+			}
+		}
+	}
+
 	// ----------------------------------------------- WebSocket Event Handlers ---------------------------------------------------------------
 
 	handleNodeTypeChanged(data: { graphId: string; nodeId: string; nodeType: string }): void {

src/views/GraphView.ts

@@ -43,6 +43,8 @@ export class GraphView extends ItemView {
 	newIndependentChatActionButton: HTMLButtonElement | null = null;
 	graphSettingsButton: HTMLButtonElement | null = null;
 
+	activeStreamControllers: Map<string, AbortController>;
+
 	constructor(leaf: WorkspaceLeaf, plugin: LLMAgentPlugin) {
 		super(leaf);
 		this.plugin = plugin;
@@ -56,6 +58,7 @@ export class GraphView extends ItemView {
 		this.selectedNodeIds = new Set<string>();
 		this.selectedGraphTitle = null;
 		this.currentGraphCustomSystemPrompt = null;
+		this.activeStreamControllers = new Map<string, AbortController>();
 	}
 
 	/**
@@ -135,6 +138,9 @@ export class GraphView extends ItemView {
 		this.plugin.eventBus.on('ws-node-title-updated', this.handleNodeTitleUpdated.bind(this));
 		this.plugin.eventBus.on('ws-node-type-changed', this.handleNodeTypeChanged.bind(this));
 
+		this.plugin.eventBus.on('stream-started', this.handleStreamStarted.bind(this));
+		this.plugin.eventBus.on('stream-ended', this.handleStreamEnded.bind(this));
+
 		// Загружаем список графов
 		await this.fetchGraphsList();
 	}
@@ -163,6 +169,9 @@ export class GraphView extends ItemView {
 		this.plugin.eventBus.off('ws-graph-title-updated', this.handleGraphTitleUpdated);
 		this.plugin.eventBus.off('ws-node-title-updated', this.handleNodeTitleUpdated);
 		this.plugin.eventBus.off('ws-node-type-changed', this.handleNodeTypeChanged);
+
+		this.plugin.eventBus.off('stream-started', this.handleStreamStarted);
+		this.plugin.eventBus.off('stream-ended', this.handleStreamEnded);
 	}
 
 	// ----------------------------------------------- Custom Title Update ---------------------------------------------------
@@ -666,6 +675,14 @@ export class GraphView extends ItemView {
 		new Notice('Операция удаления завершена.', 2000);
 	}
 
+	handleStreamStarted(data: { nodeId: string; controller: AbortController }): void {
+		this.activeStreamControllers.set(data.nodeId, data.controller);
+	}
+
+	handleStreamEnded(data: { nodeId: string }): void {
+		this.activeStreamControllers.delete(data.nodeId);
+	}
+
 	// ----------------------------------------------- WebSocket Event Handlers ---------------------------------------------------------------
 
 	/**