diff --git a/ui/desktop/src/components/GooseMessage.tsx b/ui/desktop/src/components/GooseMessage.tsx index fddcd3229c12..73b7d4bc36f0 100644 --- a/ui/desktop/src/components/GooseMessage.tsx +++ b/ui/desktop/src/components/GooseMessage.tsx @@ -13,7 +13,7 @@ import { import { Message, confirmPermission } from '../api'; import ToolCallConfirmation from './ToolCallConfirmation'; import MessageCopyLink from './MessageCopyLink'; -import { NotificationEvent } from '../hooks/useMessageStream'; +import { NotificationEvent } from '../hooks/useChatStream'; import { cn } from '../utils'; import { identifyConsecutiveToolCalls, shouldHideTimestamp } from '../utils/toolCallChaining'; diff --git a/ui/desktop/src/components/ProgressiveMessageList.tsx b/ui/desktop/src/components/ProgressiveMessageList.tsx index d95e46fb2747..66c6cbbc48e0 100644 --- a/ui/desktop/src/components/ProgressiveMessageList.tsx +++ b/ui/desktop/src/components/ProgressiveMessageList.tsx @@ -19,7 +19,7 @@ import { Message } from '../api'; import GooseMessage from './GooseMessage'; import UserMessage from './UserMessage'; import { SystemNotificationInline } from './context_management/SystemNotificationInline'; -import { NotificationEvent } from '../hooks/useMessageStream'; +import { NotificationEvent } from '../hooks/useChatStream'; import LoadingGoose from './LoadingGoose'; import { ChatType } from '../types/chat'; import { identifyConsecutiveToolCalls, isInChain } from '../utils/toolCallChaining'; diff --git a/ui/desktop/src/components/ToolCallWithResponse.tsx b/ui/desktop/src/components/ToolCallWithResponse.tsx index 8a3166e0d578..b4f4d2a69018 100644 --- a/ui/desktop/src/components/ToolCallWithResponse.tsx +++ b/ui/desktop/src/components/ToolCallWithResponse.tsx @@ -7,7 +7,7 @@ import MarkdownContent from './MarkdownContent'; import { ToolRequestMessageContent, ToolResponseMessageContent } from '../types/message'; import { cn, snakeToTitleCase } from '../utils'; import { LoadingStatus } from './ui/Dot'; -import { NotificationEvent } from '../hooks/useMessageStream'; +import { NotificationEvent } from '../hooks/useChatStream'; import { ChevronRight, FlaskConical } from 'lucide-react'; import { TooltipWrapper } from './settings/providers/subcomponents/buttons/TooltipWrapper'; import MCPUIResourceRenderer from './MCPUIResourceRenderer'; diff --git a/ui/desktop/src/components/sessions/SessionHistoryView.tsx b/ui/desktop/src/components/sessions/SessionHistoryView.tsx index dc8b69464299..3dbc4c312c07 100644 --- a/ui/desktop/src/components/sessions/SessionHistoryView.tsx +++ b/ui/desktop/src/components/sessions/SessionHistoryView.tsx @@ -33,7 +33,6 @@ import { Tooltip, TooltipContent, TooltipTrigger } from '../ui/Tooltip'; import { Message, Session } from '../../api'; import { useNavigation } from '../../hooks/useNavigation'; -// Helper function to determine if a message is a user message (same as useChatEngine) const isUserMessage = (message: Message): boolean => { if (message.role === 'assistant') { return false; diff --git a/ui/desktop/src/hooks/useChatEngine.test.ts b/ui/desktop/src/hooks/useChatEngine.test.ts deleted file mode 100644 index d62befeb9dc1..000000000000 --- a/ui/desktop/src/hooks/useChatEngine.test.ts +++ /dev/null @@ -1,167 +0,0 @@ -/** - * @vitest-environment jsdom - */ -import { act, renderHook } from '@testing-library/react'; -import type { Mock } from 'vitest'; -import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { useChatEngine } from './useChatEngine'; -import { getTextContent } from '../types/message'; -import { Message } from '../api'; -import { ChatType } from '../types/chat'; - -// Mock the useMessageStream hook which is a dependency of useChatEngine -vi.mock('./useMessageStream', () => ({ - useMessageStream: vi.fn(), -})); - -// Mock the sessions API which is another dependency -vi.mock('../sessions', () => ({ - fetchSessionDetails: vi.fn().mockResolvedValue({ metadata: {} }), -})); - -describe('useChatEngine', () => { - let mockUseMessageStream: Mock; - - beforeEach(async () => { - // Mock the appConfig and electron APIs on the existing window object - Object.defineProperty(window, 'appConfig', { - value: { - get: vi.fn((key: string) => { - if (key === 'GOOSE_API_HOST') return 'http://localhost'; - if (key === 'GOOSE_PORT') return '8000'; - return null; - }), - }, - writable: true, - }); - - Object.defineProperty(window, 'electron', { - value: { - logInfo: vi.fn(), - }, - writable: true, - }); - - // Dynamically import the hook so we can get a reference to the mock - const { useMessageStream } = await import('./useMessageStream'); - mockUseMessageStream = useMessageStream as Mock; - - // Reset all mocks before each test to ensure a clean state - vi.clearAllMocks(); - - // Provide a complete, default mock implementation for useMessageStream - mockUseMessageStream.mockReturnValue({ - messages: [], - append: vi.fn(), - stop: vi.fn(), - chatState: 'idle', - error: undefined, - setMessages: vi.fn(), - input: '', - setInput: vi.fn(), - handleInputChange: vi.fn(), - handleSubmit: vi.fn(), - updateMessageStreamBody: vi.fn(), - notifications: [], - sessionMetadata: undefined, - setError: vi.fn(), - }); - }); - - describe('onMessageUpdate', () => { - it('should truncate history and append the updated message when a message is edited', () => { - // --- 1. ARRANGE --- - const metadata = { - agentVisible: true, - userVisible: true, - }; - - const initialMessages: Message[] = [ - { - id: '1', - role: 'user', - content: [{ type: 'text', text: 'First message' }], - created: 0, - metadata, - }, - { - id: '2', - role: 'assistant', - content: [{ type: 'text', text: 'First response' }], - created: 1, - metadata, - }, - { - id: '3', - role: 'user', - content: [{ type: 'text', text: 'Message to be edited' }], - created: 2, - metadata, - }, - { - id: '4', - role: 'assistant', - content: [{ type: 'text', text: 'Response to be deleted' }], - created: 3, - metadata, - }, - ]; - - const mockSetMessages = vi.fn(); - const mockAppend = vi.fn(); - - // Configure the mock to return specific values for this test case - mockUseMessageStream.mockReturnValue({ - messages: initialMessages, - append: mockAppend, - setMessages: mockSetMessages, - notifications: [], - stop: vi.fn(), - chatState: 'idle', - error: undefined, - input: '', - setInput: vi.fn(), - handleInputChange: vi.fn(), - handleSubmit: vi.fn(), - updateMessageStreamBody: vi.fn(), - sessionMetadata: undefined, - setError: vi.fn(), - }); - - const mockChat: ChatType = { - sessionId: 'test-chat', - messages: initialMessages, - name: 'Test Chat', - messageHistoryIndex: 0, - }; - - // Render the hook with our test setup - const { result } = renderHook(() => - useChatEngine({ - chat: mockChat, - setChat: vi.fn(), - }) - ); - - const messageIdToUpdate = '3'; - const newContent = 'This is the edited message.'; - - // --- 2. ACT --- - // Call the function we want to test - act(() => { - result.current.onMessageUpdate(messageIdToUpdate, newContent); - }); - - // --- 3. ASSERT --- - // Verify that setMessages was called with the correctly truncated history - const expectedTruncatedHistory = initialMessages.slice(0, 2); - expect(mockSetMessages).toHaveBeenCalledWith(expectedTruncatedHistory); - - // Verify that append was called with the new message - expect(mockAppend).toHaveBeenCalledTimes(1); - const appendedMessage = mockAppend.mock.calls[0][0]; - expect(getTextContent(appendedMessage)).toBe(newContent); - expect(appendedMessage.role).toBe('user'); - }); - }); -}); diff --git a/ui/desktop/src/hooks/useChatEngine.ts b/ui/desktop/src/hooks/useChatEngine.ts deleted file mode 100644 index 19a333aefda9..000000000000 --- a/ui/desktop/src/hooks/useChatEngine.ts +++ /dev/null @@ -1,473 +0,0 @@ -import { useCallback, useEffect, useMemo, useState } from 'react'; -import { getApiUrl } from '../config'; -import { useMessageStream } from './useMessageStream'; -import { LocalMessageStorage } from '../utils/localMessageStorage'; -import { createUserMessage, getTextContent, ToolResponseMessageContent } from '../types/message'; -import { getSession, Message } from '../api'; -import { ChatType } from '../types/chat'; -import { ChatState } from '../types/chatState'; - -// Helper function to determine if a message is a user message -const isUserMessage = (message: Message): boolean => { - if (message.role === 'assistant') { - return false; - } - return !message.content.every((c) => c.type === 'toolConfirmationRequest'); -}; - -interface UseChatEngineProps { - chat: ChatType; - setChat: (chat: ChatType) => void; - onMessageStreamFinish?: () => void; - onMessageSent?: () => void; // Add callback for when message is sent -} - -export const useChatEngine = ({ - chat, - setChat, - onMessageStreamFinish, - onMessageSent, -}: UseChatEngineProps) => { - const [lastInteractionTime, setLastInteractionTime] = useState(Date.now()); - const [sessionTokenCount, setSessionTokenCount] = useState(0); - const [sessionInputTokens, setSessionInputTokens] = useState(0); - const [sessionOutputTokens, setSessionOutputTokens] = useState(0); - const [localInputTokens, setLocalInputTokens] = useState(0); - const [localOutputTokens, setLocalOutputTokens] = useState(0); - const [powerSaveTimeoutId, setPowerSaveTimeoutId] = useState(null); - - // Track pending edited message - const [pendingEdit, setPendingEdit] = useState<{ id: string; content: string } | null>(null); - - // Store message in global history when it's added - const storeMessageInHistory = useCallback((message: Message) => { - if (isUserMessage(message)) { - const text = getTextContent(message); - if (text) { - LocalMessageStorage.addMessage(text); - } - } - }, []); - - const stopPowerSaveBlocker = useCallback(() => { - try { - window.electron.stopPowerSaveBlocker(); - } catch (error) { - console.error('Failed to stop power save blocker:', error); - } - - // Clear timeout if it exists - if (powerSaveTimeoutId) { - window.clearTimeout(powerSaveTimeoutId); - setPowerSaveTimeoutId(null); - } - }, [powerSaveTimeoutId]); - - const { - messages, - append: originalAppend, - stop, - chatState, - error, - setMessages, - input: _input, - setInput: _setInput, - handleInputChange: _handleInputChange, - updateMessageStreamBody, - notifications, - session, - setError, - tokenState, - } = useMessageStream({ - api: getApiUrl('/reply'), - id: chat.sessionId, - initialMessages: chat.messages, - body: { - session_id: chat.sessionId, - session_working_dir: window.appConfig.get('GOOSE_WORKING_DIR'), - ...(chat.recipe?.title - ? { - recipe_name: chat.recipe.title, - recipe_version: chat.recipe?.version ?? 'unknown', - } - : {}), - }, - onFinish: async (_message, _reason) => { - stopPowerSaveBlocker(); - - const timeSinceLastInteraction = Date.now() - lastInteractionTime; - window.electron.logInfo('last interaction:' + lastInteractionTime); - if (timeSinceLastInteraction > 60000) { - // 60000ms = 1 minute - window.electron.showNotification({ - title: 'Goose finished the task.', - body: 'Click here to expand.', - }); - } - - // Always emit refresh event when message stream finishes for new sessions - // Check if this is a new session by looking at the current session ID format - const isNewSession = chat.sessionId && chat.sessionId.match(/^\d{8}_\d{6}$/); - if (isNewSession) { - console.log( - 'ChatEngine: Message stream finished for new session, emitting message-stream-finished event' - ); - // Emit event to trigger session refresh - window.dispatchEvent(new CustomEvent('message-stream-finished')); - } - - onMessageStreamFinish?.(); - }, - onError: (error) => { - stopPowerSaveBlocker(); - - console.log( - 'CHAT ENGINE RECEIVED ERROR FROM MESSAGE STREAM:', - JSON.stringify( - { - errorMessage: error.message, - errorName: error.name, - isTokenLimitError: (error as Error & { isTokenLimitError?: boolean }).isTokenLimitError, - errorStack: error.stack, - timestamp: new Date().toISOString(), - sessionId: chat.sessionId, - }, - null, - 2 - ) - ); - }, - }); - - // Wrap append to store messages in global history - const append = useCallback( - (messageOrString: Message | string) => { - const message = - typeof messageOrString === 'string' ? createUserMessage(messageOrString) : messageOrString; - storeMessageInHistory(message); - - // If this is the first message in a new session, trigger a refresh immediately - // Only trigger if we're starting a completely new session (no existing messages) - if (messages.length === 0 && chat.messages.length === 0) { - // Emit event to indicate a new session is being created - window.dispatchEvent(new CustomEvent('session-created')); - } - - return originalAppend(message); - }, - [originalAppend, storeMessageInHistory, messages.length, chat.messages.length] - ); - - // Simple token estimation function (roughly 4 characters per token) - const estimateTokens = (text: string): number => { - return Math.ceil(text.length / 4); - }; - - // Calculate token counts from messages - useEffect(() => { - let inputTokens = 0; - let outputTokens = 0; - - messages.forEach((message) => { - const textContent = getTextContent(message); - if (textContent) { - const tokens = estimateTokens(textContent); - if (message.role === 'user') { - inputTokens += tokens; - } else if (message.role === 'assistant') { - outputTokens += tokens; - } - } - }); - - setLocalInputTokens(inputTokens); - setLocalOutputTokens(outputTokens); - }, [messages]); - - // Update chat messages when they change - useEffect(() => { - // @ts-expect-error - TypeScript being overly strict about the return type - setChat((prevChat: ChatType) => ({ ...prevChat, messages })); - }, [messages, setChat]); - - useEffect(() => { - const fetchSessionTokens = async () => { - try { - const response = await getSession({ - path: { session_id: chat.sessionId }, - throwOnError: true, - }); - const sessionDetails = response.data; - setSessionTokenCount(sessionDetails.total_tokens || 0); - setSessionInputTokens(sessionDetails.accumulated_input_tokens || 0); - setSessionOutputTokens(sessionDetails.accumulated_output_tokens || 0); - } catch (err) { - console.error('Error fetching session token count:', err); - } - }; - // Only fetch session tokens when chat state is idle to avoid resetting during streaming - if (chat.sessionId && chatState === ChatState.Idle) { - fetchSessionTokens(); - } - }, [chat.sessionId, messages, chatState]); - - // Update token counts when session changes from the message stream - useEffect(() => { - if (session) { - setSessionTokenCount(session.total_tokens || 0); - setSessionInputTokens(session.accumulated_input_tokens || 0); - setSessionOutputTokens(session.accumulated_output_tokens || 0); - } - }, [session]); - - useEffect(() => { - return () => { - if (powerSaveTimeoutId) { - window.clearTimeout(powerSaveTimeoutId); - } - try { - window.electron.stopPowerSaveBlocker(); - } catch (error) { - console.error('Failed to stop power save blocker during cleanup:', error); - } - }; - }, [powerSaveTimeoutId]); - - // Handle submit - const handleSubmit = useCallback( - (combinedTextFromInput: string, onSummaryReset?: () => void) => { - if (combinedTextFromInput.trim()) { - try { - window.electron.startPowerSaveBlocker(); - } catch (error) { - console.error('Failed to start power save blocker:', error); - } - - setLastInteractionTime(Date.now()); - - // Set a timeout to automatically stop the power save blocker after 15 minutes - const timeoutId = window.setTimeout( - () => { - console.warn('Power save blocker timeout - stopping automatically after 15 minutes'); - stopPowerSaveBlocker(); - }, - 15 * 60 * 1000 - ); - - setPowerSaveTimeoutId(timeoutId); - - const userMessage = createUserMessage(combinedTextFromInput.trim()); - - if (onSummaryReset) { - onSummaryReset(); - window.setTimeout(() => { - append(userMessage); - onMessageSent?.(); - }, 150); - } else { - append(userMessage); - onMessageSent?.(); - } - } else { - // If nothing was actually submitted (e.g. empty input and no images pasted) - stopPowerSaveBlocker(); - } - }, - [append, onMessageSent, stopPowerSaveBlocker] - ); - - // Handle stopping the message stream - const onStopGoose = useCallback(() => { - stop(); - setLastInteractionTime(Date.now()); - stopPowerSaveBlocker(); - - // Handle stopping the message stream - const lastMessage = messages[messages.length - 1]; - - // Check if there are any messages before proceeding - if (!lastMessage) { - return; - } - - // check if the last user message has any tool response(s) - const isToolResponse = lastMessage.content.some( - (content): content is ToolResponseMessageContent => content.type == 'toolResponse' - ); - - // isUserMessage also checks if the message is a toolConfirmationRequest - // check if the last message is a real user's message - if (lastMessage && isUserMessage(lastMessage) && !isToolResponse) { - const textValue = getTextContent(lastMessage); - _setInput(textValue); - - // Also add to local storage history as a backup so cmd+up can retrieve it - if (textValue.trim()) { - LocalMessageStorage.addMessage(textValue.trim()); - } - - // Remove the last user message if it's the most recent one - if (messages.length > 1) { - setMessages(messages.slice(0, -1)); - } else { - setMessages([]); - } - } else if (!isUserMessage(lastMessage)) { - const toolRequests: [string, Record][] = lastMessage.content - .filter( - (content) => content.type === 'toolRequest' || content.type === 'toolConfirmationRequest' - ) - .map((content) => { - if (content.type === 'toolRequest') { - return [content.id, content.toolCall]; - } else { - const toolCall = { - status: 'success', - value: { - name: content.toolName, - arguments: content.arguments, - }, - }; - return [content.id, toolCall]; - } - }); - - if (toolRequests.length !== 0) { - // This means we were interrupted during a tool request - // Create tool responses for all interrupted tool requests - - let responseMessage: Message = { - role: 'user', - created: Date.now(), - content: [], - metadata: { userVisible: true, agentVisible: true }, - }; - - const notification = 'Interrupted by the user to make a correction'; - - // generate a response saying it was interrupted for each tool request - for (const [reqId, _] of toolRequests) { - const toolResponse: ToolResponseMessageContent = { - type: 'toolResponse', - id: reqId, - toolResult: { - status: 'error', - error: notification, - }, - }; - - responseMessage.content.push(toolResponse); - } - // Use an immutable update to add the response message to the messages array - setMessages([...messages, responseMessage]); - } - } - }, [stop, messages, _setInput, setMessages, stopPowerSaveBlocker]); - - // Since server now handles all filtering, we just use messages directly - const filteredMessages = useMemo(() => { - return messages; - }, [messages]); - - // Generate command history from messages - const commandHistory = useMemo(() => { - return filteredMessages - .reduce((history, message) => { - if (isUserMessage(message)) { - const text = getTextContent(message).trim(); - if (text) { - history.push(text); - } - } - return history; - }, []) - .reverse(); - }, [filteredMessages]); - - // Process tool call notifications - const toolCallNotifications = useMemo(() => { - return notifications.reduce((map, item) => { - const key = item.request_id; - if (!map.has(key)) { - map.set(key, []); - } - map.get(key).push(item); - return map; - }, new Map()); - }, [notifications]); - - // Handle message updates from the UI - const onMessageUpdate = useCallback( - (messageId: string, newContent: string) => { - const messageIndex = messages.findIndex((msg) => msg.id === messageId); - - if (messageIndex !== -1) { - // Truncate the history to the point *before* the edited message. - const history = messages.slice(0, messageIndex); - - // Set the truncated history. - setMessages(history); - - // Instead of setTimeout, set pendingEdit which will be handled in useEffect - setPendingEdit({ id: messageId, content: newContent }); - } - }, - [messages, setMessages, setPendingEdit] - ); - - // Listen for pending edit and append message after messages updated - useEffect(() => { - if (pendingEdit) { - const updatedMessage = createUserMessage(pendingEdit.content); - append(updatedMessage); - setPendingEdit(null); // Reset after processing - } - }, [pendingEdit, append]); - - return { - // Core message data - messages, - filteredMessages, - - // Message stream controls - append, - stop, - chatState, - error, - setMessages, - - // Input controls - input: _input, - setInput: _setInput, - handleInputChange: _handleInputChange, - - // Event handlers - handleSubmit, - onStopGoose, - - // Token and session data - sessionTokenCount, - sessionInputTokens, - sessionOutputTokens, - localInputTokens, - localOutputTokens, - tokenState, - - // UI helpers - commandHistory, - toolCallNotifications, - - // Stream utilities - updateMessageStreamBody, - sessionMetadata: session, - - // Utilities - isUserMessage, - - // Error management - clearError: () => setError(undefined), - - // New functions for message editing - onMessageUpdate, - }; -}; diff --git a/ui/desktop/src/hooks/useChatStream.ts b/ui/desktop/src/hooks/useChatStream.ts index 5fd64b215c3c..022222f8a66e 100644 --- a/ui/desktop/src/hooks/useChatStream.ts +++ b/ui/desktop/src/hooks/useChatStream.ts @@ -14,10 +14,25 @@ import { import { createUserMessage, getCompactingMessage, getThinkingMessage } from '../types/message'; import { errorMessage } from '../utils/conversionUtils'; +import { LocalMessageStorage } from '../utils/localMessageStorage'; const resultsCache = new Map(); -type NotificationEvent = Extract; +// JSON value type for notification params +type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue }; + +// Base NotificationEvent from API +type BaseNotificationEvent = Extract; + +// Extended NotificationEvent with proper typing for message structure +export type NotificationEvent = Omit & { + message: { + method: string; + params: { + [key: string]: JsonValue; + }; + }; +}; interface UseChatStreamProps { sessionId: string; @@ -114,7 +129,7 @@ async function streamFromResponse( break; } case 'Notification': { - updateNotifications(event); + updateNotifications(event as NotificationEvent); break; } case 'Ping': @@ -150,6 +165,8 @@ export function useChatStream({ }); const [notifications, setNotifications] = useState([]); const abortControllerRef = useRef(null); + const [lastInteractionTime, setLastInteractionTime] = useState(Date.now()); + const [powerSaveTimeoutId, setPowerSaveTimeoutId] = useState(null); useEffect(() => { if (session) { @@ -166,17 +183,69 @@ export function useChatStream({ setNotifications((prev) => [...prev, notification]); }, []); + const stopPowerSaveBlocker = useCallback(() => { + try { + window.electron.stopPowerSaveBlocker(); + } catch (error) { + console.error('Failed to stop power save blocker:', error); + } + + if (powerSaveTimeoutId) { + window.clearTimeout(powerSaveTimeoutId); + setPowerSaveTimeoutId(null); + } + }, [powerSaveTimeoutId]); + const onFinish = useCallback( (error?: string): void => { + stopPowerSaveBlocker(); + if (error) { setSessionLoadError(error); } + + const timeSinceLastInteraction = Date.now() - lastInteractionTime; + + if (timeSinceLastInteraction > 60000) { + try { + window.electron.showNotification({ + title: 'Goose finished the task', + body: 'Click here for details', + }); + } catch (error) { + console.error('Failed to show notification:', error); + } + } else { + console.log('useChatStream: Not showing notification (task took less than 2 seconds)'); + } + + const isNewSession = sessionId && sessionId.match(/^\d{8}_\d{6}$/); + if (isNewSession) { + console.log( + 'useChatStream: Message stream finished for new session, emitting message-stream-finished event' + ); + window.dispatchEvent(new CustomEvent('message-stream-finished')); + } + setChatState(ChatState.Idle); onStreamFinish(); }, - [onStreamFinish] + [onStreamFinish, stopPowerSaveBlocker, lastInteractionTime, sessionId] ); + useEffect(() => { + return () => { + if (powerSaveTimeoutId) { + window.clearTimeout(powerSaveTimeoutId); + } + try { + window.electron.stopPowerSaveBlocker(); + } catch (error) { + console.error('Failed to stop power save blocker during cleanup:', error); + } + }; + }, [powerSaveTimeoutId]); + // Load session on mount or sessionId change useEffect(() => { if (!sessionId) return; @@ -236,6 +305,34 @@ export function useChatStream({ return; } + if (!userMessage.trim()) { + stopPowerSaveBlocker(); + return; + } + + LocalMessageStorage.addMessage(userMessage.trim()); + + try { + window.electron.startPowerSaveBlocker(); + } catch (error) { + console.error('Failed to start power save blocker:', error); + } + + setLastInteractionTime(Date.now()); + + const timeoutId = window.setTimeout( + () => { + console.warn('Power save blocker timeout - stopping automatically after 15 minutes'); + stopPowerSaveBlocker(); + }, + 15 * 60 * 1000 + ); + setPowerSaveTimeoutId(timeoutId); + + if (messagesRef.current.length === 0) { + window.dispatchEvent(new CustomEvent('session-created')); + } + const currentMessages = [...messagesRef.current, createUserMessage(userMessage)]; updateMessages(currentMessages); setChatState(ChatState.Streaming); @@ -272,7 +369,15 @@ export function useChatStream({ } } }, - [sessionId, session, chatState, updateMessages, updateNotifications, onFinish] + [ + sessionId, + session, + chatState, + updateMessages, + updateNotifications, + onFinish, + stopPowerSaveBlocker, + ] ); const setRecipeUserParams = useCallback( @@ -315,8 +420,10 @@ export function useChatStream({ const stopStreaming = useCallback(() => { abortControllerRef.current?.abort(); + setLastInteractionTime(Date.now()); + stopPowerSaveBlocker(); setChatState(ChatState.Idle); - }, []); + }, [stopPowerSaveBlocker]); const cached = resultsCache.get(sessionId); const maybe_cached_messages = session ? messages : cached?.messages || []; diff --git a/ui/desktop/src/hooks/useMessageStream.ts b/ui/desktop/src/hooks/useMessageStream.ts deleted file mode 100644 index 4e36c7d9ec76..000000000000 --- a/ui/desktop/src/hooks/useMessageStream.ts +++ /dev/null @@ -1,672 +0,0 @@ -import { useCallback, useEffect, useId, useReducer, useRef, useState } from 'react'; -import useSWR from 'swr'; -import { - createUserMessage, - getThinkingMessage, - getCompactingMessage, - hasCompletedToolCalls, -} from '../types/message'; -import { Conversation, Message, Role, TokenState } from '../api'; - -import { getSession, Session } from '../api'; -import { ChatState } from '../types/chatState'; - -let messageIdCounter = 0; - -function generateMessageId(): string { - return `msg-${Date.now()}-${++messageIdCounter}`; -} - -// Ensure TextDecoder is available in the global scope -const TextDecoder = globalThis.TextDecoder; - -type JsonValue = string | number | boolean | null | JsonValue[] | { [key: string]: JsonValue }; - -export interface NotificationEvent { - type: 'Notification'; - request_id: string; - message: { - method: string; - params: { - [key: string]: JsonValue; - }; - }; -} - -// Event types for SSE stream -type MessageEvent = - | { type: 'Message'; message: Message; token_state: TokenState } - | { type: 'Error'; error: string } - | { type: 'Finish'; reason: string; token_state: TokenState } - | { type: 'ModelChange'; model: string; mode: string } - | { type: 'UpdateConversation'; conversation: Conversation } - | NotificationEvent; - -export interface UseMessageStreamOptions { - /** - * The API endpoint that accepts a `{ messages: Message[] }` object and returns - * a stream of messages. Defaults to `/api/chat/reply`. - */ - api?: string; - - /** - * A unique identifier for the chat. If not provided, a random one will be - * generated. When provided, the hook with the same `id` will - * have shared states across components. - */ - id?: string; - - /** - * Initial messages of the chat. Useful to load an existing chat history. - */ - initialMessages?: Message[]; - - /** - * Initial input of the chat. - */ - initialInput?: string; - - /** - * Callback function to be called when a tool call is received. - * You can optionally return a result for the tool call. - */ - _onToolCall?: (toolCall: Record) => void | Promise | unknown; - - /** - * Callback function to be called when the API response is received. - */ - onResponse?: (response: Response) => void | Promise; - - /** - * Callback function to be called when the assistant message is finished streaming. - */ - onFinish?: (message: Message, reason: string) => void; - - /** - * Callback function to be called when an error is encountered. - */ - onError?: (error: Error) => void; - - /** - * HTTP headers to be sent with the API request. - */ - headers?: Record | HeadersInit; - - /** - * Extra body object to be sent with the API request. - */ - body?: object; - - /** - * Maximum number of sequential LLM calls (steps), e.g. when you use tool calls. - * Default is 1. - */ - maxSteps?: number; -} - -export interface UseMessageStreamHelpers { - /** Current messages in the chat */ - messages: Message[]; - - /** The error object of the API request */ - error: undefined | Error; - - /** - * Append a user message to the chat list. This triggers the API call to fetch - * the assistant's response. - */ - append: (message: Message | string) => Promise; - - /** - * Reload the last AI chat response for the given chat history. - */ - reload: () => Promise; - - /** - * Abort the current request immediately. - */ - stop: () => void; - - /** - * Update the `messages` state locally. - */ - setMessages: (messages: Message[] | ((messages: Message[]) => Message[])) => void; - - /** The current value of the input */ - input: string; - - /** setState-powered method to update the input value */ - setInput: React.Dispatch>; - - /** An input/textarea-ready onChange handler to control the value of the input */ - handleInputChange: ( - e: React.ChangeEvent | React.ChangeEvent - ) => void; - - /** Form submission handler to automatically reset input and append a user message */ - handleSubmit: (event?: { preventDefault?: () => void }) => void; - - /** Current chat state (idle, thinking, streaming, waiting for user input) */ - chatState: ChatState; - - /** Add a tool result to a tool call */ - addToolResult: ({ toolCallId, result }: { toolCallId: string; result: unknown }) => void; - - /** Modify body (session id and/or work dir mid-stream) **/ - updateMessageStreamBody?: (newBody: object) => void; - - notifications: NotificationEvent[]; - - /** Current model info from the backend */ - currentModelInfo: { model: string; mode: string } | null; - - /** Session including token counts */ - session: Session | null; - - /** Clear error state */ - setError: (error: Error | undefined) => void; - - /** Real-time token state from server */ - tokenState: TokenState; -} - -/** - * Hook for streaming messages directly from the server using the native Goose message format - */ -export function useMessageStream({ - api = '/api/chat/reply', - id, - initialMessages = [], - initialInput = '', - onResponse, - onFinish, - onError, - headers, - body, - maxSteps = 1, -}: UseMessageStreamOptions = {}): UseMessageStreamHelpers { - // Generate a unique id for the chat if not provided - const hookId = useId(); - const idKey = id ?? hookId; - const chatKey = typeof api === 'string' ? [api, idKey] : idKey; - - // Store the chat state in SWR, using the chatId as the key to share states - const { data: messages, mutate } = useSWR([chatKey, 'messages'], null, { - fallbackData: initialMessages, - }); - - const [notifications, setNotifications] = useState([]); - const [currentModelInfo, setCurrentModelInfo] = useState<{ model: string; mode: string } | null>( - null - ); - const [session, setSession] = useState(null); - const [tokenState, setTokenState] = useState({ - inputTokens: 0, - outputTokens: 0, - totalTokens: 0, - accumulatedInputTokens: 0, - accumulatedOutputTokens: 0, - accumulatedTotalTokens: 0, - }); - - // expose a way to update the body so we can update the session id when CLE occurs - const updateMessageStreamBody = useCallback((newBody: object) => { - extraMetadataRef.current.body = { - ...extraMetadataRef.current.body, - ...newBody, - }; - }, []); - - // Keep the latest messages in a ref - const messagesRef = useRef(messages || []); - useEffect(() => { - messagesRef.current = messages || []; - }, [messages]); - - // Track chat state (idle, thinking, streaming, waiting for user input) - const { data: chatState = ChatState.Idle, mutate: mutateChatState } = useSWR( - [chatKey, 'chatState'], - null - ); - - const { data: error = undefined, mutate: setError } = useSWR( - [chatKey, 'error'], - null - ); - - // Abort controller to cancel the current API call - const abortControllerRef = useRef(null); - - // Extra metadata for requests - const extraMetadataRef = useRef({ - headers, - body, - }); - - useEffect(() => { - extraMetadataRef.current = { - headers, - body, - }; - }, [headers, body]); - - // TODO: not this? - const [, forceUpdate] = useReducer((x) => x + 1, 0); - - // Process the SSE stream from the server - const processMessageStream = useCallback( - async (response: Response, currentMessages: Message[]) => { - if (!response.body) { - throw new Error('Response body is empty'); - } - - const reader = response.body.getReader(); - const decoder = new TextDecoder(); - let buffer = ''; - - try { - let running = true; - while (running) { - const { done, value } = await reader.read(); - if (done) { - running = false; - break; - } - - // Decode the chunk and add it to our buffer - buffer += decoder.decode(value, { stream: true }); - - // Process complete SSE events - const events = buffer.split('\n\n'); - buffer = events.pop() || ''; // Keep the last incomplete event in the buffer - - for (const event of events) { - if (event.startsWith('data: ')) { - try { - const data = event.slice(6); // Remove 'data: ' prefix - const parsedEvent = JSON.parse(data) as MessageEvent; - - switch (parsedEvent.type) { - case 'Message': { - // Transition from waiting to streaming on first message - mutateChatState(ChatState.Streaming); - - setTokenState(parsedEvent.token_state); - - // Create a new message object with the properties preserved or defaulted - const newMessage: Message = { - ...parsedEvent.message, - id: parsedEvent.message.id || undefined, - role: parsedEvent.message.role as Role, - created: parsedEvent.message.created || Date.now(), - content: parsedEvent.message.content || [], - }; - - // Update messages with the new message - if ( - newMessage.id && - currentMessages.length > 0 && - currentMessages[currentMessages.length - 1].id === newMessage.id - ) { - // If the last message has the same ID, update it instead of adding a new one - const lastMessage = currentMessages[currentMessages.length - 1]; - lastMessage.content = [...lastMessage.content, ...newMessage.content]; - forceUpdate(); - } else { - currentMessages = [...currentMessages, newMessage]; - } - - // Check if this message contains tool confirmation requests - const hasToolConfirmation = newMessage.content.some( - (content) => content.type === 'toolConfirmationRequest' - ); - - if (hasToolConfirmation) { - mutateChatState(ChatState.WaitingForUserInput); - } - - if (getCompactingMessage(newMessage)) { - mutateChatState(ChatState.Compacting); - } else if (getThinkingMessage(newMessage)) { - mutateChatState(ChatState.Thinking); - } - - mutate(currentMessages, false); - break; - } - - case 'Notification': { - const newNotification = { - ...parsedEvent, - }; - setNotifications((prev) => [...prev, newNotification]); - break; - } - - case 'ModelChange': { - // Update the current model in the frontend - const modelInfo = { - model: parsedEvent.model, - mode: parsedEvent.mode, - }; - setCurrentModelInfo(modelInfo); - break; - } - - case 'UpdateConversation': { - // WARNING: Since Message handler uses this local variable, we need to update it here to avoid the client clobbering it. - // Longterm fix is to only send the agent the new messages, not the entire conversation. - currentMessages = parsedEvent.conversation; - setMessages(parsedEvent.conversation); - break; - } - - case 'Error': { - // Always throw the error so it gets caught and sets the error state - // This ensures the retry UI appears for ALL errors - throw new Error(parsedEvent.error); - } - - case 'Finish': { - setTokenState(parsedEvent.token_state); - - if (onFinish && currentMessages.length > 0) { - const lastMessage = currentMessages[currentMessages.length - 1]; - onFinish(lastMessage, parsedEvent.reason); - } - - const sessionId = (extraMetadataRef.current.body as Record) - ?.session_id as string; - if (sessionId) { - const sessionResponse = await getSession({ - path: { session_id: sessionId }, - throwOnError: true, - }); - - if (sessionResponse.data) { - setSession(sessionResponse.data); - } - } - break; - } - } - } catch (e) { - console.error('Error parsing SSE event:', e); - if (onError && e instanceof Error) { - onError(e); - } - // Don't re-throw here, let the error be handled by the outer catch - // Instead, set the error state directly - if (e instanceof Error) { - setError(e); - } - } - } - } - } - } catch (e) { - if (e instanceof Error && e.name !== 'AbortError') { - console.error('Error reading SSE stream:', e); - if (onError) { - onError(e); - } - // Re-throw the error so it gets caught by sendRequest and sets the error state - throw e; - } - } finally { - reader.releaseLock(); - } - - return currentMessages; - }, - // eslint-disable-next-line react-hooks/exhaustive-deps - [mutate, mutateChatState, onFinish, onError, forceUpdate, setError] - ); - - // Send a request to the server - const sendRequest = useCallback( - async (requestMessages: Message[]) => { - try { - mutateChatState(ChatState.Thinking); // Start in thinking state - setError(undefined); - - // Create abort controller - const abortController = new AbortController(); - abortControllerRef.current = abortController; - - // Send request to the server - const response = await fetch(api, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'X-Secret-Key': await window.electron.getSecretKey(), - ...extraMetadataRef.current.headers, - }, - body: JSON.stringify({ - messages: requestMessages, - ...extraMetadataRef.current.body, - }), - signal: abortController.signal, - }); - - if (onResponse) { - await onResponse(response); - } - - if (!response.ok) { - const text = await response.text(); - throw new Error(text || `Error ${response.status}: ${response.statusText}`); - } - - // Process the SSE stream - const updatedMessages = await processMessageStream(response, requestMessages); - - // Auto-submit when all tool calls in the last assistant message have results - if (maxSteps > 1 && updatedMessages.length > requestMessages.length) { - const lastMessage = updatedMessages[updatedMessages.length - 1]; - if (lastMessage.role === 'assistant' && hasCompletedToolCalls(lastMessage)) { - // Count trailing assistant messages to prevent infinite loops - let assistantCount = 0; - for (let i = updatedMessages.length - 1; i >= 0; i--) { - if (updatedMessages[i].role === 'assistant') { - assistantCount++; - } else { - break; - } - } - - if (assistantCount < maxSteps) { - await sendRequest(updatedMessages); - } - } - } - - abortControllerRef.current = null; - } catch (err) { - // Ignore abort errors as they are expected - if (err instanceof Error && err.name === 'AbortError') { - abortControllerRef.current = null; - return; - } - - if (onError && err instanceof Error) { - onError(err); - } - - setError(err as Error); - } finally { - // Check if the last message has pending tool confirmations - const currentMessages = messagesRef.current; - const lastMessage = currentMessages[currentMessages.length - 1]; - const hasPendingToolConfirmation = lastMessage?.content.some( - (content) => content.type === 'toolConfirmationRequest' - ); - - if (hasPendingToolConfirmation) { - mutateChatState(ChatState.WaitingForUserInput); - } else { - mutateChatState(ChatState.Idle); - } - } - }, - - [api, processMessageStream, mutateChatState, setError, onResponse, onError, maxSteps] - ); - - // Append a new message and send request - const append = useCallback( - async (message: Message | string) => { - // If a string is passed, convert it to a Message object - const messageToAppend = typeof message === 'string' ? createUserMessage(message) : message; - - // If we were waiting for user input and user provides input, transition away from that state - if (chatState === ChatState.WaitingForUserInput) { - mutateChatState(ChatState.Thinking); - } - - const currentMessages = [...messagesRef.current, messageToAppend]; - mutate(currentMessages, false); - await sendRequest(currentMessages); - }, - [mutate, sendRequest, chatState, mutateChatState] - ); - - // Reload the last message - const reload = useCallback(async () => { - const currentMessages = messagesRef.current; - if (currentMessages.length === 0) { - return; - } - - // Remove last assistant message if present - const lastMessage = currentMessages[currentMessages.length - 1]; - const messagesToSend = - lastMessage.role === 'assistant' ? currentMessages.slice(0, -1) : currentMessages; - - await sendRequest(messagesToSend); - }, [sendRequest]); - - // Stop the current request - const stop = useCallback(() => { - if (abortControllerRef.current) { - abortControllerRef.current.abort(); - abortControllerRef.current = null; - } - }, []); - - // Set messages directly - const setMessages = useCallback( - (messagesOrFn: Message[] | ((messages: Message[]) => Message[])) => { - if (typeof messagesOrFn === 'function') { - const newMessages = messagesOrFn(messagesRef.current); - mutate(newMessages, false); - messagesRef.current = newMessages; - } else { - mutate(messagesOrFn, false); - messagesRef.current = messagesOrFn; - } - }, - [mutate] - ); - - // Input state and handlers - const [input, setInput] = useState(initialInput); - - const handleInputChange = useCallback( - (e: React.ChangeEvent | React.ChangeEvent) => { - setInput(e.target.value); - }, - [] - ); - - const handleSubmit = useCallback( - async (event?: { preventDefault?: () => void }) => { - event?.preventDefault?.(); - if (!input.trim()) return; - - await append(input); - setInput(''); - }, - [input, append] - ); - - // Add tool result to a message - const addToolResult = useCallback( - ({ toolCallId, result }: { toolCallId: string; result: unknown }) => { - const currentMessages = messagesRef.current; - - // Find the last assistant message with the tool call - let lastAssistantIndex = -1; - for (let i = currentMessages.length - 1; i >= 0; i--) { - if (currentMessages[i].role === 'assistant') { - const toolRequests = currentMessages[i].content.filter( - (content) => content.type === 'toolRequest' && content.id === toolCallId - ); - if (toolRequests.length > 0) { - lastAssistantIndex = i; - break; - } - } - } - - if (lastAssistantIndex === -1) return; - - // Create a tool response message - const toolResponseMessage: Message = { - id: generateMessageId(), - role: 'user' as const, - created: Math.floor(Date.now() / 1000), - metadata: { userVisible: true, agentVisible: true }, - content: [ - { - type: 'toolResponse' as const, - id: toolCallId, - toolResult: { - status: 'success' as const, - value: Array.isArray(result) - ? result - : [{ type: 'text' as const, text: String(result), priority: 0 }], - }, - }, - ], - }; - - // Insert the tool response after the assistant message - const updatedMessages = [ - ...currentMessages.slice(0, lastAssistantIndex + 1), - toolResponseMessage, - ...currentMessages.slice(lastAssistantIndex + 1), - ]; - - mutate(updatedMessages, false); - messagesRef.current = updatedMessages; - - // Auto-submit if we have tool results - if (maxSteps > 1) { - sendRequest(updatedMessages); - } - }, - [mutate, maxSteps, sendRequest] - ); - - return { - messages: messages || [], - error, - append, - reload, - stop, - setMessages, - input, - setInput, - handleInputChange, - handleSubmit, - chatState, - addToolResult, - updateMessageStreamBody, - notifications, - currentModelInfo, - session, - setError, - tokenState, - }; -} diff --git a/ui/desktop/src/main.ts b/ui/desktop/src/main.ts index d4521d897911..27c583e53b64 100644 --- a/ui/desktop/src/main.ts +++ b/ui/desktop/src/main.ts @@ -2086,7 +2086,7 @@ async function appMain() { } }); - ipcMain.on('notify', (_event, data) => { + ipcMain.on('notify', (event, data) => { try { // Validate notification data if (!data || typeof data !== 'object') { @@ -2111,10 +2111,24 @@ async function appMain() { const sanitizeText = (text: string) => text.replace(/<[^>]*>/g, ''); console.log('NOTIFY', data); - new Notification({ + const notification = new Notification({ title: sanitizeText(data.title), body: sanitizeText(data.body), - }).show(); + }); + + // Add click handler to focus the window + notification.on('click', () => { + const window = BrowserWindow.fromWebContents(event.sender); + if (window) { + if (window.isMinimized()) { + window.restore(); + } + window.show(); + window.focus(); + } + }); + + notification.show(); } catch (error) { console.error('Error showing notification:', error); }