feat: reduce agent overhead

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

@@ -1,8 +1,48 @@
 import type { ToolCall } from '@ai-sdk/provider-utils';
-import { initModel } from '@onlook/ai';
-import { LLMProvider, OPENROUTER_MODELS } from '@onlook/models';
+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';
 
+export function getModelFromType(chatType: ChatType) {
+    let model: ModelConfig;
+    switch (chatType) {
+        case ChatType.CREATE:
+        case ChatType.FIX:
+            model = initModel({
+                provider: LLMProvider.OPENROUTER,
+                model: OPENROUTER_MODELS.OPEN_AI_GPT_5,
+            });
+            break;
+        case ChatType.ASK:
+        case ChatType.EDIT:
+        default:
+            model = initModel({
+                provider: LLMProvider.OPENROUTER,
+                model: OPENROUTER_MODELS.CLAUDE_4_SONNET,
+            });
+            break;
+    }
+    return model;
+}
+
+export function getSystemPromptFromType(chatType: ChatType) {
+    let systemPrompt: string;
+
+    switch (chatType) {
+        case ChatType.CREATE:
+            systemPrompt = getCreatePageSystemPrompt();
+            break;
+        case ChatType.ASK:
+            systemPrompt = getAskModeSystemPrompt();
+            break;
+        case ChatType.EDIT:
+        default:
+            systemPrompt = getSystemPrompt();
+            break;
+    }
+    return systemPrompt;
+}
+
 
 export const repairToolCall = async ({ toolCall, tools, error }: { toolCall: ToolCall<string, unknown>, tools: ToolSet, error: Error }) => {
     if (NoSuchToolError.isInstance(error)) {
@@ -20,7 +60,7 @@ export const repairToolCall = async ({ toolCall, tools, error }: { toolCall: Too
         `Invalid parameter for tool ${toolCall.toolName} with args ${JSON.stringify(toolCall.input)}, attempting to fix`,
     );
 
-    const { model } = initModel({
+    const { model } = await initModel({
         provider: LLMProvider.OPENROUTER,
         model: OPENROUTER_MODELS.OPEN_AI_GPT_5_NANO,
     });

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

@@ -1,12 +1,12 @@
 import { api } from '@/trpc/server';
 import { trackEvent } from '@/utils/analytics/server';
-import { AgentStreamer, BaseAgent, RootAgent, UserAgent } from '@onlook/ai';
+import { convertToStreamMessages } from '@onlook/ai';
+import { createRootAgent } from '@onlook/ai/src/agents/classes/root';
 import { toDbMessage } from '@onlook/db';
-import { AgentType, ChatType } from '@onlook/models';
+import { ChatType, type ChatMessage, type ChatMetadata } from '@onlook/models';
 import { type NextRequest } from 'next/server';
 import { v4 as uuidv4 } from 'uuid';
 import { checkMessageLimit, decrementUsage, errorHandler, getSupabaseUser, incrementUsage, repairToolCall } from './helpers';
-import { z } from 'zod';
 
 export async function POST(req: NextRequest) {
     try {
@@ -52,24 +52,14 @@ export async function POST(req: NextRequest) {
     }
 }
 
-const streamResponseSchema = z.object({
-    agentType: z.enum(AgentType).optional().default(AgentType.ROOT),
-    messages: z.array(z.any()),
-    chatType: z.enum(ChatType).optional(),
-    conversationId: z.string(),
-    projectId: z.string(),
-}).refine((data) => {
-    // Only allow chatType if agentType is ROOT
-    if (data.chatType !== undefined && data.agentType !== AgentType.ROOT) {
-        return false;
-    }
-    return true;
-}, { message: "chatType is only allowed if agentType is root" });
-
 export const streamResponse = async (req: NextRequest, userId: string) => {
     const body = await req.json();
-    const { agentType, messages, chatType, conversationId, projectId } = streamResponseSchema.parse(body);
-
+    const { messages, chatType, conversationId, projectId } = body as {
+        messages: ChatMessage[],
+        chatType: ChatType,
+        conversationId: string,
+        projectId: string,
+    };
     // Updating the usage record and rate limit is done here to avoid
     // abuse in the case where a single user sends many concurrent requests.
     // If the call below fails, the user will not be penalized.
@@ -82,53 +72,35 @@ export const streamResponse = async (req: NextRequest, userId: string) => {
         const lastUserMessage = messages.findLast((message) => message.role === 'user');
         const traceId = lastUserMessage?.id ?? uuidv4();
 
-        // Create RootAgent instance
-        let agent: BaseAgent;
-        if (agentType === AgentType.ROOT) {
-            if (chatType === ChatType.EDIT) {
-                usageRecord = await incrementUsage(req, traceId);
-            }
-
-            agent = new RootAgent(chatType!);
-        } else if (agentType === AgentType.USER) {
-            agent = new UserAgent();
-        } else {
-            // agent = new WeatherAgent();
-            throw new Error('Agent type not supported');
+        if (chatType === ChatType.EDIT) {
+            usageRecord = await incrementUsage(req, traceId);
         }
-        const streamer = new AgentStreamer(agent, conversationId);
-
-        return streamer.streamText(messages, {
-            streamTextConfig: {
-                experimental_telemetry: {
-                    isEnabled: true,
-                    metadata: {
+        const { agent, modelConfig } = createRootAgent({
+            chatType,
+            conversationId,
+            projectId,
+            userId,
+            traceId,
+            repairToolCall,
+        });
+        const result = agent.stream({
+            providerOptions: modelConfig.providerOptions,
+            messages: convertToStreamMessages(messages)
+        });
+        return result.toUIMessageStreamResponse<ChatMessage>(
+            {
+                originalMessages: messages,
+                generateMessageId: () => uuidv4(),
+                messageMetadata: ({ part }) => {
+                    return {
+                        createdAt: new Date(),
                         conversationId,
-                        projectId,
-                        userId,
-                        agentType: agentType ?? AgentType.ROOT,
-                        chatType: chatType ?? "null",
-                        tags: ['chat'],
-                        langfuseTraceId: traceId,
-                        sessionId: conversationId,
-                    },
-                },
-                experimental_repairToolCall: repairToolCall,
-                onError: async (error) => {
-                    console.error('Error in chat stream call', error);
-                    // if there was an error with the API, do not penalize the user
-                    await decrementUsage(req, usageRecord);
-
-                    // Ensure the stream stops on error by re-throwing
-                    if (error instanceof Error) {
-                        throw error;
-                    } else {
-                        const errorMessage = typeof error === 'string' ? error : JSON.stringify(error);
-                        throw new Error(errorMessage);
-                    }
+                        context: [],
+                        checkpoints: [],
+                        finishReason: part.type === 'finish-step' ? part.finishReason : undefined,
+                        usage: part.type === 'finish-step' ? part.usage : undefined,
+                    } satisfies ChatMetadata;
                 },
-            },
-            toUIMessageStreamResponseConfig: {
                 onFinish: async ({ messages: finalMessages }) => {
                     const messagesToStore = finalMessages
                         .filter(msg =>
@@ -142,8 +114,8 @@ export const streamResponse = async (req: NextRequest, userId: string) => {
                     });
                 },
                 onError: errorHandler,
-            },
-        });
+            }
+        );
     } catch (error) {
         console.error('Error in streamResponse setup', error);
         // If there was an error setting up the stream and we incremented usage, revert it

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

@@ -3,7 +3,7 @@
 import { useEditorEngine } from '@/components/store/editor';
 import { handleToolCall } from '@/components/tools';
 import { useChat as useAiChat } from '@ai-sdk/react';
-import { AgentType, ChatType, type ChatMessage, type MessageContext, type QueuedMessage } from '@onlook/models';
+import { ChatType, type ChatMessage, type MessageContext, type QueuedMessage } from '@onlook/models';
 import { jsonClone } from '@onlook/utility';
 import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls, type FinishReason } from 'ai';
 import { usePostHog } from 'posthog-js/react';
@@ -31,7 +31,6 @@ interface UseChatProps {
     projectId: string;
     initialMessages: ChatMessage[];
 }
-const agentType = AgentType.ROOT;
 
 export function useChat({ conversationId, projectId, initialMessages }: UseChatProps) {
     const editorEngine = useEditorEngine();
@@ -52,12 +51,11 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
                 body: {
                     conversationId,
                     projectId,
-                    agentType,
                 },
             }),
             onToolCall: async (toolCall) => {
                 setIsExecutingToolCall(true);
-                void handleToolCall(agentType, toolCall.toolCall, editorEngine, addToolResult).then(() => {
+                void handleToolCall(toolCall.toolCall, editorEngine, addToolResult).then(() => {
                     setIsExecutingToolCall(false);
                 });
             },
@@ -90,7 +88,6 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
                     chatType: type,
                     conversationId,
                     context: messageContext,
-                    agentType,
                 },
             });
             void editorEngine.chat.conversation.generateTitle(content);
@@ -165,7 +162,6 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
                 body: {
                     chatType,
                     conversationId,
-                    agentType,
                 },
             });
 
@@ -185,16 +181,16 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
 
     const processNextInQueue = useCallback(async () => {
         if (isProcessingQueue.current || isStreaming || queuedMessages.length === 0) return;
-        
+
         const nextMessage = queuedMessages[0];
         if (!nextMessage) return;
-        
+
         isProcessingQueue.current = true;
-        
+
         try {
             const refreshedContext = await editorEngine.chat.context.getRefreshedContext(nextMessage.context);
             await processMessage(nextMessage.content, nextMessage.type, refreshedContext);
-            
+
             // Remove only after successful processing
             setQueuedMessages(prev => prev.slice(1));
         } catch (error) {

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

@@ -141,13 +141,6 @@ export class ConversationManager {
         return api.chat.conversation.getAll.query({ projectId: id });
     }
 
-    async upsertConversationInStorage(conversation: Partial<ChatConversation>): Promise<ChatConversation> {
-        return await api.chat.conversation.upsert.mutate({
-            ...conversation,
-            projectId: this.editorEngine.projectId,
-        });
-    }
-
     async updateConversationInStorage(conversation: Partial<ChatConversation> & { id: string }) {
         await api.chat.conversation.update.mutate(conversation);
     }

apps/web/client/src/components/tools/tools.ts

@@ -1,34 +1,29 @@
 import type { EditorEngine } from '@/components/store/editor/engine';
 import type { ToolCall } from '@ai-sdk/provider-utils';
-import type { AbstractChat } from 'ai';
-import { getAvailableTools, type OnToolCallHandler } from '@onlook/ai';
+import { getToolClassesFromType } from '@onlook/ai';
 import { toast } from '@onlook/ui/sonner';
-import type { AgentType } from '@onlook/models';
 
-export async function handleToolCall(agentType: AgentType, toolCall: ToolCall<string, unknown>, editorEngine: EditorEngine, addToolResult: typeof AbstractChat.prototype.addToolResult) {
+export async function handleToolCall(toolCall: ToolCall<string, unknown>, editorEngine: EditorEngine, addToolResult: (toolResult: { tool: string, toolCallId: string, output: any }) => Promise<void>) {
     const toolName = toolCall.toolName;
     const currentChatMode = editorEngine.state.chatMode;
-    const availableTools = getAvailableTools(agentType, currentChatMode) as any[];
-    let output: any = null;
+    const availableTools = getToolClassesFromType(currentChatMode);
+    let output: unknown = null;
 
     try {
-        const tool = availableTools.find((tool: any) => tool.toolName === toolName);
+        const tool = availableTools.find(tool => tool.toolName === toolName);
         if (!tool) {
             toast.error(`Tool "${toolName}" not available in ask mode`, {
                 description: `Switch to build mode to use this tool.`,
                 duration: 2000,
             });
 
-            throw new Error(`Tool "${toolName}" is not available in ${currentChatMode} mode!!!!`);
+            throw new Error(`Tool "${toolName}" is not available in ${currentChatMode} mode`);
         }
         // Parse the input to the tool parameters. Throws if invalid.
         const validatedInput = tool.parameters.parse(toolCall.input);
         const toolInstance = new tool();
-        const getOnToolCall: OnToolCallHandler = (subAgentType, addSubAgentToolResult) => (toolCall) =>
-            void handleToolCall(subAgentType, toolCall.toolCall, editorEngine, addSubAgentToolResult);
-
         // Can force type with as any because we know the input is valid.
-        output = await toolInstance.handle(validatedInput as any, editorEngine, getOnToolCall);
+        output = await toolInstance.handle(validatedInput as any, editorEngine);
     } catch (error) {
         output = 'error handling tool call ' + error;
     } finally {

apps/web/client/src/server/api/routers/chat/conversation.ts

@@ -36,13 +36,7 @@ export const conversationRouter = createTRPCRouter({
     upsert: protectedProcedure
         .input(conversationInsertSchema)
         .mutation(async ({ ctx, input }) => {
-            const [conversation] = await ctx.db.insert(conversations).values(input).onConflictDoUpdate({
-                target: [conversations.id],
-                set: {
-                    ...input,
-                    updatedAt: new Date(),
-                },
-            }).returning();
+            const [conversation] = await ctx.db.insert(conversations).values(input).returning();
             if (!conversation) {
                 throw new Error('Conversation not created');
             }
@@ -51,11 +45,10 @@ export const conversationRouter = createTRPCRouter({
     update: protectedProcedure
         .input(conversationUpdateSchema)
         .mutation(async ({ ctx, input }) => {
-            const [conversation] = await ctx.db.update(conversations)
-                .set({
-                    ...input,
-                    updatedAt: new Date(),
-                })
+            const [conversation] = await ctx.db.update({
+                ...conversations,
+                updatedAt: new Date(),
+            }).set(input)
                 .where(eq(conversations.id, input.id)).returning();
             if (!conversation) {
                 throw new Error('Conversation not updated');
@@ -75,7 +68,7 @@ export const conversationRouter = createTRPCRouter({
             content: z.string(),
         }))
         .mutation(async ({ ctx, input }) => {
-            const { model, providerOptions, headers } = initModel({
+            const { model, providerOptions, headers } = await initModel({
                 provider: LLMProvider.OPENROUTER,
                 model: OPENROUTER_MODELS.CLAUDE_3_5_HAIKU,
             });

apps/web/client/src/server/api/routers/chat/suggestion.ts

@@ -18,7 +18,7 @@ export const suggestionsRouter = createTRPCRouter({
             })),
         }))
         .mutation(async ({ ctx, input }) => {
-            const { model, headers } = initModel({
+            const { model, headers } = await initModel({
                 provider: LLMProvider.OPENROUTER,
                 model: OPENROUTER_MODELS.OPEN_AI_GPT_5_NANO,
             });

apps/web/client/src/server/api/routers/project/project.ts

@@ -321,7 +321,7 @@ export const projectRouter = createTRPCRouter({
         }))
         .mutation(async ({ ctx, input }): Promise<string> => {
             try {
-                const { model, providerOptions, headers } = initModel({
+                const { model, providerOptions, headers } = await initModel({
                     provider: LLMProvider.OPENROUTER,
                     model: OPENROUTER_MODELS.OPEN_AI_GPT_5_NANO,
                 });