From 1bb1170cdaa37d718416118e3fa8780fc2828ebd Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 1 Dec 2025 03:14:23 +0000 Subject: [PATCH 1/4] Checkpoint before follow-up message Co-authored-by: daniel --- apps/web/client/messages/en.json | 6 ++++++ .../client/src/app/project/[id]/_hooks/use-chat/index.tsx | 6 ++++++ packages/ai/package.json | 1 + packages/ai/src/agents/root.ts | 3 ++- packages/constants/src/editor.ts | 5 +++++ 5 files changed, 20 insertions(+), 1 deletion(-) diff --git a/apps/web/client/messages/en.json b/apps/web/client/messages/en.json index 5907b83fbe..748ea00cc0 100644 --- a/apps/web/client/messages/en.json +++ b/apps/web/client/messages/en.json @@ -260,6 +260,12 @@ }, "openInCode": { "button": "Open in Code" + }, + "stepLimit": { + "title": "Task paused", + "description": "The AI has completed {stepCount} steps. Would you like to continue?", + "continue": "Continue", + "stop": "Stop here" } }, "styles": { diff --git a/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx b/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx index fd2d431dc1..b3fe50326a 100644 --- a/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx +++ b/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx @@ -4,6 +4,7 @@ import { useEditorEngine } from '@/components/store/editor'; import { handleToolCall } from '@/components/tools'; import { api } from '@/trpc/client'; import { useChat as useAiChat } from '@ai-sdk/react'; +import { MAX_AGENT_STEPS } from '@onlook/constants'; import { ChatType, type ChatMessage, type GitMessageCheckpoint, type MessageContext, type QueuedMessage } from '@onlook/models'; import { jsonClone } from '@onlook/utility'; import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls, type FinishReason } from 'ai'; @@ -40,6 +41,7 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP const [finishReason, setFinishReason] = useState(null); const [isExecutingToolCall, setIsExecutingToolCall] = useState(false); const [queuedMessages, setQueuedMessages] = useState([]); + const [hitStepLimit, setHitStepLimit] = useState(false); const isProcessingQueue = useRef(false); const { addToolResult, messages, error, stop, setMessages, regenerate, status } = @@ -63,6 +65,10 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP onFinish: ({ message }) => { const finishReason = message.metadata?.finishReason; setFinishReason(finishReason ?? null); + // Detect when the step limit is reached + if (finishReason === 'step-limit') { + setHitStepLimit(true); + } }, }); diff --git a/packages/ai/package.json b/packages/ai/package.json index c81655cff8..ce729ab970 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -35,6 +35,7 @@ "dependencies": { "@mendable/firecrawl-js": "^1.29.1", "@openrouter/ai-sdk-provider": "^1.2.0", + "@onlook/constants": "*", "ai": "5.0.60", "exa-js": "^1.8.26", "fg": "^0.0.3", diff --git a/packages/ai/src/agents/root.ts b/packages/ai/src/agents/root.ts index 134b05e52a..1efdb19233 100644 --- a/packages/ai/src/agents/root.ts +++ b/packages/ai/src/agents/root.ts @@ -1,4 +1,5 @@ import type { ToolCall } from '@ai-sdk/provider-utils'; +import { MAX_AGENT_STEPS } from '@onlook/constants'; import { ChatType, LLMProvider, OPENROUTER_MODELS, type ChatMessage, type ModelConfig } from '@onlook/models'; import { NoSuchToolError, generateObject, smoothStream, stepCountIs, streamText, type ToolSet } from 'ai'; import { convertToStreamMessages, getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt, getToolSetFromType, initModel } from '../index'; @@ -28,7 +29,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: { diff --git a/packages/constants/src/editor.ts b/packages/constants/src/editor.ts index dfcaca60fc..df4fb78b40 100644 --- a/packages/constants/src/editor.ts +++ b/packages/constants/src/editor.ts @@ -68,3 +68,8 @@ export const DefaultSettings = { }; export const DEFAULT_COLOR_NAME = 'DEFAULT'; + +// AI Agent limits +// The maximum number of steps (tool calls) the AI agent can take before pausing +// and asking the user if they want to continue. This prevents runaway costs. +export const MAX_AGENT_STEPS = 10; From 0944b84fa05460383affc4a39485a71d91f24fe3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 1 Dec 2025 03:24:26 +0000 Subject: [PATCH 2/4] feat: Add AI step limit banner and logic Co-authored-by: daniel --- apps/web/client/messages/en.json | 6 -- .../chat-tab/chat-tab-content/index.tsx | 20 ++++++- .../chat-tab/step-limit-banner.tsx | 57 +++++++++++++++++++ .../project/[id]/_hooks/use-chat/index.tsx | 20 ++++++- bun.lock | 1 + packages/ai/package.json | 1 - packages/ai/src/agents/root.ts | 4 +- packages/constants/src/editor.ts | 5 -- 8 files changed, 97 insertions(+), 17 deletions(-) create mode 100644 apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/step-limit-banner.tsx diff --git a/apps/web/client/messages/en.json b/apps/web/client/messages/en.json index 748ea00cc0..5907b83fbe 100644 --- a/apps/web/client/messages/en.json +++ b/apps/web/client/messages/en.json @@ -260,12 +260,6 @@ }, "openInCode": { "button": "Open in Code" - }, - "stepLimit": { - "title": "Task paused", - "description": "The AI has completed {stepCount} steps. Would you like to continue?", - "continue": "Continue", - "stop": "Stop here" } }, "styles": { diff --git a/apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-tab-content/index.tsx b/apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-tab-content/index.tsx index 6080dc8e9f..8efcd6844c 100644 --- a/apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-tab-content/index.tsx +++ b/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} /> + void continueAfterStepLimit()} + onDismiss={dismissStepLimit} + /> void; + onDismiss: () => void; +} + +export const StepLimitBanner = ({ show, onContinue, onDismiss }: StepLimitBannerProps) => { + return ( + + {show && ( + +
+
+ +
+

+ Task paused +

+

+ The AI has completed several steps. Would you like to continue? +

+
+
+
+ + +
+
+
+ )} +
+ ); +}; diff --git a/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx b/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx index b3fe50326a..e5daeab035 100644 --- a/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx +++ b/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx @@ -4,7 +4,6 @@ import { useEditorEngine } from '@/components/store/editor'; import { handleToolCall } from '@/components/tools'; import { api } from '@/trpc/client'; import { useChat as useAiChat } from '@ai-sdk/react'; -import { MAX_AGENT_STEPS } from '@onlook/constants'; import { ChatType, type ChatMessage, type GitMessageCheckpoint, type MessageContext, type QueuedMessage } from '@onlook/models'; import { jsonClone } from '@onlook/utility'; import { DefaultChatTransport, lastAssistantMessageIsCompleteWithToolCalls, type FinishReason } from 'ai'; @@ -65,8 +64,8 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP onFinish: ({ message }) => { const finishReason = message.metadata?.finishReason; setFinishReason(finishReason ?? null); - // Detect when the step limit is reached - if (finishReason === 'step-limit') { + // Detect when the step limit is reached (stepCountIs returns 'step-limit') + if ((finishReason as string) === 'step-limit') { setHitStepLimit(true); } }, @@ -226,6 +225,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') { @@ -327,5 +338,8 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP isStreaming, queuedMessages, removeFromQueue, + hitStepLimit, + continueAfterStepLimit, + dismissStepLimit, }; } diff --git a/bun.lock b/bun.lock index 02aa9c097b..fcbb451473 100644 --- a/bun.lock +++ b/bun.lock @@ -1,5 +1,6 @@ { "lockfileVersion": 1, + "configVersion": 0, "workspaces": { "": { "name": "@onlook/repo", diff --git a/packages/ai/package.json b/packages/ai/package.json index ce729ab970..c81655cff8 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -35,7 +35,6 @@ "dependencies": { "@mendable/firecrawl-js": "^1.29.1", "@openrouter/ai-sdk-provider": "^1.2.0", - "@onlook/constants": "*", "ai": "5.0.60", "exa-js": "^1.8.26", "fg": "^0.0.3", diff --git a/packages/ai/src/agents/root.ts b/packages/ai/src/agents/root.ts index 1efdb19233..6064336c17 100644 --- a/packages/ai/src/agents/root.ts +++ b/packages/ai/src/agents/root.ts @@ -1,9 +1,11 @@ import type { ToolCall } from '@ai-sdk/provider-utils'; -import { MAX_AGENT_STEPS } from '@onlook/constants'; import { ChatType, LLMProvider, OPENROUTER_MODELS, type ChatMessage, type ModelConfig } from '@onlook/models'; import { NoSuchToolError, generateObject, smoothStream, stepCountIs, streamText, type ToolSet } from 'ai'; import { convertToStreamMessages, getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt, getToolSetFromType, initModel } from '../index'; +// Max steps before pausing and asking user to continue (prevents runaway costs) +const MAX_AGENT_STEPS = 10; + export const createRootAgentStream = ({ chatType, conversationId, diff --git a/packages/constants/src/editor.ts b/packages/constants/src/editor.ts index df4fb78b40..dfcaca60fc 100644 --- a/packages/constants/src/editor.ts +++ b/packages/constants/src/editor.ts @@ -68,8 +68,3 @@ export const DefaultSettings = { }; export const DEFAULT_COLOR_NAME = 'DEFAULT'; - -// AI Agent limits -// The maximum number of steps (tool calls) the AI agent can take before pausing -// and asking the user if they want to continue. This prevents runaway costs. -export const MAX_AGENT_STEPS = 10; From 0684318231d5980f28f30615e7c55064a7f353c2 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 1 Dec 2025 18:06:58 +0000 Subject: [PATCH 3/4] feat: Reduce MAX_AGENT_STEPS to 2 for testing Co-authored-by: daniel --- packages/ai/src/agents/root.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/ai/src/agents/root.ts b/packages/ai/src/agents/root.ts index 6064336c17..a2e0a20da6 100644 --- a/packages/ai/src/agents/root.ts +++ b/packages/ai/src/agents/root.ts @@ -4,7 +4,8 @@ import { NoSuchToolError, generateObject, smoothStream, stepCountIs, streamText, import { convertToStreamMessages, getAskModeSystemPrompt, getCreatePageSystemPrompt, getSystemPrompt, getToolSetFromType, initModel } from '../index'; // Max steps before pausing and asking user to continue (prevents runaway costs) -const MAX_AGENT_STEPS = 10; +// TODO: Change back to 10 after testing +const MAX_AGENT_STEPS = 2; export const createRootAgentStream = ({ chatType, From 69dfe032b23f0f159aaa3ec21119f01472851273 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 1 Dec 2025 19:59:05 +0000 Subject: [PATCH 4/4] feat: Limit tool calls and reset count on new messages Co-authored-by: daniel --- .../project/[id]/_hooks/use-chat/index.tsx | 28 ++++++++++++++++--- packages/ai/src/agents/root.ts | 5 ++-- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx b/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx index e5daeab035..a2974553a8 100644 --- a/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx +++ b/apps/web/client/src/app/project/[id]/_hooks/use-chat/index.tsx @@ -41,13 +41,28 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP const [isExecutingToolCall, setIsExecutingToolCall] = useState(false); const [queuedMessages, setQueuedMessages] = useState([]); 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({ 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: { @@ -57,6 +72,7 @@ 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); }); @@ -64,8 +80,8 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP onFinish: ({ message }) => { const finishReason = message.metadata?.finishReason; setFinishReason(finishReason ?? null); - // Detect when the step limit is reached (stepCountIs returns 'step-limit') - if ((finishReason as string) === 'step-limit') { + // Check if we've hit the tool call limit + if (toolCallCountRef.current >= MAX_TOOL_CALLS) { setHitStepLimit(true); } }, @@ -84,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])); diff --git a/packages/ai/src/agents/root.ts b/packages/ai/src/agents/root.ts index a2e0a20da6..b15153fe10 100644 --- a/packages/ai/src/agents/root.ts +++ b/packages/ai/src/agents/root.ts @@ -3,9 +3,8 @@ 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 before pausing and asking user to continue (prevents runaway costs) -// TODO: Change back to 10 after testing -const MAX_AGENT_STEPS = 2; +// Max steps per single request (server-side safety limit) +const MAX_AGENT_STEPS = 20; export const createRootAgentStream = ({ chatType,