diff --git a/apps/studio/src/lib/editor/engine/chat/index.ts b/apps/studio/src/lib/editor/engine/chat/index.ts index c21427567..c4e8db942 100644 --- a/apps/studio/src/lib/editor/engine/chat/index.ts +++ b/apps/studio/src/lib/editor/engine/chat/index.ts @@ -5,14 +5,16 @@ import { ChatMessageRole, StreamRequestType, type AssistantChatMessage, + type ChatMessageContext, type CompletedStreamResponse, type ErrorStreamResponse, + type QueuedMessage, type RateLimitedStreamResponse, } from '@onlook/models/chat'; import { MainChannels } from '@onlook/models/constants'; import type { ParsedError } from '@onlook/utility'; import type { CoreMessage } from 'ai'; -import { makeAutoObservable } from 'mobx'; +import { makeAutoObservable, runInAction } from 'mobx'; import { nanoid } from 'nanoid/non-secure'; import type { EditorEngine } from '..'; import { ChatCodeManager } from './code'; @@ -31,13 +33,17 @@ export class ChatManager { context: ChatContext; stream: StreamResolver; suggestions: SuggestionManager; + messageQueue: QueuedMessage[] = []; + private maxQueueSize = 5; constructor( private editorEngine: EditorEngine, private projectsManager: ProjectsManager, private userManager: UserManager, ) { - makeAutoObservable(this); + makeAutoObservable(this, { + messageQueue: true, + }); this.context = new ChatContext(this.editorEngine, this.projectsManager); this.conversation = new ConversationManager(this.editorEngine, this.projectsManager); this.stream = new StreamResolver(); @@ -49,14 +55,41 @@ export class ChatManager { window.dispatchEvent(new Event(FOCUS_CHAT_INPUT_EVENT)); } - async sendNewMessage(content: string): Promise { + get queueSize(): number { + return this.messageQueue.length; + } + + async processMessageQueue() { + if (this.messageQueue.length === 0 || this.isWaiting) { + return; + } + const nextMessage = this.messageQueue.shift()!; + await this.processMessage(nextMessage.content, nextMessage.context); + + if (this.messageQueue.length > 0) { + await this.processMessageQueue(); + } + } + + private async processMessage(content: string, context?: ChatMessageContext[]) { if (!this.conversation.current) { console.error('No conversation found'); return; } - const context = await this.context.getChatContext(); - const userMessage = this.conversation.addUserMessage(content, context); + if (this.isWaiting) { + if (this.messageQueue.length >= this.maxQueueSize) { + console.warn('Message queue is full'); + return; + } + runInAction(() => { + this.messageQueue.push({ content, context }); + }); + return; + } + + const messageContext = context ?? (await this.context.getChatContext()); + const userMessage = this.conversation.addUserMessage(content, messageContext); this.conversation.current.updateName(content); if (!userMessage) { console.error('Failed to add user message'); @@ -68,6 +101,27 @@ export class ChatManager { await this.sendChatToAi(StreamRequestType.CHAT, content); } + async sendNewMessage(content: string): Promise { + if (!this.conversation.current) { + console.error('No conversation found'); + return; + } + + if (this.isWaiting) { + if (this.messageQueue.length >= this.maxQueueSize) { + console.warn('Message queue is full'); + return; + } + runInAction(() => { + this.messageQueue.push({ content }); + }); + return; + } + + await this.processMessage(content); + this.processMessageQueue(); + } + async sendFixErrorToAi(errors: ParsedError[]): Promise { if (!this.conversation.current) { console.error('No conversation found'); @@ -77,6 +131,22 @@ export class ChatManager { const prompt = `How can I resolve these errors? If you propose a fix, please make it concise.`; const errorContexts = this.context.getMessageContext(errors); const projectContexts = this.context.getProjectContext(); + + if (this.isWaiting) { + if (this.messageQueue.length >= this.maxQueueSize) { + console.warn('Message queue is full'); + return false; + } + runInAction(() => { + this.messageQueue.push({ + content: prompt, + context: [...errorContexts, ...projectContexts], + }); + }); + console.log(`Error-fix message queued. Queue size: ${this.messageQueue.length}`); + return true; + } + const userMessage = this.conversation.addUserMessage(prompt, [ ...errorContexts, ...projectContexts, @@ -139,6 +209,10 @@ export class ChatManager { invokeMainChannel(MainChannels.SEND_STOP_STREAM_REQUEST, { requestId, }); + + runInAction(() => { + this.messageQueue = []; + }); sendAnalytics('stop chat stream'); } @@ -204,6 +278,8 @@ export class ChatManager { } this.context.clearAttachments(); + + await this.processMessageQueue(); } handleNewCoreMessages(messages: CoreMessage[]) { diff --git a/apps/studio/src/routes/editor/EditPanel/ChatTab/ChatInput.tsx b/apps/studio/src/routes/editor/EditPanel/ChatTab/ChatInput.tsx index 7b632be68..3baa23b2e 100644 --- a/apps/studio/src/routes/editor/EditPanel/ChatTab/ChatInput.tsx +++ b/apps/studio/src/routes/editor/EditPanel/ChatTab/ChatInput.tsx @@ -29,6 +29,9 @@ export const ChatInput = observer(() => { const [actionTooltipOpen, setActionTooltipOpen] = useState(false); const [isDragging, setIsDragging] = useState(false); + const queueSize = editorEngine.chat.queueSize; + const MAX_QUEUE_SIZE = 5; + const focusInput = () => { requestAnimationFrame(() => { textareaRef.current?.focus(); @@ -428,6 +431,11 @@ export const ChatInput = observer(() => { File Reference + {queueSize > 0 && ( + + {queueSize} message{queueSize > 1 ? 's' : ''} queued + + )} {editorEngine.chat.isWaiting ? ( @@ -450,7 +458,10 @@ export const ChatInput = observer(() => { size={'icon'} variant={'secondary'} className="text-smallPlus w-fit h-full py-0.5 px-2.5 text-primary" - disabled={inputEmpty || editorEngine.chat.isWaiting} + disabled={ + inputEmpty || + (editorEngine.chat.isWaiting && queueSize >= MAX_QUEUE_SIZE) + } onClick={sendMessage} > diff --git a/packages/models/src/chat/conversation/index.ts b/packages/models/src/chat/conversation/index.ts index 6fd12d905..b56ced242 100644 --- a/packages/models/src/chat/conversation/index.ts +++ b/packages/models/src/chat/conversation/index.ts @@ -1,4 +1,9 @@ -import type { AssistantChatMessage, ChatMessage, TokenUsage } from '../message/index.ts'; +import type { + AssistantChatMessage, + ChatMessage, + ChatMessageContext, + TokenUsage, +} from '../message/index.ts'; export type ChatConversation = { id: string; @@ -10,3 +15,8 @@ export type ChatConversation = { summaryMessage?: AssistantChatMessage | null; tokenUsage?: TokenUsage; }; + +export type QueuedMessage = { + content: string; + context?: ChatMessageContext[]; +};