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 fd2d431dc1..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 @@ -40,13 +40,29 @@ 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 [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: { @@ -56,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); }); @@ -63,6 +80,10 @@ export function useChat({ conversationId, projectId, initialMessages }: UseChatP 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, }; } 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/src/agents/root.ts b/packages/ai/src/agents/root.ts index 134b05e52a..b15153fe10 100644 --- a/packages/ai/src/agents/root.ts +++ b/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: {