Limit agent response length and cost

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

@@ -3,6 +3,7 @@ import { useChat } from '../../../../_hooks/use-chat';
 import { ChatInput } from '../chat-input';
 import { ChatMessages } from '../chat-messages';
 import { ErrorSection } from '../error';
+import { StepLimitBanner } from '../step-limit-banner';
 
 interface ChatTabContentProps {
     conversationId: string;
@@ -15,7 +16,19 @@ export const ChatTabContent = ({
     projectId,
     initialMessages,
 }: ChatTabContentProps) => {
-    const { isStreaming, sendMessage, editMessage, messages, error, stop, queuedMessages, removeFromQueue } = useChat({
+    const {
+        isStreaming,
+        sendMessage,
+        editMessage,
+        messages,
+        error,
+        stop,
+        queuedMessages,
+        removeFromQueue,
+        hitStepLimit,
+        continueAfterStepLimit,
+        dismissStepLimit,
+    } = useChat({
         conversationId,
         projectId,
         initialMessages,
@@ -30,6 +43,11 @@ export const ChatTabContent = ({
                 onEditMessage={editMessage}
             />
             <ErrorSection isStreaming={isStreaming} onSendMessage={sendMessage} />
+            <StepLimitBanner
+                show={hitStepLimit}
+                onContinue={() => void continueAfterStepLimit()}
+                onDismiss={dismissStepLimit}
+            />
             <ChatInput
                 messages={messages}
                 isStreaming={isStreaming}

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/step-limit-banner.tsx

@@ -0,0 +1,57 @@
+import { Button } from '@onlook/ui/button';
+import { Icons } from '@onlook/ui/icons';
+import { AnimatePresence, motion } from 'motion/react';
+
+interface StepLimitBannerProps {
+    show: boolean;
+    onContinue: () => void;
+    onDismiss: () => void;
+}
+
+export const StepLimitBanner = ({ show, onContinue, onDismiss }: StepLimitBannerProps) => {
+    return (
+        <AnimatePresence>
+            {show && (
+                <motion.div
+                    initial={{ opacity: 0, y: 10 }}
+                    animate={{ opacity: 1, y: 0 }}
+                    exit={{ opacity: 0, y: 10 }}
+                    transition={{ duration: 0.2 }}
+                    className="mx-2 mb-2"
+                >
+                    <div className="flex flex-col gap-2 p-3 rounded-lg border border-blue-200 dark:border-blue-500/30 bg-blue-50 dark:bg-blue-950/50">
+                        <div className="flex items-start gap-2">
+                            <Icons.InfoCircled className="h-4 w-4 mt-0.5 text-blue-600 dark:text-blue-400 shrink-0" />
+                            <div className="flex-1 min-w-0">
+                                <p className="text-sm font-medium text-blue-800 dark:text-blue-200">
+                                    Task paused
+                                </p>
+                                <p className="text-xs text-blue-600 dark:text-blue-300 mt-0.5">
+                                    The AI has completed several steps. Would you like to continue?
+                                </p>
+                            </div>
+                        </div>
+                        <div className="flex items-center gap-2 justify-end">
+                            <Button
+                                variant="ghost"
+                                size="sm"
+                                onClick={onDismiss}
+                                className="h-7 px-2 text-blue-600 dark:text-blue-400 hover:text-blue-800 hover:bg-blue-100 dark:hover:text-blue-200 dark:hover:bg-blue-900"
+                            >
+                                Stop here
+                            </Button>
+                            <Button
+                                size="sm"
+                                onClick={onContinue}
+                                className="h-7 px-3 bg-blue-600 hover:bg-blue-700 text-white"
+                            >
+                                <Icons.Play className="h-3 w-3 mr-1" />
+                                Continue
+                            </Button>
+                        </div>
+                    </div>
+                </motion.div>
+            )}
+        </AnimatePresence>
+    );
+};

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

@@ -40,13 +40,29 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
     const [finishReason, setFinishReason] = useState<FinishReason | null>(null);
     const [isExecutingToolCall, setIsExecutingToolCall] = useState(false);
     const [queuedMessages, setQueuedMessages] = useState<QueuedMessage[]>([]);
+    const [hitStepLimit, setHitStepLimit] = useState(false);
+    const [toolCallCount, setToolCallCount] = useState(0);
     const isProcessingQueue = useRef(false);
 
+    // Max tool calls before pausing and asking user to continue
+    // TODO: Change back to 10 after testing
+    const MAX_TOOL_CALLS = 2;
+
+    // Track tool call count in a ref to avoid stale closures
+    const toolCallCountRef = useRef(toolCallCount);
+    toolCallCountRef.current = toolCallCount;
+
     const { addToolResult, messages, error, stop, setMessages, regenerate, status } =
         useAiChat<ChatMessage>({
             id: 'user-chat',
             messages: initialMessages,
-            sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
+            // Only auto-send if we haven't hit the tool call limit
+            sendAutomaticallyWhen: (messages) => {
+                if (toolCallCountRef.current >= MAX_TOOL_CALLS) {
+                    return false;
+                }
+                return lastAssistantMessageIsCompleteWithToolCalls(messages);
+            },
             transport: new DefaultChatTransport({
                 api: '/api/chat',
                 body: {
@@ -56,13 +72,18 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
             }),
             onToolCall: async (toolCall) => {
                 setIsExecutingToolCall(true);
+                setToolCallCount(prev => prev + 1);
                 void handleToolCall(toolCall.toolCall, editorEngine, addToolResult).then(() => {
                     setIsExecutingToolCall(false);
                 });
             },
             onFinish: ({ message }) => {
                 const finishReason = message.metadata?.finishReason;
                 setFinishReason(finishReason ?? null);
+                // Check if we've hit the tool call limit
+                if (toolCallCountRef.current >= MAX_TOOL_CALLS) {
+                    setHitStepLimit(true);
+                }
             },
         });
 
@@ -79,7 +100,11 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
     }, [messages]);
 
     const processMessage = useCallback(
-        async (content: string, type: ChatType, context?: MessageContext[]) => {
+        async (content: string, type: ChatType, context?: MessageContext[], resetToolCount = true) => {
+            // Reset tool call count for new user messages
+            if (resetToolCount) {
+                setToolCallCount(0);
+            }
             const messageContext = context || await editorEngine.chat.context.getContextByChatType(type);
             const newMessage = getUserChatMessageFromString(content, messageContext, conversationId);
             setMessages(jsonClone([...messagesRef.current, newMessage]));
@@ -220,6 +245,18 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
         [processMessageEdit, posthog, isStreaming, stop, editorEngine.chat.context],
     );
 
+    // Continue after hitting the step limit
+    const continueAfterStepLimit = useCallback(() => {
+        setHitStepLimit(false);
+        posthog.capture('user_continue_after_step_limit');
+        return sendMessage('Continue where you left off.', ChatType.EDIT);
+    }, [sendMessage, posthog]);
+
+    // Dismiss the step limit banner without continuing
+    const dismissStepLimit = useCallback(() => {
+        setHitStepLimit(false);
+    }, []);
+
     useEffect(() => {
         // Actions to handle when the chat is finished
         if (finishReason && finishReason !== 'tool-calls') {
@@ -321,5 +358,8 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP
         isStreaming,
         queuedMessages,
         removeFromQueue,
+        hitStepLimit,
+        continueAfterStepLimit,
+        dismissStepLimit,
     };
 }

bun.lock

@@ -1,5 +1,6 @@
 {
   "lockfileVersion": 1,
+  "configVersion": 0,
   "workspaces": {
     "": {
       "name": "@onlook/repo",

packages/ai/src/agents/root.ts

@@ -3,6 +3,9 @@ import { ChatType, LLMProvider, OPENROUTER_MODELS, type ChatMessage, type ModelC
 import { NoSuchToolError, generateObject, smoothStream, stepCountIs, streamText, type ToolSet } from 'ai';
 import { convertToStreamMessages, getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt, getToolSetFromType, initModel } from '../index';
 
+// Max steps per single request (server-side safety limit)
+const MAX_AGENT_STEPS = 20;
+
 export const createRootAgentStream = ({
     chatType,
     conversationId,
@@ -28,7 +31,7 @@ export const createRootAgentStream = ({
         system: systemPrompt,
         tools: toolSet,
         headers: modelConfig.headers,
-        stopWhen: stepCountIs(20),
+        stopWhen: stepCountIs(MAX_AGENT_STEPS),
         experimental_repairToolCall: repairToolCall,
         experimental_transform: smoothStream(),
         experimental_telemetry: {