feat: mastra no storage

apps/web/client/.gitignore

@@ -45,5 +45,8 @@ yarn-error.log*
 # idea files
 .idea
 
+# mastra
+.mastra/
+
 # Ignore preload script changes unless the script should be updated. Uncomment to update.
-/public/onlook-preload-script.js
+/public/onlook-preload-script.js

apps/web/client/next.config.ts

@@ -8,11 +8,12 @@ import path from 'node:path';
 import './src/env';
 
 const nextConfig: NextConfig = {
-    devIndicators: false,
-    // TODO: Remove this once we have a proper ESLint and TypeScript config
+    devIndicators: {
+        buildActivity: false,
+    },
     eslint: {
         ignoreDuringBuilds: true,
-    }
+    },
 };
 
 if (process.env.NODE_ENV === 'development') {

apps/web/client/package.json

@@ -34,6 +34,9 @@
         "@codemirror/lang-javascript": "^6.2.3",
         "@codemirror/lang-json": "^6.0.1",
         "@codemirror/lang-markdown": "^6.1.7",
+        "@mastra/core": "^0.13.1",
+        "@mastra/memory": "^0.12.1",
+        "@mastra/pg": "^0.13.3",
         "@onlook/constants": "*",
         "@onlook/code-provider": "*",
         "@onlook/db": "*",
@@ -100,7 +103,7 @@
         "use-resize-observer": "^9.1.0",
         "uuid": "^11.1.0",
         "webfontloader": "^1.6.28",
-        "zod": "^3.24.3"
+        "zod": "^3"
     },
     "devDependencies": {
         "@eslint/eslintrc": "^3.3.1",

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

@@ -0,0 +1,125 @@
+import { createClient as createTRPCClient } from '@/trpc/request-server';
+import { createClient as createSupabaseClient } from '@/utils/supabase/request-server';
+import { initModel } from '@onlook/ai';
+import { LLMProvider, OPENROUTER_MODELS, UsageType, type Usage } from '@onlook/models';
+import { generateObject, NoSuchToolError, type ToolCall, type ToolSet } from 'ai';
+import { type NextRequest } from 'next/server';
+
+export const checkMessageLimit = async (req: NextRequest): Promise<{
+    exceeded: boolean;
+    usage: Usage;
+}> => {
+    const { api } = await createTRPCClient(req);
+    const usage = await api.usage.get();
+
+    const dailyUsage = usage.daily;
+    const dailyExceeded = dailyUsage.usageCount >= dailyUsage.limitCount;
+    if (dailyExceeded) {
+        return {
+            exceeded: true,
+            usage: dailyUsage,
+        };
+    }
+
+    const monthlyUsage = usage.monthly;
+    const monthlyExceeded = monthlyUsage.usageCount >= monthlyUsage.limitCount;
+    if (monthlyExceeded) {
+        return {
+            exceeded: true,
+            usage: monthlyUsage,
+        };
+    }
+
+    return {
+        exceeded: false,
+        usage: monthlyUsage,
+    };
+}
+
+export const getSupabaseUser = async (request: NextRequest) => {
+    const supabase = await createSupabaseClient(request);
+    const { data: { user } } = await supabase.auth.getUser();
+    return user;
+}
+
+export const repairToolCall = async ({ toolCall, tools, error }: { toolCall: ToolCall<string, any>, 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 as keyof typeof tools];
+
+    console.warn(
+        `Invalid parameter for tool ${toolCall.toolName} with args ${JSON.stringify(toolCall.args)}, attempting to fix`,
+    );
+
+    const { model } = await initModel({
+        provider: LLMProvider.OPENROUTER,
+        model: OPENROUTER_MODELS.CLAUDE_4_SONNET,
+    });
+
+    const { object: repairedArgs } = await generateObject({
+        model,
+        schema: tool?.parameters,
+        prompt: [
+            `The model tried to call the tool "${toolCall.toolName}"` +
+            ` with the following arguments:`,
+            JSON.stringify(toolCall.args),
+            `The tool accepts the following schema:`,
+            JSON.stringify(tool?.parameters),
+            'Please fix the arguments.',
+        ].join('\n'),
+    });
+
+    return {
+        ...toolCall,
+        args: JSON.stringify(repairedArgs),
+        toolCallType: 'function' as const
+    };
+}
+
+export const incrementUsage = async (req: NextRequest): Promise<{
+    usageRecordId: string | undefined,
+    rateLimitId: string | undefined,
+} | null> => {
+    try {
+        const user = await getSupabaseUser(req);
+        if (!user) {
+            throw new Error('User not found');
+        }
+        const { api } = await createTRPCClient(req);
+        const incrementRes = await api.usage.increment({
+            type: UsageType.MESSAGE,
+        });
+        return {
+            usageRecordId: incrementRes?.usageRecordId,
+            rateLimitId: incrementRes?.rateLimitId,
+        };
+    } catch (error) {
+        console.error('Error in chat usage increment', error);
+    }
+    return null;
+}
+
+export const decrementUsage = async (
+    req: NextRequest,
+    usageRecord: {
+        usageRecordId: string | undefined,
+        rateLimitId: string | undefined,
+    } | null
+): Promise<void> => {
+    try {
+        if (!usageRecord) {
+            return;
+        }
+        const { usageRecordId, rateLimitId } = usageRecord;
+        if (!usageRecordId || !rateLimitId) {
+            return;
+        }
+        const { api } = await createTRPCClient(req);
+        await api.usage.revertIncrement({ usageRecordId, rateLimitId });
+    } catch (error) {
+        console.error('Error in chat usage decrement', error);
+    }
+}

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

@@ -1,21 +1,11 @@
-import { env } from '@/env';
-import { createClient as createTRPCClient } from '@/trpc/request-server';
+import { mastra } from '@/mastra';
+import { CHAT_TYPE_KEY, ONLOOK_AGENT_KEY, type OnlookAgentRuntimeContext } from '@/mastra/agents';
 import { trackEvent } from '@/utils/analytics/server';
-import { createClient as createSupabaseClient } from '@/utils/supabase/request-server';
-import { ASK_TOOL_SET, BUILD_TOOL_SET, getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt, initModel } from '@onlook/ai';
-import { ChatType, type InitialModelPayload, LLMProvider, OPENROUTER_MODELS, type Usage, UsageType } from '@onlook/models';
-import { generateObject, NoSuchToolError, streamText } from 'ai';
+import { RuntimeContext } from '@mastra/core/runtime-context';
+import { convertToStreamMessages } from '@onlook/ai';
+import { ChatType, type ChatMessage } from '@onlook/models';
 import { type NextRequest } from 'next/server';
-
-const isProd = env.NODE_ENV === 'production';
-
-const MainModelConfig: InitialModelPayload = isProd ? {
-    provider: LLMProvider.OPENROUTER,
-    model: OPENROUTER_MODELS.OPEN_AI_GPT_5,
-} : {
-    provider: LLMProvider.OPENROUTER,
-    model: OPENROUTER_MODELS.OPEN_AI_GPT_5,
-};
+import { checkMessageLimit, decrementUsage, getSupabaseUser, incrementUsage, repairToolCall } from './helpers';
 
 export async function POST(req: NextRequest) {
     try {
@@ -29,7 +19,6 @@ export async function POST(req: NextRequest) {
                 headers: { 'Content-Type': 'application/json' }
             });
         }
-
         const usageCheckResult = await checkMessageLimit(req);
         if (usageCheckResult.exceeded) {
             trackEvent({
@@ -63,130 +52,37 @@ export async function POST(req: NextRequest) {
     }
 }
 
-export const checkMessageLimit = async (req: NextRequest): Promise<{
-    exceeded: boolean;
-    usage: Usage;
-}> => {
-    const { api } = await createTRPCClient(req);
-    const usage = await api.usage.get();
-
-    const dailyUsage = usage.daily;
-    const dailyExceeded = dailyUsage.usageCount >= dailyUsage.limitCount;
-    if (dailyExceeded) {
-        return {
-            exceeded: true,
-            usage: dailyUsage,
-        };
-    }
-
-    const monthlyUsage = usage.monthly;
-    const monthlyExceeded = monthlyUsage.usageCount >= monthlyUsage.limitCount;
-    if (monthlyExceeded) {
-        return {
-            exceeded: true,
-            usage: monthlyUsage,
-        };
-    }
-
-    return {
-        exceeded: false,
-        usage: monthlyUsage,
-    };
-}
-
-export const getSupabaseUser = async (request: NextRequest) => {
-    const supabase = await createSupabaseClient(request);
-    const { data: { user } } = await supabase.auth.getUser();
-    return user;
-}
-
 export const streamResponse = async (req: NextRequest) => {
-    const { messages, maxSteps, chatType } = await req.json();
-    const { model, providerOptions, headers, maxTokens } = await initModel(MainModelConfig);
+    const { messages, maxSteps, chatType }: { messages: ChatMessage[], maxSteps: number, chatType: ChatType } = await req.json();
+    const agent = mastra.getAgent(ONLOOK_AGENT_KEY);
+    const runtimeContext = new RuntimeContext<OnlookAgentRuntimeContext>()
+    runtimeContext.set(CHAT_TYPE_KEY, chatType);
 
     // 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.
-    let usageRecordId: string | undefined;
-    let rateLimitId: string | undefined;
+    let usageRecord: {
+        usageRecordId: string | undefined;
+        rateLimitId: string | undefined;
+    } | null;
     if (chatType === ChatType.EDIT) {
-        const user = await getSupabaseUser(req);
-        if (!user) {
-            throw new Error('User not found');
-        }
-        const { api } = await createTRPCClient(req);
-        const incrementRes = await api.usage.increment({
-            type: UsageType.MESSAGE,
-        });
-        usageRecordId = incrementRes?.usageRecordId;
-        rateLimitId = incrementRes?.rateLimitId;
-    }
-
-    let systemPrompt: string;
-    switch (chatType) {
-        case ChatType.CREATE:
-            systemPrompt = getCreatePageSystemPrompt();
-            break;
-        case ChatType.ASK:
-            systemPrompt = getAskModeSystemPrompt();
-            break;
-        case ChatType.EDIT:
-        default:
-            systemPrompt = getSystemPrompt();
-            break;
+        usageRecord = await incrementUsage(req);
     }
-    const toolSet = chatType === ChatType.ASK ? ASK_TOOL_SET : BUILD_TOOL_SET;
-    const result = streamText({
-        model,
-        headers,
-        messages: [
-            {
-                role: 'system',
-                content: systemPrompt,
-                providerOptions,
-            },
-            ...messages,
-        ],
+    const result = await agent.stream(convertToStreamMessages(messages), {
+        headers: {
+            'HTTP-Referer': 'https://onlook.com',
+            'X-Title': 'Onlook',
+        },
         maxSteps,
-        tools: toolSet,
+        runtimeContext,
         toolCallStreaming: true,
-        experimental_repairToolCall: async ({ toolCall, tools, parameterSchema, error }) => {
-            if (NoSuchToolError.isInstance(error)) {
-                throw new Error(
-                    `Tool "${toolCall.toolName}" not found. Available tools: ${Object.keys(tools).join(', ')}`,
-                );
-            }
-            const tool = tools[toolCall.toolName as keyof typeof tools];
-
-            console.warn(
-                `Invalid parameter for tool ${toolCall.toolName} with args ${JSON.stringify(toolCall.args)}, attempting to fix`,
-            );
-
-            const { object: repairedArgs } = await generateObject({
-                model,
-                schema: tool?.parameters,
-                prompt: [
-                    `The model tried to call the tool "${toolCall.toolName}"` +
-                    ` with the following arguments:`,
-                    JSON.stringify(toolCall.args),
-                    `The tool accepts the following schema:`,
-                    JSON.stringify(parameterSchema(toolCall)),
-                    'Please fix the arguments.',
-                ].join('\n'),
-            });
-
-            return { ...toolCall, args: JSON.stringify(repairedArgs) };
-        },
+        experimental_repairToolCall: repairToolCall,
         onError: async (error) => {
             console.error('Error in chat', error);
             // if there was an error with the API, do not penalize the user
-            if (usageRecordId && rateLimitId) {
-                await createTRPCClient(req)
-                    .then(({ api }) => api.usage.revertIncrement({ usageRecordId, rateLimitId }))
-                    .catch(error => console.error('Error in chat usage decrement', error));
-            }
-        },
-    });
+            await decrementUsage(req, usageRecord);
+        }
+    })
 
     return result.toDataStreamResponse(
         {
@@ -196,17 +92,22 @@ export const streamResponse = async (req: NextRequest) => {
 }
 
 export function errorHandler(error: unknown) {
-    if (error == null) {
-        return 'unknown error';
-    }
+    try {
+        console.error('Error in chat', error);
+        if (!error) {
+            return 'unknown error';
+        }
 
-    if (typeof error === 'string') {
-        return error;
-    }
+        if (typeof error === 'string') {
+            return error;
+        }
 
-    if (error instanceof Error) {
-        return error.message;
+        if (error instanceof Error) {
+            return error.message;
+        }
+        return JSON.stringify(error);
+    } catch (error) {
+        console.error('Error in errorHandler', error);
+        return 'unknown error';
     }
-
-    return JSON.stringify(error);
 }