feat: add agent rules to context (#2946)

Author: Kitenite

apps/web/client/public/onlook-preload-script.js

@@ -33,12 +33,12 @@ export function getContextIcon(context: MessageContext) {
             return (
                 <NodeIcon tagName={context.displayName} iconClass="w-3 h-3 ml-1 mr-2 flex-none" />
             );
-        case MessageContextType.PROJECT:
-            icon = Icons.Cube;
-            break;
         case MessageContextType.BRANCH:
             icon = Icons.Branch;
             break;
+        case MessageContextType.AGENT_RULE:
+            icon = Icons.Cube;
+            break;
         default:
             assertNever(context);
     }

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/context-pills/helpers.tsx

@@ -9,11 +9,11 @@ import { DraftImagePill } from './draft-image-pill';
 
 const typeOrder = {
     [MessageContextType.BRANCH]: 0,
-    [MessageContextType.PROJECT]: 1,
-    [MessageContextType.FILE]: 2,
-    [MessageContextType.HIGHLIGHT]: 3,
-    [MessageContextType.ERROR]: 4,
-    [MessageContextType.IMAGE]: 5,
+    [MessageContextType.FILE]: 1,
+    [MessageContextType.HIGHLIGHT]: 2,
+    [MessageContextType.ERROR]: 3,
+    [MessageContextType.IMAGE]: 4,
+    [MessageContextType.AGENT_RULE]: 5,
 };
 
 export const InputContextPills = observer(() => {

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/context-pills/input-context-pills.tsx

@@ -1,13 +1,13 @@
 import { ChatType, type DomElement } from '@onlook/models';
 import {
     MessageContextType,
+    type AgentRuleMessageContext,
     type BranchMessageContext,
     type ErrorMessageContext,
     type FileMessageContext,
     type HighlightMessageContext,
     type ImageMessageContext,
-    type MessageContext,
-    type ProjectMessageContext,
+    type MessageContext
 } from '@onlook/models/chat';
 import { assertNever, type ParsedError } from '@onlook/utility';
 import { makeAutoObservable, reaction } from 'mobx';
@@ -51,16 +51,16 @@ export class ChatContext {
             case ChatType.EDIT:
             case ChatType.CREATE:
             case ChatType.ASK:
-                return await this.getLatestContext();
+                return await this.getChatEditContext();
             case ChatType.FIX:
                 return this.getErrorContext();
             default:
                 assertNever(type);
         }
     }
 
-    async getLatestContext(): Promise<MessageContext[]> {
-        return await this.getRefreshedContext(this.context);
+    async getChatEditContext(): Promise<MessageContext[]> {
+        return [...await this.getRefreshedContext(this.context), ...await this.getAgentRuleContext()];
     }
 
     private async generateContextFromReaction({ elements, frames }: { elements: DomElement[], frames: FrameData[] }): Promise<MessageContext[]> {
@@ -215,16 +215,34 @@ export class ChatContext {
         };
     }
 
-    // TODO: Enhance with custom rules
-    getProjectContext(): ProjectMessageContext[] {
-        return [
-            {
-                type: MessageContextType.PROJECT,
-                content: '',
-                displayName: 'Project',
-                path: './',
-            },
-        ];
+    private async getAgentRuleContext(): Promise<AgentRuleMessageContext[]> {
+        try {
+            const agentRuleFilesPaths = ['agents.md', 'claude.md', 'AGENTS.md', 'CLAUDE.md'];
+            const agentRuleContexts: AgentRuleMessageContext[] = (await Promise.all(
+                agentRuleFilesPaths.map(async (filePath) => {
+                    const file = await this.editorEngine.activeSandbox.readFile(`./${filePath}`);
+                    if (file === null) {
+                        return null;
+                    }
+                    if (file.type === 'binary') {
+                        return null;
+                    }
+                    if (file.content.trim().length === 0) {
+                        return null;
+                    }
+                    return {
+                        type: MessageContextType.AGENT_RULE,
+                        content: file.content,
+                        displayName: filePath,
+                        path: filePath,
+                    } satisfies AgentRuleMessageContext;
+                })
+            )).filter((context) => context !== null);
+            return agentRuleContexts
+        } catch (error) {
+            console.error('Error getting agent rule context', error);
+            return [];
+        }
     }
 
     getErrorContext(): ErrorMessageContext[] {

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

@@ -21,12 +21,12 @@ IMPORTANT: This project uses Bun as the package manager. Always use Bun commands
 
 NEVER SUGGEST THE "bun run dev" command. Assume the user is already running the app.`;
 
-const projectContextPrefix = `This is a Nextjs project with TailwindCSS`;
+const agentRulesContextPrefix = `These are user provided rules for the project`;
 
 export const CONTEXT_PROMPTS = {
     filesContentPrefix,
     truncatedFilesContentPrefix,
     highlightPrefix,
     errorsContentPrefix,
-    projectContextPrefix,
+    agentRulesContextPrefix,
 };

packages/ai/src/prompt/constants/context.ts

@@ -1,12 +1,12 @@
 import {
     MessageContextType,
+    type AgentRuleMessageContext,
     type BranchMessageContext,
     type ChatMessage,
     type ErrorMessageContext,
     type FileMessageContext,
     type HighlightMessageContext,
-    type MessageContext,
-    type ProjectMessageContext,
+    type MessageContext
 } 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';
@@ -67,7 +67,7 @@ export function getHydratedUserMessage(
     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 agentRules = context.filter((c) => c.type === MessageContextType.AGENT_RULE).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);
 
@@ -94,11 +94,9 @@ export function getHydratedUserMessage(
         prompt += errorPrompt;
     }
 
-    if (project.length > 0) {
-        const projectContext = project[0];
-        if (projectContext) {
-            prompt += getProjectContext(projectContext);
-        }
+    if (agentRules.length > 0) {
+        const agentRulePrompt = getAgentRulesContent(agentRules);
+        prompt += agentRulePrompt;
     }
 
     if (branches.length > 0) {
@@ -242,7 +240,11 @@ export function getSummaryPrompt() {
     return prompt;
 }
 
-export function getProjectContext(project: ProjectMessageContext) {
-    const content = `${CONTEXT_PROMPTS.projectContextPrefix} ${project.path}`;
-    return wrapXml('project-info', content);
+export function getAgentRulesContent(agentRules: AgentRuleMessageContext[]) {
+    let content = `${CONTEXT_PROMPTS.agentRulesContextPrefix}\n`;
+    for (const agentRule of agentRules) {
+        const agentRulePrompt = `${agentRule.path}\n${agentRule.content}`;
+        content += agentRulePrompt;
+    }
+    return wrapXml('agent-rules', content);
 }

packages/ai/src/prompt/provider.ts

@@ -5,8 +5,8 @@ export enum MessageContextType {
     HIGHLIGHT = 'highlight',
     IMAGE = 'image',
     ERROR = 'error',
-    PROJECT = 'project',
     BRANCH = 'branch',
+    AGENT_RULE = 'agent_rule',
 }
 
 type BaseMessageContext = {
@@ -45,15 +45,15 @@ export type ErrorMessageContext = BaseMessageContext & {
     branchId: string;
 };
 
-export type ProjectMessageContext = BaseMessageContext & {
-    type: MessageContextType.PROJECT;
+export type AgentRuleMessageContext = BaseMessageContext & {
+    type: MessageContextType.AGENT_RULE;
     path: string;
 };
 
 export type MessageContext =
     | HighlightMessageContext
     | ImageMessageContext
     | ErrorMessageContext
-    | ProjectMessageContext
+    | AgentRuleMessageContext
     | BranchMessageContext
     | FileMessageContext;