feat: use single onlook message format

apps/web/client/src/app/api/chat/helpers/stream.ts

@@ -1,5 +1,5 @@
 import type { ToolCall } from '@ai-sdk/provider-utils';
-import { ASK_TOOL_SET, BUILD_TOOL_SET, getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt, initModel } from '@onlook/ai';
+import { getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt, initModel } from '@onlook/ai';
 import { ChatType, LLMProvider, OPENROUTER_MODELS, type ModelConfig } from '@onlook/models';
 import { generateObject, NoSuchToolError, type ToolSet } from 'ai';
 
@@ -25,11 +25,7 @@ export async function getModelFromType(chatType: ChatType) {
     return model;
 }
 
-export function getToolSetFromType(chatType: ChatType) {
-    return chatType === ChatType.ASK ? ASK_TOOL_SET : BUILD_TOOL_SET;
-}
-
-export async function getSystemPromptFromType(chatType: ChatType) {
+export function getSystemPromptFromType(chatType: ChatType) {
     let systemPrompt: string;
 
     switch (chatType) {

apps/web/client/src/app/api/chat/route.ts

@@ -1,9 +1,10 @@
 import { trackEvent } from '@/utils/analytics/server';
-import { ChatType } from '@onlook/models';
-import { convertToModelMessages, stepCountIs, streamText, type UIMessage } from 'ai';
+import { convertToStreamMessages, getToolSetFromType } from '@onlook/ai';
+import { ChatType, type ChatMessage } from '@onlook/models';
+import { stepCountIs, streamText } from 'ai';
 import { type NextRequest } from 'next/server';
 import { v4 as uuidv4 } from 'uuid';
-import { checkMessageLimit, decrementUsage, errorHandler, getModelFromType, getSupabaseUser, getSystemPromptFromType, getToolSetFromType, incrementUsage, repairToolCall } from './helpers';
+import { checkMessageLimit, decrementUsage, errorHandler, getModelFromType, getSupabaseUser, getSystemPromptFromType, incrementUsage, repairToolCall } from './helpers';
 
 const MAX_STEPS = 20;
 
@@ -54,7 +55,7 @@ export async function POST(req: NextRequest) {
 export const streamResponse = async (req: NextRequest, userId: string) => {
     const body = await req.json();
     const { messages, chatType, conversationId, projectId } = body as {
-        messages: UIMessage[],
+        messages: ChatMessage[],
         chatType: ChatType,
         conversationId: string,
         projectId: string,
@@ -68,16 +69,16 @@ export const streamResponse = async (req: NextRequest, userId: string) => {
     } | null = null;
 
     try {
-        const lastUserMessage = messages.findLast((message: UIMessage) => message.role === 'user');
+        const lastUserMessage = messages.findLast((message) => message.role === 'user');
         const traceId = lastUserMessage?.id ?? uuidv4();
 
         if (chatType === ChatType.EDIT) {
             usageRecord = await incrementUsage(req, traceId);
         }
         const modelConfig = await getModelFromType(chatType);
         const { model, providerOptions, headers } = modelConfig;
-        const systemPrompt = await getSystemPromptFromType(chatType);
-        const tools = await getToolSetFromType(chatType);
+        const systemPrompt = getSystemPromptFromType(chatType);
+        const tools = getToolSetFromType(chatType);
 
         const result = streamText({
             model,
@@ -90,7 +91,7 @@ export const streamResponse = async (req: NextRequest, userId: string) => {
                     content: systemPrompt,
                     providerOptions,
                 },
-                ...convertToModelMessages(messages),
+                ...convertToStreamMessages(messages),
             ],
             experimental_telemetry: {
                 isEnabled: true,
@@ -123,13 +124,14 @@ export const streamResponse = async (req: NextRequest, userId: string) => {
         return result.toUIMessageStreamResponse(
             {
                 originalMessages: messages,
-                messageMetadata: ({
-                    part
-                }) => {
-                    if (part.type === 'finish-step') {
-                        return {
-                            finishReason: part.finishReason,
-                        }
+                generateMessageId: () => uuidv4(),
+                messageMetadata: ({ part }) => {
+                    return {
+                        createdAt: new Date(),
+                        conversationId,
+                        context: [],
+                        checkpoints: [],
+                        finishReason: part.type === 'finish-step' ? part.finishReason : undefined,
                     }
                 },
                 onError: errorHandler,

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/assistant-message.tsx

@@ -1,7 +1,7 @@
-import type { AssistantChatMessage } from '@onlook/models';
+import type { ChatMessage } from '@onlook/models';
 import { MessageContent } from './message-content';
 
-export const AssistantMessage = ({ message }: { message: AssistantChatMessage }) => {
+export const AssistantMessage = ({ message }: { message: ChatMessage }) => {
     return (
         <div className="px-4 py-2 text-small content-start flex flex-col text-wrap gap-2">
             <MessageContent

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/index.tsx

@@ -1,7 +1,7 @@
 import { useChatContext } from '@/app/project/[id]/_hooks/use-chat';
 import { useEditorEngine } from '@/components/store/editor';
 import { transKeys } from '@/i18n/keys';
-import { ChatMessageRole, type ChatMessage } from '@onlook/models/chat';
+import { type ChatMessage } from '@onlook/models/chat';
 import { ChatMessageList } from '@onlook/ui/chat/chat-message-list';
 import { Icons } from '@onlook/ui/icons';
 import { assertNever } from '@onlook/utility';
@@ -23,21 +23,24 @@ export const ChatMessages = observer(() => {
     const renderMessage = useCallback((message: ChatMessage, index: number) => {
         let messageNode;
         switch (message.role) {
-            case ChatMessageRole.ASSISTANT:
+            case 'assistant':
                 messageNode = <AssistantMessage message={message} />;
                 break;
-            case ChatMessageRole.USER:
+            case 'user':
                 messageNode = <UserMessage message={message} />;
                 break;
+            case 'system':
+                messageNode = null;
+                break;
             default:
-                assertNever(message);
+                assertNever(message.role);
         }
         return <div key={`message-${message.id}-${index}`}>{messageNode}</div>;
     }, []);
 
     // Exclude the currently streaming assistant message (rendered by <StreamMessage />)
     const messagesToRender = useMemo(() => {
-        if (!engineMessages || engineMessages.length === 0) return [] as ChatMessage[];
+        if (!engineMessages || engineMessages.length === 0) return [];
 
         const lastUiMessage = uiMessages?.[uiMessages.length - 1];
         const streamingAssistantId = isWaiting && lastUiMessage?.role === 'assistant' ? lastUiMessage.id : undefined;

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/message-content/index.tsx

@@ -1,8 +1,8 @@
 import { Icons } from '@onlook/ui/icons/index';
 import { cn } from '@onlook/ui/utils';
-import type { ToolUIPart, UIMessage } from 'ai';
+import type { ToolUIPart } from 'ai';
+import type { ChatMessage } from '@onlook/models';
 import { observer } from 'mobx-react-lite';
-import { useMemo } from 'react';
 import { MarkdownRenderer } from '../markdown-renderer';
 import { ToolCallDisplay } from './tool-call-display';
 
@@ -14,7 +14,7 @@ export const MessageContent = observer(
         isStream,
     }: {
         messageId: string;
-        parts: UIMessage['parts'];
+        parts: ChatMessage['parts'];
         applied: boolean;
         isStream: boolean;
     }) => {

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/stream-message.tsx

@@ -1,5 +1,4 @@
 import { useChatContext } from '@/app/project/[id]/_hooks/use-chat';
-import { ChatMessageRole } from '@onlook/models/chat';
 import { Icons } from '@onlook/ui/icons';
 import { useMemo } from 'react';
 import { MessageContent } from './message-content';
@@ -8,7 +7,7 @@ export const StreamMessage = () => {
     const { messages, isWaiting } = useChatContext();
     const streamMessage = messages.length > 0 ? messages[messages.length - 1] : null;
     const isAssistantStreamMessage = useMemo(() =>
-        streamMessage?.role === ChatMessageRole.ASSISTANT,
+        streamMessage?.role === 'assistant',
         [streamMessage?.role]
     );
 

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/user-message.tsx

@@ -1,6 +1,8 @@
+'use client';
+
 import { useChatContext } from '@/app/project/[id]/_hooks/use-chat';
 import { useEditorEngine } from '@/components/store/editor';
-import { ChatType, MessageCheckpointType, type UserChatMessage } from '@onlook/models';
+import { ChatType, MessageCheckpointType, type ChatMessage } from '@onlook/models';
 import { Button } from '@onlook/ui/button';
 import { Icons } from '@onlook/ui/icons';
 import { toast } from '@onlook/ui/sonner';
@@ -12,11 +14,7 @@ import React, { useEffect, useRef, useState } from 'react';
 import { SentContextPill } from '../context-pills/sent-context-pill';
 import { MessageContent } from './message-content';
 
-interface UserMessageProps {
-    message: UserChatMessage;
-}
-
-export const getUserMessageContent = (message: UserChatMessage) => {
+export const getUserMessageContent = (message: ChatMessage) => {
     return message.parts.map((part) => {
         if (part.type === 'text') {
             return part.text;
@@ -25,7 +23,7 @@ export const getUserMessageContent = (message: UserChatMessage) => {
     }).join('');
 }
 
-export const UserMessage = ({ message }: UserMessageProps) => {
+export const UserMessage = ({ message }: { message: ChatMessage }) => {
     const editorEngine = useEditorEngine();
     const { sendMessageToChat } = useChatContext();
     const [isCopied, setIsCopied] = useState(false);

apps/web/client/src/app/project/[id]/_hooks/use-chat.tsx

@@ -3,25 +3,26 @@
 import { useEditorEngine } from '@/components/store/editor';
 import { handleToolCall } from '@/components/tools';
 import { useChat, type UseChatHelpers } from '@ai-sdk/react';
-import { toVercelMessageFromOnlook } from '@onlook/ai';
-import { toOnlookMessageFromVercel } from '@onlook/db';
-import { ChatMessageRole, ChatType } from '@onlook/models';
-import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls, type UIMessage } from 'ai';
+import { ChatType, type ChatMessage } from '@onlook/models';
+import { jsonClone } from '@onlook/utility';
+import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls } from 'ai';
 import { observer } from 'mobx-react-lite';
 import { usePostHog } from 'posthog-js/react';
 import { createContext, useContext, useRef } from 'react';
+import { v4 as uuidv4 } from 'uuid';
 
-type ExtendedUseChatHelpers = UseChatHelpers<UIMessage> & { sendMessageToChat: (type: ChatType) => Promise<void> };
+type ExtendedUseChatHelpers = UseChatHelpers<ChatMessage> & { sendMessageToChat: (type: ChatType) => Promise<void> };
 const ChatContext = createContext<ExtendedUseChatHelpers | null>(null);
 
 export const ChatProvider = observer(({ children }: { children: React.ReactNode }) => {
     const editorEngine = useEditorEngine();
-    const lastMessageRef = useRef<UIMessage | null>(null);
+    const lastMessageRef = useRef<ChatMessage | null>(null);
     const posthog = usePostHog();
 
     const conversationId = editorEngine.chat.conversation.current?.conversation.id;
-    const chat = useChat({
+    const chat = useChat<ChatMessage>({
         id: 'user-chat',
+        generateId: () => uuidv4(),
         sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
         transport: new DefaultChatTransport({
             api: '/api/chat',
@@ -42,8 +43,7 @@ export const ChatProvider = observer(({ children }: { children: React.ReactNode
             }
 
             if (finishReason !== 'tool-calls') {
-                const currentConversationId = editorEngine.chat.conversation.current?.conversation.id;
-                editorEngine.chat.conversation.addOrReplaceMessage(toOnlookMessageFromVercel(message, currentConversationId ?? ''));
+                editorEngine.chat.conversation.addOrReplaceMessage(message);
                 editorEngine.chat.suggestions.generateSuggestions();
                 lastMessageRef.current = null;
             }
@@ -59,8 +59,7 @@ export const ChatProvider = observer(({ children }: { children: React.ReactNode
             console.error('Error in chat', error);
             editorEngine.chat.error.handleChatError(error);
             if (lastMessageRef.current) {
-                const currentConversationId = editorEngine.chat.conversation.current?.conversation.id;
-                editorEngine.chat.conversation.addOrReplaceMessage(toOnlookMessageFromVercel(lastMessageRef.current, currentConversationId ?? ''));
+                editorEngine.chat.conversation.addOrReplaceMessage(lastMessageRef.current);
                 lastMessageRef.current = null;
             }
         }
@@ -74,16 +73,7 @@ export const ChatProvider = observer(({ children }: { children: React.ReactNode
         editorEngine.chat.error.clear();
 
         const messages = editorEngine.chat.conversation.current?.messages ?? [];
-        const uiMessages = messages.map((message, index) =>
-            toVercelMessageFromOnlook(message, {
-                totalMessages: messages.length,
-                currentMessageIndex: index,
-                lastUserMessageIndex: messages.findLastIndex(m => m.role === ChatMessageRole.USER),
-                lastAssistantMessageIndex: messages.findLastIndex(m => m.role === ChatMessageRole.ASSISTANT),
-            })
-        );
-
-        chat.setMessages(uiMessages);
+        chat.setMessages(jsonClone(messages));
         try {
             posthog.capture('user_send_message', {
                 type,

apps/web/client/src/components/store/editor/chat/conversation.ts

@@ -1,7 +1,7 @@
 import { api } from '@/trpc/client';
 import { toDbMessage } from '@onlook/db';
 import type { GitCommit } from '@onlook/git';
-import { ChatMessageRole, MessageCheckpointType, type ChatConversation, type ChatMessage, type MessageContext, type UserChatMessage } from '@onlook/models';
+import { MessageCheckpointType, type ChatConversation, type ChatMessage, type MessageContext } from '@onlook/models';
 import { makeAutoObservable } from 'mobx';
 import { toast } from 'sonner';
 import type { EditorEngine } from '../engine';
@@ -114,7 +114,7 @@ export class ConversationManager {
     async addUserMessage(
         content: string,
         context: MessageContext[],
-    ): Promise<UserChatMessage> {
+    ): Promise<ChatMessage> {
         if (!this.current) {
             console.error('No conversation found');
             throw new Error('No conversation found');
@@ -147,34 +147,43 @@ export class ConversationManager {
     }
 
     async attachCommitToUserMessage(id: string, commit: GitCommit): Promise<void> {
-        const message = this.current?.messages.find((m) => m.id === id && m.role === ChatMessageRole.USER);
+        if (!this.current) {
+            console.error('No conversation found');
+            return;
+        }
+        const message = this.current.messages.find((m) => m.id === id && m.role === 'user');
         if (!message) {
             console.error('No message found with id', id);
             return;
         }
-        const userMessage = message as UserChatMessage;
         const newCheckpoints = [
-            ...userMessage.metadata.checkpoints,
+            ...(message.metadata?.checkpoints ?? []),
             {
                 type: MessageCheckpointType.GIT,
                 oid: commit.oid,
                 createdAt: new Date(),
             },
         ];
-        userMessage.metadata.checkpoints = newCheckpoints;
+        message.metadata = {
+            ...message.metadata,
+            createdAt: message.metadata?.createdAt ?? new Date(),
+            conversationId: message.metadata?.conversationId || this.current.conversation.id,
+            checkpoints: newCheckpoints,
+            context: message.metadata?.context ?? [],
+        };
         await api.chat.message.updateCheckpoints.mutate({
             messageId: message.id,
             checkpoints: newCheckpoints,
         });
-        await this.addOrReplaceMessage(userMessage);
+        await this.addOrReplaceMessage(message);
     }
 
     async addOrReplaceMessage(message: ChatMessage) {
         if (!this.current) {
             console.error('No conversation found');
             return;
         }
-        const index = this.current.messages.findIndex((m) => m.id === message.id || (m.metadata?.vercelId && m.metadata?.vercelId === message.metadata?.vercelId));
+        const index = this.current.messages.findIndex((m) => m.id === message.id);
         if (index === -1) {
             this.current.messages.push(message);
         } else {
@@ -209,7 +218,11 @@ export class ConversationManager {
     }
 
     async upsertMessageInStorage(message: ChatMessage) {
-        await api.chat.message.upsert.mutate({ message: toDbMessage(message) });
+        if (!this.current) {
+            console.error('No conversation found');
+            return;
+        }
+        await api.chat.message.upsert.mutate({ message: toDbMessage(message, this.current.conversation.id) });
     }
 
     clear() {