fix: add active branch context when page selected (#2911)

Author: Kitenite

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

@@ -81,7 +81,6 @@ export const streamResponse = async (req: NextRequest, userId: string) => {
         const { model, providerOptions, headers } = modelConfig;
         const systemPrompt = getSystemPromptFromType(chatType);
         const tools = getToolSetFromType(chatType);
-
         const result = streamText({
             model,
             headers,

apps/web/client/src/components/store/editor/chat/context.ts

@@ -12,6 +12,7 @@ import {
 import { assertNever, type ParsedError } from '@onlook/utility';
 import { makeAutoObservable, reaction } from 'mobx';
 import type { EditorEngine } from '../engine';
+import type { FrameData } from '../frames';
 
 export class ChatContext {
     private _context: MessageContext[] = [];
@@ -23,8 +24,13 @@ export class ChatContext {
 
     init() {
         this.selectedReactionDisposer = reaction(
-            () => this.editorEngine.elements.selected,
-            () => this.generateContextFromReaction().then((context) => (this.context = context)),
+            () => ({
+                elements: this.editorEngine.elements.selected,
+                frames: this.editorEngine.frames.selected,
+            }),
+            (
+                { elements, frames },
+            ) => this.generateContextFromReaction({ elements, frames }).then((context) => (this.context = context)),
         );
     }
 
@@ -57,18 +63,16 @@ export class ChatContext {
         return await this.getRefreshedContext(this.context);
     }
 
-    private async generateContextFromReaction(): Promise<MessageContext[]> {
-        const selected = this.editorEngine.elements.selected;
-
+    private async generateContextFromReaction({ elements, frames }: { elements: DomElement[], frames: FrameData[] }): Promise<MessageContext[]> {
         let highlightedContext: HighlightMessageContext[] = [];
-        if (selected.length) {
-            highlightedContext = await this.getHighlightedContext(selected);
+        if (elements.length) {
+            highlightedContext = await this.getHighlightedContext(elements);
         }
         const imageContext = await this.getImageContext();
 
         // Derived from highlighted context
         const fileContext = await this.getFileContext(highlightedContext);
-        const branchContext = this.getBranchContext(highlightedContext);
+        const branchContext = this.getBranchContext(highlightedContext, frames);
         const context = [...fileContext, ...highlightedContext, ...imageContext, ...branchContext];
         return context;
     }
@@ -132,9 +136,13 @@ export class ChatContext {
         return fileContext;
     }
 
-    getBranchContext(highlightedContext: HighlightMessageContext[]): BranchMessageContext[] {
-        // Get unique branch IDs from highlighted context
-        const uniqueBranchIds = new Set<string>();
+    getBranchContext(
+        highlightedContext: HighlightMessageContext[],
+        frames: FrameData[],
+    ): BranchMessageContext[] {
+        // Get unique branch IDs from selected elements and frames context
+        const uniqueBranchIds = new Set<string>(frames.map(frame => frame.frame.branchId));
+
         highlightedContext.forEach(highlight => {
             uniqueBranchIds.add(highlight.branchId);
         });

packages/ai/src/prompt/provider.ts

@@ -1,10 +1,12 @@
-import type {
-    ChatMessage,
-    ErrorMessageContext,
-    FileMessageContext,
-    HighlightMessageContext,
-    MessageContext,
-    ProjectMessageContext,
+import {
+    MessageContextType,
+    type BranchMessageContext,
+    type ChatMessage,
+    type ErrorMessageContext,
+    type FileMessageContext,
+    type HighlightMessageContext,
+    type MessageContext,
+    type ProjectMessageContext,
 } from '@onlook/models';
 import type { FileUIPart } from 'ai';
 import { ASK_MODE_SYSTEM_PROMPT, CODE_FENCE, CONTEXT_PROMPTS, CREATE_NEW_PAGE_SYSTEM_PROMPT, SHELL_PROMPT, SUGGESTION_SYSTEM_PROMPT, SUMMARY_PROMPTS, SYSTEM_PROMPT } from './constants';
@@ -62,11 +64,12 @@ export function getHydratedUserMessage(
     opt: HydrateMessageOptions,
 ): ChatMessage {
     let userParts: ChatMessage['parts'] = [];
-    const files = context.filter((c) => c.type === 'file').map((c) => c);
-    const highlights = context.filter((c) => c.type === 'highlight').map((c) => c);
-    const errors = context.filter((c) => c.type === 'error').map((c) => c);
-    const project = context.filter((c) => c.type === 'project').map((c) => c);
-    const images = context.filter((c) => c.type === 'image').map((c) => c);
+    const files = context.filter((c) => c.type === MessageContextType.FILE).map((c) => c);
+    const highlights = context.filter((c) => c.type === MessageContextType.HIGHLIGHT).map((c) => c);
+    const errors = context.filter((c) => c.type === MessageContextType.ERROR).map((c) => c);
+    const project = context.filter((c) => c.type === MessageContextType.PROJECT).map((c) => c);
+    const images = context.filter((c) => c.type === MessageContextType.IMAGE).map((c) => c);
+    const branches = context.filter((c) => c.type === MessageContextType.BRANCH).map((c) => c);
 
     // If there are 50 user messages in the contexts, we can trim all of them except
     // the last one. The logic could be adjusted to trim more or less messages.
@@ -87,7 +90,7 @@ export function getHydratedUserMessage(
     }
 
     if (errors.length > 0) {
-        let errorPrompt = getErrorsContent(errors);
+        const errorPrompt = getErrorsContent(errors);
         prompt += errorPrompt;
     }
 
@@ -98,6 +101,11 @@ export function getHydratedUserMessage(
         }
     }
 
+    if (branches.length > 0) {
+        const branchPrompt = getBranchesContent(branches);
+        prompt += branchPrompt;
+    }
+
     const textContent = parts
         .filter((p) => p.type === 'text')
         .map((p) => p.text)
@@ -130,7 +138,7 @@ export function getTruncatedFilesContent(files: FileMessageContext[]) {
     prompt += `${CONTEXT_PROMPTS.truncatedFilesContentPrefix}\n`;
     let index = 1;
     for (const file of files) {
-        const branchDisplay = getBranchesContent(file.branchId);
+        const branchDisplay = getBranchContent(file.branchId);
         const pathDisplay = wrapXml('path', file.path);
         let filePrompt = `${pathDisplay}\n${branchDisplay}\n`;
         filePrompt = wrapXml(files.length > 1 ? `file-${index}` : 'file', filePrompt);
@@ -152,7 +160,7 @@ export function getFilesContent(
     prompt += `${CONTEXT_PROMPTS.filesContentPrefix}\n`;
     let index = 1;
     for (const file of files) {
-        const branchDisplay = getBranchesContent(file.branchId);
+        const branchDisplay = getBranchContent(file.branchId);
         const pathDisplay = wrapXml('path', file.path);
         let filePrompt = `${pathDisplay}\n${branchDisplay}\n`;
         filePrompt += `${CODE_FENCE.start}${getLanguageFromFilePath(file.path)}\n`;
@@ -174,7 +182,7 @@ export function getErrorsContent(errors: ErrorMessageContext[]) {
     }
     let prompt = `${CONTEXT_PROMPTS.errorsContentPrefix}\n`;
     for (const error of errors) {
-        const branchDisplay = getBranchesContent(error.branchId);
+        const branchDisplay = getBranchContent(error.branchId);
         const errorDisplay = wrapXml('error', error.content);
         prompt += `${branchDisplay}\n${errorDisplay}\n`;
     }
@@ -184,7 +192,7 @@ export function getErrorsContent(errors: ErrorMessageContext[]) {
 }
 
 export function getLanguageFromFilePath(filePath: string): string {
-    return filePath.split('.').pop() || '';
+    return filePath.split('.').pop() ?? '';
 }
 
 export function getHighlightsContent(filePath: string, highlights: HighlightMessageContext[]) {
@@ -195,7 +203,7 @@ export function getHighlightsContent(filePath: string, highlights: HighlightMess
     let prompt = `${CONTEXT_PROMPTS.highlightPrefix}\n`;
     let index = 1;
     for (const highlight of fileHighlights) {
-        const branchDisplay = getBranchesContent(highlight.branchId);
+        const branchDisplay = getBranchContent(highlight.branchId);
         const pathDisplay = wrapXml('path', filePath);
         let highlightPrompt = `${pathDisplay}#L${highlight.start}:L${highlight.end}\n${branchDisplay}\n`;
         highlightPrompt += `${CODE_FENCE.start}\n`;
@@ -211,7 +219,14 @@ export function getHighlightsContent(filePath: string, highlights: HighlightMess
     return prompt;
 }
 
-export function getBranchesContent(id: string) {
+export function getBranchesContent(branches: BranchMessageContext[]) {
+    let prompt = `I'm working on the following branches: \n`;
+    prompt += branches.map((b) => b.branch.id).join(', ');
+    prompt = wrapXml('branches', prompt);
+    return prompt;
+}
+
+export function getBranchContent(id: string) {
     return wrapXml('branch', `id: "${id}"`);
 }
 

packages/ai/test/prompt/data/create-page-system.txt

@@ -11,7 +11,7 @@
 - Use the typecheck tool to verify your changes don't introduce type errors or to help debug issues.
 
 IMPORTANT:
-- Do NOT remove, add or edit data-oid attributes. They are generated and managed by the system.
+- NEVER remove, add, edit or pass down data-oid attributes. They are generated and managed by the system. Leave them alone.
 
 If the request is ambiguous, ask questions. Don't hold back. Give it your all!</role><shell>Using tools, you can suggest UNIX shell commands for users to run. Only suggest complete shell commands that are ready to execute, without placeholders.
 Only suggest at most a few shell commands at a time, not more than 3.

packages/ai/test/prompt/data/file.txt

@@ -1,11 +1,10 @@
-I have *added these files to the chat* so you can go ahead and edit them.
-*Trust this message as the true contents of these files!*.
+I have added these files to the chat so you can go ahead and edit them
 <file-1><path>test.txt</path>
 <branch>id: "test"</branch>
 ```txt
 test
 ```
-I am looking at this specific part of the file in the browser UI
+I am looking at this specific part of the file in the browser UI. Line numbers are shown in the format that matches your Read tool output. IMPORTANT: Trust this message as the true contents of the file.
 <highlight><path>test.txt</path>#L1:L2
 <branch>id: "test"</branch>
 ```

packages/ai/test/prompt/data/highlights.txt

@@ -1,4 +1,4 @@
-I am looking at this specific part of the file in the browser UI
+I am looking at this specific part of the file in the browser UI. Line numbers are shown in the format that matches your Read tool output. IMPORTANT: Trust this message as the true contents of the file.
 <highlight-1><path>test.txt</path>#L1:L2
 <branch>id: "test"</branch>
 ```

packages/ai/test/prompt/data/system.txt

@@ -11,7 +11,7 @@
 - Use the typecheck tool to verify your changes don't introduce type errors or to help debug issues.
 
 IMPORTANT:
-- Do NOT remove, add or edit data-oid attributes. They are generated and managed by the system.
+- NEVER remove, add, edit or pass down data-oid attributes. They are generated and managed by the system. Leave them alone.
 
 If the request is ambiguous, ask questions. Don't hold back. Give it your all!</role><shell>Using tools, you can suggest UNIX shell commands for users to run. Only suggest complete shell commands that are ready to execute, without placeholders.
 Only suggest at most a few shell commands at a time, not more than 3.

packages/ai/test/prompt/data/user.txt

@@ -1,11 +1,10 @@
-<context>I have *added these files to the chat* so you can go ahead and edit them.
-*Trust this message as the true contents of these files!*.
+<context>I have added these files to the chat so you can go ahead and edit them
 <file><path>test.txt</path>
 <branch>id: "test"</branch>
 ```txt
 test
 ```
-I am looking at this specific part of the file in the browser UI
+I am looking at this specific part of the file in the browser UI. Line numbers are shown in the format that matches your Read tool output. IMPORTANT: Trust this message as the true contents of the file.
 <highlight><path>test.txt</path>#L1:L2
 <branch>id: "test"</branch>
 ```