feat: reduce agent overhead (#2953)

Author: Kitenite

apps/web/client/package.json

@@ -28,8 +28,8 @@
         "preview:standalone": "bun run build:standalone && bun run start:standalone"
     },
     "dependencies": {
-        "@ai-sdk/provider-utils": "^3.0.5",
-        "@ai-sdk/react": "2.0.25",
+        "@ai-sdk/provider-utils": "^3.0.10",
+        "@ai-sdk/react": "2.0.60",
         "@codemirror/lang-css": "^6.2.0",
         "@codemirror/lang-html": "^6.4.7",
         "@codemirror/lang-javascript": "^6.2.3",

apps/web/client/public/onlook-preload-script.js

@@ -1,51 +1,3 @@
-import type { ToolCall } from '@ai-sdk/provider-utils';
-import { initModel } from '@onlook/ai';
-import { LLMProvider, OPENROUTER_MODELS } from '@onlook/models';
-import { generateObject, NoSuchToolError, type ToolSet } from 'ai';
-
-
-export const repairToolCall = async ({ toolCall, tools, error }: { toolCall: ToolCall<string, unknown>, tools: ToolSet, error: Error }) => {
-    if (NoSuchToolError.isInstance(error)) {
-        throw new Error(
-            `Tool "${toolCall.toolName}" not found. Available tools: ${Object.keys(tools).join(', ')}`,
-        );
-    }
-    const tool = tools[toolCall.toolName];
-
-    if (!tool?.inputSchema) {
-        throw new Error(`Tool "${toolCall.toolName}" has no input schema`);
-    }
-
-    console.warn(
-        `Invalid parameter for tool ${toolCall.toolName} with args ${JSON.stringify(toolCall.input)}, attempting to fix`,
-    );
-
-    const { model } = initModel({
-        provider: LLMProvider.OPENROUTER,
-        model: OPENROUTER_MODELS.OPEN_AI_GPT_5_NANO,
-    });
-
-    const { object: repairedArgs } = await generateObject({
-        model,
-        schema: tool.inputSchema,
-        prompt: [
-            `The model tried to call the tool "${toolCall.toolName}"` +
-            ` with the following arguments:`,
-            JSON.stringify(toolCall.input),
-            `The tool accepts the following schema:`,
-            JSON.stringify(tool?.inputSchema),
-            'Please fix the inputs. Return the fixed inputs as a JSON object, DO NOT include any other text.',
-        ].join('\n'),
-    });
-
-    return {
-        type: 'tool-call' as const,
-        toolCallId: toolCall.toolCallId,
-        toolName: toolCall.toolName,
-        input: JSON.stringify(repairedArgs),
-    };
-}
-
 export function errorHandler(error: unknown) {
     try {
         console.error('Error in chat', error);

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

@@ -1,12 +1,11 @@
 import { api } from '@/trpc/server';
 import { trackEvent } from '@/utils/analytics/server';
-import { AgentStreamer, BaseAgent, RootAgent, UserAgent } from '@onlook/ai';
+import { createRootAgentStream } from '@onlook/ai';
 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';
+import { checkMessageLimit, decrementUsage, errorHandler, getSupabaseUser, incrementUsage } from './helpers';
 
 export async function POST(req: NextRequest) {
     try {
@@ -52,24 +51,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 +71,31 @@ 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 stream = createRootAgentStream({
+            chatType,
+            conversationId,
+            projectId,
+            userId,
+            traceId,
+            messages,
+        });
+        return stream.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 +109,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/api/chat/route.ts

@@ -19,7 +19,7 @@ import { useMemo } from 'react';
 export const ChatContextWindow = ({ usage }: { usage: LanguageModelUsage }) => {
     const showCost = false;
     // Hardcoded for now, but should be dynamic based on the model used
-    const maxTokens = MODEL_MAX_TOKENS[OPENROUTER_MODELS.CLAUDE_4_SONNET];
+    const maxTokens = MODEL_MAX_TOKENS[OPENROUTER_MODELS.CLAUDE_4_5_SONNET];
     const usedTokens = useMemo(() => {
         if (!usage) return 0;
         const input = usage.inputTokens ?? 0;

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-input/chat-context.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/app/project/[id]/_hooks/use-chat/index.tsx

@@ -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/components/tools/tools.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');